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