def __call__(
self,
definition_ir: StencilDefinition) -> Dict[str, Dict[str, str]]:
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
base_oir = gtir_to_oir.GTIRToOIR().visit(gtir)
oir_pipeline = self.backend.builder.options.backend_opts.get(
"oir_pipeline",
DefaultPipeline(skip=[MaskStmtMerging, MaskInlining]),
)
oir = oir_pipeline.run(base_oir)
sdfg = OirSDFGBuilder().visit(oir)
sdfg.expand_library_nodes(recursive=True)
sdfg.apply_strict_transformations(validate=True)
implementation = DaCeComputationCodegen.apply(gtir, sdfg)
bindings = DaCeBindingsCodegen.apply(gtir,
sdfg,
module_name=self.module_name,
backend=self.backend)
bindings_ext = ".cu" if self.backend.GT_BACKEND_T == "gpu" else ".cpp"
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings" + bindings_ext: bindings
},
}
def __call__(self, definition_ir) -> Dict[str, Dict[str, str]]:
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
base_oir = gtir_to_oir.GTIRToOIR().visit(gtir)
oir_pipeline = self.backend.builder.options.backend_opts.get(
"oir_pipeline", DefaultPipeline(skip=[FillFlushToLocalKCaches]))
oir = oir_pipeline.run(base_oir)
gtcpp = oir_to_gtcpp.OIRToGTCpp().visit(oir)
format_source = self.backend.builder.options.format_source
implementation = gtcpp_codegen.GTCppCodegen.apply(
gtcpp,
gt_backend_t=self.backend.GT_BACKEND_T,
format_source=format_source)
bindings = GTCppBindingsCodegen.apply(gtcpp,
module_name=self.module_name,
backend=self.backend,
format_source=format_source)
bindings_ext = ".cu" if self.backend.GT_BACKEND_T == "gpu" else ".cpp"
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings" + bindings_ext: bindings
},
}
def __call__(self, definition_ir) -> Dict[str, Dict[str, str]]:
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
base_oir = gtir_to_oir.GTIRToOIR().visit(gtir)
oir_pipeline = self.backend.builder.options.backend_opts.get(
"oir_pipeline", DefaultPipeline(skip=[NoFieldAccessPruning]))
oir = oir_pipeline.run(base_oir)
oir = FillFlushToLocalKCaches().visit(oir)
cuir = oir_to_cuir.OIRToCUIR().visit(oir)
cuir = kernel_fusion.FuseKernels().visit(cuir)
cuir = extent_analysis.CacheExtents().visit(cuir)
format_source = self.backend.builder.options.format_source
implementation = cuir_codegen.CUIRCodegen.apply(
cuir, format_source=format_source)
bindings = GTCCudaBindingsCodegen.apply(cuir,
module_name=self.module_name,
backend=self.backend,
format_source=format_source)
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings.cu": bindings
},
}
def gtir_pipeline(self) -> GtirPipeline:
return self._build_data.get(
"gtir_pipeline") or self._build_data.setdefault(
"gtir_pipeline",
GtirPipeline(
self.frontend.generate(self.definition, self.externals,
self.options)),
)
def stencil_def_to_oir(stencil_def, externals):
build_options = BuildOptions(
name=stencil_def.__name__, module=__name__, rebuild=True, backend_opts={}, build_info=None
)
definition_ir = GTScriptFrontend.generate(
stencil_def, externals=externals, options=build_options
)
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
return GTIRToOIR().visit(gtir)
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 __call__(self, definition_ir) -> Dict[str, Dict[str, str]]:
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
oir = gtir_to_oir.GTIRToOIR().visit(gtir)
oir = self._optimize_oir(oir)
gtcpp = oir_to_gtcpp.OIRToGTCpp().visit(oir)
implementation = gtcpp_codegen.GTCppCodegen.apply(
gtcpp, gt_backend_t=self.gt_backend_t)
bindings = GTCppBindingsCodegen.apply(gtcpp,
module_name=self.module_name,
gt_backend_t=self.gt_backend_t)
bindings_ext = ".cu" if self.gt_backend_t == "gpu" else ".cpp"
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings" + bindings_ext: bindings
},
}
def __call__(self, definition_ir) -> Dict[str, Dict[str, str]]:
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
oir = OirPipeline(gtir_to_oir.GTIRToOIR().visit(gtir)).full(
skip=[FillFlushToLocalKCaches])
gtcpp = oir_to_gtcpp.OIRToGTCpp().visit(oir)
implementation = gtcpp_codegen.GTCppCodegen.apply(
gtcpp, gt_backend_t=self.backend.GT_BACKEND_T)
bindings = GTCppBindingsCodegen.apply(gtcpp,
module_name=self.module_name,
backend=self.backend)
bindings_ext = ".cu" if self.backend.GT_BACKEND_T == "gpu" else ".cpp"
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings" + bindings_ext: bindings
},
}
def __call__(self, definition_ir) -> Dict[str, Dict[str, str]]:
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
oir = OirPipeline(gtir_to_oir.GTIRToOIR().visit(gtir)).full(
skip=[NoFieldAccessPruning])
cuir = oir_to_cuir.OIRToCUIR().visit(oir)
cuir = kernel_fusion.FuseKernels().visit(cuir)
cuir = extent_analysis.ComputeExtents().visit(cuir)
cuir = extent_analysis.CacheExtents().visit(cuir)
implementation = cuir_codegen.CUIRCodegen.apply(cuir)
bindings = GTCCudaBindingsCodegen.apply(cuir,
module_name=self.module_name,
backend=self.backend)
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings.cu": bindings
},
}
def _expand_and_finalize_sdfg(stencil_ir: gtir.Stencil, sdfg: dace.SDFG,
layout_map) -> dace.SDFG:
args_data = make_args_data_from_gtir(GtirPipeline(stencil_ir))
# stencils without effect
if all(info is None for info in args_data.field_info.values()):
sdfg = dace.SDFG(stencil_ir.name)
sdfg.add_state(stencil_ir.name)
return sdfg
for array in sdfg.arrays.values():
if array.transient:
array.lifetime = dace.AllocationLifetime.Persistent
_pre_expand_trafos(sdfg)
sdfg.expand_library_nodes(recursive=True)
_specialize_transient_strides(sdfg, layout_map=layout_map)
_post_expand_trafos(sdfg)
return sdfg
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 __call__(self, stencil_ir: gtir.Stencil) -> Dict[str, Dict[str, str]]:
stencil_ir = GtirPipeline(stencil_ir).full()
base_oir = GTIRToOIR().visit(stencil_ir)
oir_pipeline = self.backend.builder.options.backend_opts.get(
"oir_pipeline", DefaultPipeline()
)
oir_node = oir_pipeline.run(base_oir)
gtcpp_ir = OIRToGTCpp().visit(oir_node)
format_source = self.backend.builder.options.format_source
implementation = gtcpp_codegen.GTCppCodegen.apply(
gtcpp_ir, gt_backend_t=self.backend.GT_BACKEND_T, format_source=format_source
)
bindings = GTCppBindingsCodegen.apply(
gtcpp_ir,
module_name=self.module_name,
backend=self.backend,
format_source=format_source,
)
bindings_ext = ".cu" if self.backend.GT_BACKEND_T == "gpu" else ".cpp"
return {
"computation": {"computation.hpp": implementation},
"bindings": {"bindings" + bindings_ext: bindings},
}
def gtir_is_not_emtpy(pipeline: GtirPipeline) -> bool:
node = pipeline.full()
return bool(node.iter_tree().if_isinstance(gtir.ParAssignStmt).to_list())
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
def gtir_pipeline(self) -> GtirPipeline:
return self._build_data.get(
"gtir_pipeline") or self._build_data.setdefault(
"gtir_pipeline",
GtirPipeline(DefIRToGTIR.apply(self.definition_ir)))