Example #1
0
    def visit_BinOp(self, t):
        self.visit(t.left)
        self.visit(t.right)

        if util.only_scalars_involed(self.defined_symbols, t.left, t.right):
            return self.generic_visit(t)

        # Detect fused operations

        # MAD: These functions multiply the first two floating-point inputs and add the result to the third input.
        # MLA: These functions multiply the second and third floating-point inputs and add the result to the first input.
        # MSB: These functions multiply the first two floating-point inputs and subtract the result from the third input.
        # MLS: These functions multiply the second and third floating-point inputs and subtract the result from the first input.

        parent_op = t.op.__class__
        left_op = None
        right_op = None

        if isinstance(t.left, ast.BinOp):
            left_op = t.left.op.__class__
        if isinstance(t.right, ast.BinOp):
            right_op = t.right.op.__class__

        args = []
        name = None

        if parent_op == ast.Add:
            if left_op == ast.Mult:
                name = '__svmad_'
                args = [t.left.left, t.left.right, t.right]
            elif right_op == ast.Mult:
                name = '__svmla_'
                args = [t.left, t.right.left, t.right.right]
        elif parent_op == ast.Sub:
            if left_op == ast.Mult:
                name = '__svmsb_'
                args = [t.left.left, t.left.right, t.right]
            elif right_op == ast.Mult:
                name = '__svmls_'
                args = [t.left, t.right.left, t.right.right]

        # Fused ops need at least two of three arguments to be a vector
        if name:
            inferred = util.infer_ast(self.defined_symbols, *args)
            scalar_args = sum([util.is_scalar(tp) for tp in inferred])
            if scalar_args > 1:
                return self.generic_visit(t)
            # Add the type suffix for internal representation
            name += util.TYPE_TO_SVE_SUFFIX[util.get_base_type(
                dace.dtypes.result_type_of(*inferred))]
            return ast.copy_location(
                ast.Call(func=ast.Name(name, ast.Load()),
                         args=args,
                         keywords=[]), t)

        return self.generic_visit(t)
Example #2
0
    def generate_out_register(self,
                              sdfg: SDFG,
                              state: SDFGState,
                              edge: graph.MultiConnectorEdge[mm.Memlet],
                              code: CodeIOStream,
                              use_data_name: bool = False) -> bool:
        """
            Responsible for generating temporary out registers in a Tasklet, given an outgoing edge.
            Returns `True` if a writeback of this register is needed.
        """
        if edge.src_conn is None:
            return

        dst_node = state.memlet_path(edge)[-1].dst

        src_type = edge.src.out_connectors[edge.src_conn]
        src_name = edge.src_conn

        if use_data_name:
            src_name = edge.data.data

        if isinstance(dst_node, nodes.AccessNode) and isinstance(
                dst_node.desc(sdfg), data.Stream):
            # Streams don't need writeback and are treated differently
            self.stream_associations[edge.src_conn] = (edge.data.data,
                                                       src_type.base_type)
            return False
        elif edge.data.wcr is not None:
            # WCR is addressed within the unparser to capture conditionals
            self.wcr_associations[edge.src_conn] = (dst_node, edge,
                                                    src_type.base_type)
            return False

        # Create temporary registers
        ctype = None
        if util.is_vector(src_type):
            ctype = util.TYPE_TO_SVE[src_type.type]
        elif util.is_scalar(src_type):
            ctype = src_type.ctype
        else:
            raise util.NotSupportedError(
                'Unsupported Code->Code edge (pointer)')

        self.dispatcher.defined_vars.add(src_name, DefinedType.Scalar, ctype)
        code.write(f'{ctype} {src_name};')

        return True
Example #3
0
    def _UnaryOp(self, t):
        inf_type = self.infer(t.operand)[0]

        if util.is_scalar(inf_type):
            return super()._UnaryOp(t)

        if isinstance(t.op, ast.UAdd):
            # A + in front is just ignored
            t.dispatch(t.operand)
            return

        if t.op.__class__ not in util.UN_OP_TO_SVE:
            raise NotImplementedError(
                f'Unary operation {t.op.__class__.__name__} not implemented')

        self.write('{}_x({}, '.format(util.UN_OP_TO_SVE[t.op.__class__],
                                      self.pred_name))
        self.dispatch(t.operand)
        self.write(')')
Example #4
0
    def generate_writeback(self, sdfg: SDFG, state: SDFGState, map: nodes.Map,
                           edge: graph.MultiConnectorEdge[mm.Memlet], code: CodeIOStream):
        """
            Responsible for generating code for a writeback in a Tasklet, given the outgoing edge.
            This is mainly taking the temporary register and writing it back.
        """
        if edge.src_conn is None:
            return

        dst_node = state.memlet_path(edge)[-1].dst

        src_type = edge.src.out_connectors[edge.src_conn]
        src_name = edge.src_conn

        if isinstance(dst_node, nodes.Tasklet):
            ##################
            # Code->Code edges
            dst_type = edge.dst.in_connectors[edge.dst_conn]

            if (util.is_vector(src_type) and util.is_vector(dst_type)) or (util.is_scalar(src_type)
                                                                           and util.is_scalar(dst_type)):
                # Simply write back to shared register
                code.write(f'{edge.data.data} = {src_name};')
            elif util.is_scalar(src_type) and util.is_vector(dst_type):
                # Scalar broadcast to shared vector register
                code.write(f'{edge.data.data} = svdup_{util.TYPE_TO_SVE_SUFFIX[dst_type.type]}({src_name});')
            else:
                raise util.NotSupportedError('Unsupported Code->Code edge')
        elif isinstance(dst_node, nodes.AccessNode):
            ##################
            # Write to AccessNode
            desc = dst_node.desc(sdfg)
            if isinstance(desc, data.Array):
                ##################
                # Write into Array
                if util.is_pointer(src_type):
                    raise util.NotSupportedError('Unsupported writeback')
                elif util.is_vector(src_type):
                    ##################
                    # Scatter vector store into array

                    stride = edge.data.get_stride(sdfg, map)

                    # long long fix
                    ptr_cast = ''
                    if src_type.type == np.int64:
                        ptr_cast = '(int64_t*) '
                    elif src_type.type == np.uint64:
                        ptr_cast = '(uint64_t*) '

                    store_args = '{}, {}'.format(
                        util.get_loop_predicate(sdfg, state, edge.src),
                        ptr_cast + cpp.cpp_ptr_expr(sdfg, edge.data, DefinedType.Pointer, codegen=self.frame),
                    )

                    if stride == 1:
                        code.write(f'svst1({store_args}, {src_name});')
                    else:
                        code.write(
                            f'svst1_scatter_index({store_args}, svindex_s{util.get_base_type(src_type).bytes * 8}(0, {sym2cpp(stride)}), {src_name});'
                        )
                else:
                    ##################
                    # Scalar write into array
                    code.write(f'{cpp.cpp_array_expr(sdfg, edge.data, codegen=self.frame)} = {src_name};')
            elif isinstance(desc, data.Scalar):
                ##################
                # Write into Scalar
                if util.is_pointer(src_type):
                    raise util.NotSupportedError('Unsupported writeback')
                elif util.is_vector(src_type):
                    if util.is_vector(desc.dtype):
                        ##################
                        # Vector write into vector Scalar access node
                        code.write(f'{edge.data.data} = {src_name};')
                    else:
                        raise util.NotSupportedError('Unsupported writeback')
                else:
                    if util.is_vector(desc.dtype):
                        ##################
                        # Broadcast into scalar AccessNode
                        code.write(f'{edge.data.data} = svdup_{util.TYPE_TO_SVE_SUFFIX[src_type]}({src_name});')
                    else:
                        ##################
                        # Scalar write into scalar AccessNode
                        code.write(f'{edge.data.data} = {src_name};')

        else:
            raise util.NotSupportedError('Only writeback to Tasklets and AccessNodes is supported')
Example #5
0
    def generate_read(self, sdfg: SDFG, state: SDFGState, map: nodes.Map, edge: graph.MultiConnectorEdge[mm.Memlet],
                      code: CodeIOStream):
        """
            Responsible for generating code for reads into a Tasklet, given the ingoing edge.
        """
        if edge.dst_conn is None:
            return
        src_node = state.memlet_path(edge)[0].src
        dst_type = edge.dst.in_connectors[edge.dst_conn]
        dst_name = edge.dst_conn
        if isinstance(src_node, nodes.Tasklet):
            ##################
            # Code->Code edges
            src_type = edge.src.out_connectors[edge.src_conn]
            if util.is_vector(src_type) and util.is_vector(dst_type):
                # Directly read from shared vector register
                code.write(f'{util.TYPE_TO_SVE[dst_type.type]} {dst_name} = {edge.data.data};')
            elif util.is_scalar(src_type) and util.is_scalar(dst_type):
                # Directly read from shared scalar register
                code.write(f'{dst_type} {dst_name} = {edge.data.data};')
            elif util.is_scalar(src_type) and util.is_vector(dst_type):
                # Scalar broadcast from shared scalar register
                code.write(
                    f'{util.TYPE_TO_SVE[dst_type.type]} {dst_name} = svdup_{util.TYPE_TO_SVE_SUFFIX[dst_type.type]}({edge.data.data});'
                )
            else:
                raise util.NotSupportedError('Unsupported Code->Code edge')
        elif isinstance(src_node, nodes.AccessNode):
            ##################
            # Read from AccessNode
            desc = src_node.desc(sdfg)
            if isinstance(desc, data.Array):
                # Copy from array
                if util.is_pointer(dst_type):
                    ##################
                    # Pointer reference
                    code.write(
                        f'{dst_type} {dst_name} = {cpp.cpp_ptr_expr(sdfg, edge.data, None, codegen=self.frame)};')
                elif util.is_vector(dst_type):
                    ##################
                    # Vector load

                    stride = edge.data.get_stride(sdfg, map)

                    # First part of the declaration is `type name`
                    load_lhs = '{} {}'.format(util.TYPE_TO_SVE[dst_type.type], dst_name)

                    # long long issue casting
                    ptr_cast = ''
                    if dst_type.type == np.int64:
                        ptr_cast = '(int64_t*) '
                    elif dst_type.type == np.uint64:
                        ptr_cast = '(uint64_t*) '

                    # Regular load and gather share the first arguments
                    load_args = '{}, {}'.format(
                        util.get_loop_predicate(sdfg, state, edge.dst),
                        ptr_cast + cpp.cpp_ptr_expr(sdfg, edge.data, DefinedType.Pointer, codegen=self.frame))

                    if stride == 1:
                        code.write('{} = svld1({});'.format(load_lhs, load_args))
                    else:
                        code.write('{} = svld1_gather_index({}, svindex_s{}(0, {}));'.format(
                            load_lhs, load_args,
                            util.get_base_type(dst_type).bytes * 8, sym2cpp(stride)))
                else:
                    ##################
                    # Scalar read from array
                    code.write(f'{dst_type} {dst_name} = {cpp.cpp_array_expr(sdfg, edge.data, codegen=self.frame)};')
            elif isinstance(desc, data.Scalar):
                # Refer to shared variable
                src_type = desc.dtype
                if util.is_vector(src_type) and util.is_vector(dst_type):
                    # Directly read from shared vector register
                    code.write(f'{util.TYPE_TO_SVE[dst_type.type]} {dst_name} = {edge.data.data};')
                elif util.is_scalar(src_type) and util.is_scalar(dst_type):
                    # Directly read from shared scalar register
                    code.write(f'{dst_type} {dst_name} = {edge.data.data};')
                elif util.is_scalar(src_type) and util.is_vector(dst_type):
                    # Scalar broadcast from shared scalar register
                    code.write(
                        f'{util.TYPE_TO_SVE[dst_type.type]} {dst_name} = svdup_{util.TYPE_TO_SVE_SUFFIX[dst_type.type]}({edge.data.data});'
                    )
                else:
                    raise util.NotSupportedError('Unsupported Scalar->Code edge')
        else:
            raise util.NotSupportedError('Only copy from Tasklets and AccessNodes is supported')