def visit_FieldDecl(self, node: cuir.FieldDecl, **kwargs):
if "external_arg" in kwargs:
if kwargs["external_arg"]:
return "py::buffer {name}, std::array<gt::uint_t,{ndim}> {name}_origin".format(
name=node.name,
ndim=node.dimensions.count(True),
)
else:
num_dims = node.dimensions.count(True)
sid_def = """gt::as_cuda_sid<{dtype}, {num_dims},
std::integral_constant<int, {unique_index}>>({name})""".format(
name=node.name,
dtype=self.visit(node.dtype),
unique_index=self.unique_index(),
num_dims=num_dims,
)
if num_dims != 3:
gt_dims = [
f"gt::stencil::dim::{dim}"
for dim in gtc_utils.dimension_flags_to_names(node.dimensions)
]
sid_def = "gt::sid::rename_numbered_dimensions<{gt_dims}>({sid_def})".format(
gt_dims=", ".join(gt_dims), sid_def=sid_def
)
return "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format(
sid_def=sid_def,
name=node.name,
)
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData:
data = ModuleData()
node = pipeline.full()
field_extents = compute_legacy_extents(node)
write_fields = (
node.iter_tree()
.if_isinstance(gtir.ParAssignStmt)
.getattr("left")
.if_isinstance(gtir.FieldAccess)
.getattr("name")
.to_set()
)
read_fields: Set[str] = set()
for expr in node.iter_tree().if_isinstance(gtir.ParAssignStmt).getattr("right"):
read_fields |= expr.iter_tree().if_isinstance(gtir.FieldAccess).getattr("name").to_set()
referenced_field_params = [
param.name for param in node.params if isinstance(param, gtir.FieldDecl)
]
for name in sorted(referenced_field_params):
access = AccessKind.NONE
if name in read_fields:
access |= AccessKind.READ
if name in write_fields:
access |= AccessKind.WRITE
data.field_info[name] = FieldInfo(
access=access,
boundary=field_extents[name].to_boundary(),
axes=tuple(dimension_flags_to_names(node.symtable_[name].dimensions).upper()),
data_dims=tuple(node.symtable_[name].data_dims),
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()),
)
referenced_scalar_params = [
param.name for param in node.params if param.name not in referenced_field_params
]
for name in sorted(referenced_scalar_params):
data.parameter_info[name] = ParameterInfo(
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower())
)
unref_params = get_unused_params_from_gtir(pipeline)
for param in sorted(unref_params, key=lambda decl: decl.name):
if isinstance(param, gtir.FieldDecl):
data.field_info[param.name] = None
elif isinstance(param, gtir.ScalarDecl):
data.parameter_info[param.name] = None
data.unreferenced = [*sorted(param.name for param in unref_params)]
return data
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData:
"""
Compute module data containing information about stencil arguments from gtir.
This is no longer compatible with the legacy backends.
"""
data = ModuleData()
# NOTE: pipeline.gtir has not had prune_unused_parameters applied.
all_params = pipeline.gtir.params
node = pipeline.full()
oir = gtir_to_oir.GTIRToOIR().visit(node)
field_extents = compute_fields_extents(oir)
accesses = compute_access_kinds(oir)
for decl in (param for param in all_params
if isinstance(param, gtir.FieldDecl)):
access = accesses[decl.name]
dtype = numpy.dtype(decl.dtype.name.lower())
if access != AccessKind.NONE:
k_boundary = compute_k_boundary(node)[decl.name]
boundary = Boundary(*field_extents[decl.name].to_boundary()[0:2],
k_boundary)
else:
boundary = Boundary.zeros(ndims=3)
data.field_info[decl.name] = FieldInfo(
access=access,
boundary=boundary,
axes=tuple(dimension_flags_to_names(decl.dimensions).upper()),
data_dims=tuple(decl.data_dims),
dtype=dtype,
)
for decl in (param for param in all_params
if isinstance(param, gtir.ScalarDecl)):
access = accesses[decl.name]
dtype = numpy.dtype(decl.dtype.name.lower())
data.parameter_info[decl.name] = ParameterInfo(access=access,
dtype=dtype)
data.unreferenced = [
*sorted(name for name in accesses if accesses[name] == AccessKind.NONE)
]
return data
def visit_FieldDecl(self, node: gtcpp.FieldDecl, **kwargs):
backend = kwargs["backend"]
if "external_arg" in kwargs:
domain_ndim = node.dimensions.count(True)
data_ndim = len(node.data_dims)
sid_ndim = domain_ndim + data_ndim
if kwargs["external_arg"]:
return "py::buffer {name}, std::array<gt::uint_t,{sid_ndim}> {name}_origin".format(
name=node.name,
sid_ndim=sid_ndim,
)
else:
make_layout_map = backend.storage_info["layout_map"]
layout_map = [
x for x in make_layout_map(node.dimensions +
(True, ) * data_ndim)
if x is not None
]
sid_def = """gt::as_{sid_type}<{dtype}, {sid_ndim},
gt::integral_constant<int, {unique_index}>,
{unit_stride_dim}>({name})""".format(
sid_type="cuda_sid"
if backend.GT_BACKEND_T == "gpu" else "sid",
name=node.name,
dtype=self.visit(node.dtype),
unique_index=self.unique_index(),
sid_ndim=sid_ndim,
unit_stride_dim=layout_map.index(max(layout_map)),
)
sid_def = "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format(
sid_def=sid_def,
name=node.name,
)
if domain_ndim != 3:
gt_dims = [
f"gt::stencil::dim::{dim}" for dim in
gtc_utils.dimension_flags_to_names(node.dimensions)
]
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)
return sid_def
def test_make_args_data_from_gtir(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals(
{"MODE": mode})
args_data = make_args_data_from_gtir(builder.gtir_pipeline)
assert set(args_data.unreferenced) == set(unreferenced_val[mode])
field_info_from_gtir = {(
p.name,
np.dtype(p.dtype.name.lower()),
utils.dimension_flags_to_names(p.dimensions).upper(),
p.data_dims,
)
for p in builder.gtir.params
if isinstance(p, gtir.FieldDecl)}
field_info_from_args_data = {(name, d.dtype, "".join(d.axes), d.data_dims)
for name, d in args_data.field_info.items()
if name not in args_data.unreferenced}
assert field_info_from_gtir == field_info_from_args_data
param_info_from_gtir = {(p.name, np.dtype(p.dtype.name.lower()))
for p in builder.gtir.params
if isinstance(p, gtir.ScalarDecl)}
param_info_from_args_data = {
(name, d.dtype)
for name, d in args_data.parameter_info.items()
if name not in args_data.unreferenced
}
assert param_info_from_gtir == param_info_from_args_data
for name, field_info in args_data.field_info.items():
if name == "out":
access = AccessKind.WRITE
elif name in field_info_val[mode]:
access = AccessKind.READ
else:
access = AccessKind.NONE
assert field_info.access == access
for name, param_info in args_data.parameter_info.items():
if name in parameter_info_val[mode]:
access = AccessKind.READ
else:
access = AccessKind.NONE
assert param_info.access == access
def pybuffer_to_sid(
*,
name: str,
ctype: str,
domain_dim_flags: Tuple[bool, bool, bool],
data_ndim: int,
stride_kind_index: int,
backend: Backend,
check_layout: bool = True,
):
domain_ndim = domain_dim_flags.count(True)
sid_ndim = domain_ndim + data_ndim
as_sid = "as_cuda_sid" if backend.storage_info[
"device"] == "gpu" else "as_sid"
sid_def = """gt::{as_sid}<{ctype}, {sid_ndim},
gt::integral_constant<int, {unique_index}>{unit_stride_dim}>({name})""".format(
name=name,
ctype=ctype,
unique_index=stride_kind_index,
sid_ndim=sid_ndim,
as_sid=as_sid,
unit_stride_dim=
f", {_get_unit_stride_dim(backend, domain_dim_flags, data_ndim)}"
if check_layout else "",
)
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 gtc_utils.dimension_flags_to_names(domain_dim_flags)
]
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)
return sid_def
def visit_AssignStmt(
self,
node: AssignStmt,
*,
loop_order: LoopOrder,
symtable: Dict[str, Any],
**kwargs: Any,
) -> None:
if not isinstance(symtable[node.left.name], self.decl_type):
return None
allowed_flags = self._allowed_flags(loop_order)
flags = symtable[node.left.name].dimensions
if flags not in allowed_flags:
dims = dimension_flags_to_names(flags)
raise ValueError(
f"Not allowed to assign to {dims}-field `{node.left.name}` in {loop_order.name}."
)
return None
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData:
data = ModuleData()
node = pipeline.full()
field_extents = compute_legacy_extents(node)
write_fields = (node.iter_tree().if_isinstance(
gtir.ParAssignStmt).getattr("left").if_isinstance(
gtir.FieldAccess).getattr("name").to_list())
referenced_field_params = {
param.name
for param in node.params if isinstance(param, gtir.FieldDecl)
}
for name in referenced_field_params:
data.field_info[name] = FieldInfo(
access=AccessKind.READ_WRITE
if name in write_fields else AccessKind.READ_ONLY,
boundary=field_extents[name].to_boundary(),
axes=list(
dimension_flags_to_names(
node.symtable_[name].dimensions).upper()),
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()),
)
referenced_scalar_params = set(
node.param_names).difference(referenced_field_params)
for name in referenced_scalar_params:
data.parameter_info[name] = ParameterInfo(
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()))
unref_params = get_unused_params_from_gtir(pipeline)
for param in unref_params:
if isinstance(param, gtir.FieldDecl):
data.field_info[param.name] = None
elif isinstance(param, gtir.ScalarDecl):
data.parameter_info[param.name] = None
data.unreferenced = {param.name for param in unref_params}
return data
def visit_AssignStmt(
self,
node: AssignStmt,
*,
loop_order: LoopOrder,
symtable: Dict[str, Any],
**kwargs: Any,
) -> None:
decl = symtable.get(node.left.name, None)
if decl is None:
raise ValueError("Symbol {} not found.".format(node.left.name))
if not isinstance(decl, self.decl_type):
return None
allowed_flags = self._allowed_flags(loop_order)
flags = decl.dimensions
if flags not in allowed_flags:
dims = dimension_flags_to_names(flags)
raise ValueError(
f"Not allowed to assign to {dims}-field `{node.left.name}` in {loop_order.name}."
)
return None
def visit_FieldDecl(self, node: gtcpp.FieldDecl, **kwargs):
assert "gt_backend_t" in kwargs
if "external_arg" in kwargs:
domain_ndim = node.dimensions.count(True)
data_ndim = len(node.data_dims)
sid_ndim = domain_ndim + data_ndim
if kwargs["external_arg"]:
return "py::buffer {name}, std::array<gt::uint_t,{sid_ndim}> {name}_origin".format(
name=node.name,
sid_ndim=sid_ndim,
)
else:
sid_def = """gt::as_{sid_type}<{dtype}, {sid_ndim},
gt::integral_constant<int, {unique_index}>>({name})""".format(
sid_type="cuda_sid"
if kwargs["gt_backend_t"] == "gpu" else "sid",
name=node.name,
dtype=self.visit(node.dtype),
unique_index=self.unique_index(),
sid_ndim=sid_ndim,
)
if domain_ndim != 3:
gt_dims = [
f"gt::stencil::dim::{dim}" for dim in
gtc_utils.dimension_flags_to_names(node.dimensions)
]
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)
return "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format(
sid_def=sid_def,
name=node.name,
)
def make_args_data_from_gtir(pipeline: GtirPipeline,
legacy=False) -> ModuleData:
"""
Compute module data containing information about stencil arguments from gtir.
Use `legacy` parameter to ensure equality with values from :func:`make_args_data_from_iir`.
"""
data = ModuleData()
node = pipeline.full()
write_fields = (node.iter_tree().if_isinstance(
gtir.ParAssignStmt).getattr("left").if_isinstance(
gtir.FieldAccess).getattr("name").to_set())
read_fields: Set[str] = set()
for expr in node.iter_tree().if_isinstance(
gtir.ParAssignStmt).getattr("right"):
read_fields |= expr.iter_tree().if_isinstance(
gtir.FieldAccess).getattr("name").to_set()
referenced_field_params = [
param.name for param in node.params
if isinstance(param, gtir.FieldDecl)
]
field_extents = compute_legacy_extents(node, mask_inwards=legacy)
k_boundary = (compute_k_boundary(node) if not legacy else
{v: (0, 0)
for v in referenced_field_params})
for name in sorted(referenced_field_params):
access = AccessKind.NONE
if name in read_fields:
access |= AccessKind.READ
if name in write_fields:
access |= AccessKind.WRITE
boundary = Boundary(*field_extents[name].to_boundary()[0:2],
k_boundary[name])
data.field_info[name] = FieldInfo(
access=access,
boundary=boundary,
axes=tuple(
dimension_flags_to_names(
node.symtable_[name].dimensions).upper()),
data_dims=tuple(node.symtable_[name].data_dims),
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()),
)
referenced_scalar_params = [
param.name for param in node.params
if param.name not in referenced_field_params
]
for name in sorted(referenced_scalar_params):
data.parameter_info[name] = ParameterInfo(
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()))
unref_params = get_unused_params_from_gtir(pipeline)
for param in sorted(unref_params, key=lambda decl: decl.name):
if isinstance(param, gtir.FieldDecl):
data.field_info[param.name] = None
elif isinstance(param, gtir.ScalarDecl):
data.parameter_info[param.name] = None
data.unreferenced = [*sorted(param.name for param in unref_params)]
return data