def __call__(
self,
definition_ir: StencilDefinition) -> Dict[str, Dict[str, str]]:
gtir = DefIRToGTIR.apply(definition_ir)
gtir_without_unused_params = prune_unused_parameters(gtir)
dtype_deduced = resolve_dtype(gtir_without_unused_params)
upcasted = upcast(dtype_deduced)
oir = gtir_to_oir.GTIRToOIR().visit(upcasted)
oir = self._optimize_oir(oir)
sdfg = OirSDFGBuilder().visit(oir)
sdfg.expand_library_nodes(recursive=True)
# TODO uncomment once the branch dace/linus-fixes-8 is merged into dace/master
# sdfg.apply_strict_transformations(validate=True) # noqa: E800 Found commented out code
implementation = DaCeComputationCodegen.apply(gtir, sdfg)
bindings = DaCeBindingsCodegen.apply(gtir,
sdfg,
module_name=self.module_name)
bindings_ext = ".cu" if self.backend.GT_BACKEND_T == "gpu" else ".cpp"
return {
"computation": {
"computation.hpp": implementation
},
"bindings": {
"bindings" + bindings_ext: bindings
},
}
def test_stencils_roundtrip(stencil_name):
stencil_def = stencil_registry[stencil_name]
externals = externals_registry[stencil_name]
oir = stencil_def_to_oir(stencil_def, externals)
sdfg = OirSDFGBuilder().visit(oir)
sdfg_pre = deepcopy(sdfg)
oir = convert(sdfg)
sdfg_post = OirSDFGBuilder().visit(oir)
assert_sdfg_equal(sdfg_pre, sdfg_post)
def optimize_horizontal_executions(
stencil: oir.Stencil, transformation: Transformation
) -> oir.Stencil:
sdfg = OirSDFGBuilder().visit(stencil)
api_fields = {param.name for param in stencil.params}
for subgraph in iter_vertical_loop_section_sub_sdfgs(sdfg):
subgraph.apply_transformations_repeated(
transformation, validate=False, options=dict(api_fields=api_fields)
)
return dace_to_oir.convert(sdfg)
def __call__(self, stencil_ir: gtir.Stencil) -> Dict[str, Dict[str, str]]:
base_oir = GTIRToOIR().visit(stencil_ir)
oir_pipeline = self.backend.builder.options.backend_opts.get(
"oir_pipeline",
DefaultPipeline(skip=[MaskInlining]),
)
oir_node = oir_pipeline.run(base_oir)
sdfg = OirSDFGBuilder().visit(oir_node)
_to_device(sdfg, self.backend.storage_info["device"])
sdfg = _expand_and_finalize_sdfg(
stencil_ir, sdfg, self.backend.storage_info["layout_map"])
# strip history from SDFG for faster save/load
for tmp_sdfg in sdfg.all_sdfgs_recursive():
tmp_sdfg.transformation_hist = []
tmp_sdfg.orig_sdfg = None
sources: Dict[str, Dict[str, str]]
implementation = DaCeComputationCodegen.apply(stencil_ir, sdfg)
bindings = DaCeBindingsCodegen.apply(stencil_ir,
sdfg,
module_name=self.module_name,
backend=self.backend)
bindings_ext = "cu" if self.backend.storage_info[
"device"] == "gpu" else "cpp"
sources = {
"computation": {
"computation.hpp": implementation
},
"bindings": {
f"bindings.{bindings_ext}": bindings
},
"info": {
self.backend.builder.module_name + ".sdfg":
dumps(sdfg.to_json())
},
}
return sources
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 optimize_horizontal_executions(
stencil: oir.Stencil, transformation: Transformation
) -> oir.Stencil:
if has_variable_access(stencil):
warnings.warn(
"oir dace optimize_horizontal_executions is not yet supported with variable vertical accesses. See https://github.com/GridTools/gt4py/issues/517"
)
return stencil
sdfg = OirSDFGBuilder().visit(stencil)
api_fields = {param.name for param in stencil.params}
for subgraph in iter_vertical_loop_section_sub_sdfgs(sdfg):
subgraph.apply_transformations_repeated(
transformation, validate=False, options=dict(api_fields=api_fields)
)
return dace_to_oir.convert(sdfg)
def test_same_node_read_write_not_overlap():
oir = StencilFactory(vertical_loops=[
VerticalLoopFactory(sections__0=VerticalLoopSectionFactory(
interval=Interval(start=AxisBound.start(),
end=AxisBound.from_start(1)),
horizontal_executions__0__body__0=AssignStmtFactory(
left__name="field", right__name="other"),
)),
VerticalLoopFactory(sections__0=VerticalLoopSectionFactory(
interval=Interval(start=AxisBound.from_start(1),
end=AxisBound.from_start(2)),
horizontal_executions__0__body__0=AssignStmtFactory(
left__name="field", right__name="field", right__offset__k=-1),
)),
])
sdfg = OirSDFGBuilder().visit(oir)
convert(sdfg, oir.loc)
def test_two_vertical_loops_no_read():
oir_pre = StencilFactory(vertical_loops=[
VerticalLoopFactory(sections__0=VerticalLoopSectionFactory(
horizontal_executions=[
HorizontalExecutionFactory(body__0=AssignStmtFactory(
left__name="field",
right=Literal(value="42.0", dtype=DataType.FLOAT32),
))
],
interval__end=AxisBound.from_start(3),
), ),
VerticalLoopFactory(sections__0=VerticalLoopSectionFactory(
horizontal_executions=[
HorizontalExecutionFactory(body__0=AssignStmtFactory(
left__name="field",
right=Literal(value="43.0", dtype=DataType.FLOAT32),
))
],
interval__start=AxisBound.from_start(3),
), ),
])
sdfg = OirSDFGBuilder().visit(oir_pre)
convert(sdfg, oir_pre.loc)
def optimize_horizontal_executions(
stencil: oir.Stencil, transformation: Transformation) -> oir.Stencil:
sdfg = OirSDFGBuilder().visit(stencil)
for subgraph in iter_vertical_loop_section_sub_sdfgs(sdfg):
subgraph.apply_transformations_repeated(transformation, validate=False)
return dace_to_oir.convert(sdfg)