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
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) ]
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
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
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
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