Ejemplo n.º 1
0
    def generate_sid_params(self, sdfg: dace.SDFG):
        res = []
        import dace.data

        for name, array in sdfg.arrays.items():
            if array.transient:
                continue
            domain_dim_flags = tuple(True if any(
                dace.symbolic.pystr_to_symbolic(f"__{dim.upper()}") in
                s.free_symbols for s in array.shape
                if hasattr(s, "free_symbols")) else False for dim in "ijk")
            data_ndim = len(array.shape) - sum(array_dimensions(array))
            sid_def = pybuffer_to_sid(
                name=name,
                ctype=array.dtype.ctype,
                domain_dim_flags=domain_dim_flags,
                data_ndim=data_ndim,
                stride_kind_index=self.unique_index(),
                backend=self.backend,
            )

            res.append(sid_def)
        # pass scalar parameters as variables
        for name in (n for n in sdfg.symbols.keys() if not n.startswith("__")):
            res.append(name)
        return res
Ejemplo n.º 2
0
    def generate_dace_args(self, gtir, sdfg):
        offset_dict: Dict[str, Tuple[int, int, int]] = {
            k: (-v[0][0], -v[1][0], -v[2][0])
            for k, v in compute_legacy_extents(gtir).items()
        }
        symbols = {f"__{var}": f"__{var}" for var in "IJK"}
        for name, array in sdfg.arrays.items():
            if array.transient:
                symbols[f"__{name}_K_stride"] = "1"
                symbols[f"__{name}_J_stride"] = str(array.shape[2])
                symbols[f"__{name}_I_stride"] = str(array.shape[1] *
                                                    array.shape[2])
            else:
                dims = [
                    dim for dim, select in zip("IJK", array_dimensions(array))
                    if select
                ]
                data_ndim = len(array.shape) - len(dims)

                # api field strides
                fmt = "gt::sid::get_stride<{dim}>(gt::sid::get_strides(__{name}_sid))"

                symbols.update({
                    f"__{name}_{dim}_stride":
                    fmt.format(dim=f"gt::stencil::dim::{dim.lower()}",
                               name=name)
                    for dim in dims
                })
                symbols.update({
                    f"__{name}_d{dim}_stride":
                    fmt.format(dim=f"gt::integral_constant<int, {3 + dim}>",
                               name=name)
                    for dim in range(data_ndim)
                })

                # api field pointers
                fmt = """gt::sid::multi_shifted(
                             gt::sid::get_origin(__{name}_sid)(),
                             gt::sid::get_strides(__{name}_sid),
                             std::array<gt::int_t, {ndim}>{{{origin}}}
                         )"""
                origin = tuple(-offset_dict[name][idx]
                               for idx, var in enumerate("IJK") if any(
                                   dace.symbolic.pystr_to_symbolic(f"__{var}")
                                   in s.free_symbols for s in array.shape
                                   if hasattr(s, "free_symbols")))
                symbols[name] = fmt.format(name=name,
                                           ndim=len(array.shape),
                                           origin=",".join(
                                               str(o) for o in origin))
        # the remaining arguments are variables and can be passed by name
        for sym in sdfg.signature_arglist(with_types=False, for_call=True):
            if sym not in symbols:
                symbols[sym] = sym

        # return strings in order of sdfg signature
        return [
            symbols[s]
            for s in sdfg.signature_arglist(with_types=False, for_call=True)
        ]
Ejemplo n.º 3
0
    def get_mapped_subsets_dicts(self, interval: Interval, section: dace.SDFG):

        in_subsets = dict()
        out_subsets = dict()
        section_origins: Dict[str, Tuple[int, int]] = dict()
        min_k_offsets: Dict[str, int] = dict()
        for he in (ln for ln, _ in section.all_nodes_recursive()
                   if isinstance(ln, (HorizontalExecutionLibraryNode,
                                      VerticalLoopLibraryNode))):
            access_collection: AccessCollector.CartesianAccessCollection = get_access_collection(
                he)

            for name, offsets in access_collection.offsets().items():
                off: Tuple[int, int, int]
                for off in offsets:
                    origin = (
                        -off[0] - he.iteration_space.i_interval.start.offset,
                        -off[1] - he.iteration_space.j_interval.start.offset,
                    )
                    if name not in section_origins:
                        section_origins[name] = origin
                    if name not in min_k_offsets:
                        min_k_offsets[name] = off[2]
                    section_origins[name] = (
                        max(section_origins[name][0], origin[0]),
                        max(section_origins[name][1], origin[1]),
                    )
                    min_k_offsets[name] = min(min_k_offsets[name], off[2])
        access_collection = get_access_collection(section)
        for name, section_origin in section_origins.items():
            vl_origin = self.origins[name]
            shape = section.arrays[name].shape
            dimensions = array_dimensions(section.arrays[name])
            subset_strs = []
            idx = iter(range(3))
            if dimensions[0]:
                subset_strs.append("{i:+d}:{i:+d}+({I})".format(
                    i=vl_origin[0] - section_origin[0],
                    I=shape[next(idx)],
                ))
            if dimensions[1]:
                subset_strs.append("{j:+d}:{j:+d}+({J})".format(
                    j=vl_origin[1] - section_origin[1], J=shape[next(idx)]))
            if dimensions[2]:
                subset_strs.append(
                    "k-({k_orig}){k:+d}:k-({k_orig}){k:+d}{K:+d}".format(
                        k_orig=get_axis_bound_str(vl_origin[2], "__K"),
                        k=min_k_offsets[name],
                        K=shape[next(idx)],
                    ))
            data_dims = shape[sum(dimensions):]
            subset_strs.extend([f"0:{d}" for d in data_dims])
            subset_str = ",".join(subset_strs)
            if name in access_collection.read_fields():
                in_subsets[name] = subset_str
            if name in access_collection.write_fields():
                out_subsets[name] = subset_str
        return in_subsets, out_subsets
Ejemplo n.º 4
0
    def get_innermost_memlets(self):

        access_collection: AccessCollector.CartesianAccessCollection = get_access_collection(
            self.node)

        in_memlets = dict()
        for name, offsets in access_collection.read_offsets().items():
            dimensions = array_dimensions(self.parent_sdfg.arrays[name])
            data_dims = self.parent_sdfg.arrays[name].shape[sum(dimensions):]

            for off in offsets:
                subset_strs = [
                    f"{var}{self.origins[name][dim] + off[dim]:+d}"
                    for dim, var in enumerate("ij0") if dimensions[dim]
                ]
                subset_strs.extend(f"0:{dim}" for dim in data_dims)
                acc_name = get_tasklet_symbol(name, off, is_target=False)
                in_memlets[acc_name] = dace.memlet.Memlet.simple(
                    name, ",".join(subset_strs))

        out_memlets = dict()
        for name in access_collection.write_fields():
            dimensions = array_dimensions(self.parent_sdfg.arrays[name])
            data_dims = self.parent_sdfg.arrays[name].shape[sum(dimensions):]
            subset_strs = [
                f"{var}{self.origins[name][dim]:+d}"
                for dim, var in enumerate("ij0") if dimensions[dim]
            ]
            subset_strs.extend(f"0:{dim}" for dim in data_dims)
            acc_name = "__" + name
            out_memlets[acc_name] = dace.memlet.Memlet.simple(
                name,
                ",".join(subset_strs),
                dynamic=any(
                    isinstance(stmt, oir.MaskStmt)
                    for stmt in self.node.oir_node.body),
            )

        return in_memlets, out_memlets
Ejemplo n.º 5
0
    def generate_sid_params(self, sdfg: dace.SDFG):
        res = []
        import dace.data

        for name, array in sdfg.arrays.items():
            if array.transient:
                continue
            dimensions = array_dimensions(array)
            domain_ndim = sum(dimensions)
            data_ndim = len(array.shape) - domain_ndim
            sid_def = """gt::as_{sid_type}<{dtype}, {num_dims},
                gt::integral_constant<int, {unique_index}>>({name})""".format(
                sid_type="cuda_sid" if array.storage in [
                    dace.StorageType.GPU_Global, dace.StorageType.GPU_Shared
                ] else "sid",
                name=name,
                dtype=array.dtype.ctype,
                unique_index=self.unique_index(),
                num_dims=len(array.shape),
            )
            sid_def = "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format(
                sid_def=sid_def, name=name)

            if domain_ndim != 3:
                gt_dims = [
                    f"gt::stencil::dim::{dim}" for dim in "ijk" if any(
                        dace.symbolic.pystr_to_symbolic(f"__{dim.upper()}") in
                        s.free_symbols for s in array.shape
                        if hasattr(s, "free_symbols"))
                ]
                if data_ndim:
                    gt_dims += [
                        f"gt::integral_constant<int, {3 + dim}>"
                        for dim in range(data_ndim)
                    ]
                sid_def = "gt::sid::rename_numbered_dimensions<{gt_dims}>({sid_def})".format(
                    gt_dims=", ".join(gt_dims), sid_def=sid_def)

            res.append(sid_def)
        # pass scalar parameters as variables
        for name in (n for n in sdfg.symbols.keys() if not n.startswith("__")):
            res.append(name)
        return res
Ejemplo n.º 6
0
def sdfg_arrays_to_oir_decls(
        sdfg: dace.SDFG) -> Tuple[List[oir.Decl], List[oir.Temporary]]:
    params = list()
    decls = list()

    array: dace.data.Data
    for name, array in sdfg.arrays.items():
        dtype = common.typestr_to_data_type(dace_dtype_to_typestr(array.dtype))
        if isinstance(array, dace.data.Array):
            dimensions = array_dimensions(array)
            if not array.transient:
                params.append(
                    oir.FieldDecl(
                        name=name,
                        dtype=dtype,
                        dimensions=dimensions,
                        data_dims=array.shape[sum(dimensions):],
                    ))
            else:
                decls.append(
                    oir.Temporary(
                        name=name,
                        dtype=dtype,
                        dimensions=dimensions,
                        data_dims=array.shape[sum(dimensions):],
                    ))
        else:
            assert isinstance(array, dace.data.Scalar)
            params.append(oir.ScalarDecl(name=name, dtype=dtype))

    reserved_symbols = internal_symbols(sdfg)
    for sym, stype in sdfg.symbols.items():
        if sym not in reserved_symbols:
            params.append(
                oir.ScalarDecl(name=sym,
                               dtype=common.typestr_to_data_type(
                                   stype.as_numpy_dtype().str)))
    return params, decls