def test_visit_FieldIfStmt_nesting():
testee = (FieldIfStmtBuilder().cond(
FieldAccessBuilder("cond").dtype(DataType.BOOL).build()).add_true_stmt(
FieldIfStmtBuilder().cond(
FieldAccessBuilder("cond2").dtype(
DataType.BOOL).build()).build()).build())
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_FieldIfStmt():
testee = FieldIfStmtFactory(true_branch__body__0=ParAssignStmtFactory())
mask_stmts = GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
assert len(mask_stmts) == 2
assert "mask" in mask_stmts[0].left.name
assert testee.cond.name == mask_stmts[0].right.name
assert mask_stmts[1].body[0].left.name == testee.true_branch.body[0].left.name
def test_visit_Assign_VariableKOffset():
testee = gtir.ParAssignStmt(
left=FieldAccessFactory(), right=FieldAccessFactory(offset=VariableKOffsetFactory())
)
ctx = GTIRToOIR.Context()
GTIRToOIR().visit(testee, ctx=ctx)
assert len(ctx.horizontal_executions) == 1
assert ctx.horizontal_executions[0].iter_tree().if_isinstance(oir.VariableKOffset).to_list()
def test_visit_ParAssignStmt():
out_name = "out"
in_name = "in"
testee = gtir.ParAssignStmt(left=FieldAccessBuilder(out_name).build(),
right=FieldAccessBuilder(in_name).build())
ctx = GTIRToOIR.Context()
GTIRToOIR().visit(testee, ctx=ctx)
result_horizontal_executions = ctx.horizontal_executions
assert len(result_horizontal_executions) == 1
assign = isinstance_and_return(result_horizontal_executions[0].body[0],
oir.AssignStmt)
left = isinstance_and_return(assign.left, oir.FieldAccess)
right = isinstance_and_return(assign.right, oir.FieldAccess)
assert left.name == out_name
assert right.name == in_name
def test_visit_ParAssignStmt():
out_name = "out"
in_name = "in"
testee = ParAssignStmtFactory(left__name=out_name, right__name=in_name)
assign = GTIRToOIR().visit(testee)
left = isinstance_and_return(assign.left, oir.FieldAccess)
right = isinstance_and_return(assign.right, oir.FieldAccess)
assert left.name == out_name
assert right.name == in_name
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_npir(self) -> npir.Computation:
base_oir = GTIRToOIR().visit(self.builder.gtir)
oir_pipeline = self.builder.options.backend_opts.get(
"oir_pipeline",
DefaultPipeline(skip=[
IJCacheDetection,
KCacheDetection,
PruneKCacheFills,
PruneKCacheFlushes,
FillFlushToLocalKCaches,
]),
)
oir = oir_pipeline.run(base_oir)
return OirToNpir().visit(oir)
def _make_npir(self) -> npir.Computation:
base_oir = GTIRToOIR().visit(self.builder.gtir)
oir_pipeline = self.builder.options.backend_opts.get(
"oir_pipeline",
DefaultPipeline(skip=[
IJCacheDetection,
KCacheDetection,
PruneKCacheFills,
PruneKCacheFlushes,
]),
)
oir_node = oir_pipeline.run(base_oir)
base_npir = OirToNpir().visit(oir_node)
npir_node = ScalarsToTemporaries().visit(base_npir)
return npir_node
def test_visit_gtir_Stencil():
out_name = "out"
in_name = "in"
testee = StencilFactory(vertical_loops__0__body__0=ParAssignStmtFactory(
left__name=out_name, right__name=in_name))
oir_stencil = GTIRToOIR().visit(testee)
hexecs = oir_stencil.vertical_loops[0].sections[0].horizontal_executions
assert len(hexecs) == 1
assert len(hexecs[0].body) == 1
assign = hexecs[0].body[0]
left = isinstance_and_return(assign.left, oir.FieldAccess)
right = isinstance_and_return(assign.right, oir.FieldAccess)
assert left.name == out_name
assert right.name == in_name
def test_create_mask():
mask_name = "mask"
cond = oir_utils.FieldAccessFactory(dtype=DataType.BOOL)
ctx = GTIRToOIR.Context()
result_decl = gtir_to_oir._create_mask(ctx, mask_name, cond)
result_assign = ctx.horizontal_executions[0]
assert isinstance(result_decl, oir.Temporary)
assert result_decl.name == mask_name
horizontal_exec = isinstance_and_return(result_assign, oir.HorizontalExecution)
assign = isinstance_and_return(horizontal_exec.body[0], oir.AssignStmt)
left = isinstance_and_return(assign.left, oir.FieldAccess)
right = isinstance_and_return(assign.right, oir.FieldAccess)
assert left.name == mask_name
assert right == cond
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, 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 __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(skip=[NoFieldAccessPruning])
)
oir_node = oir_pipeline.run(base_oir)
oir_node = FillFlushToLocalKCaches().visit(oir_node)
cuir_node = OIRToCUIR().visit(oir_node)
cuir_node = kernel_fusion.FuseKernels().visit(cuir_node)
cuir_node = extent_analysis.CacheExtents().visit(cuir_node)
format_source = self.backend.builder.options.format_source
implementation = cuir_codegen.CUIRCodegen.apply(cuir_node, format_source=format_source)
bindings = CudaBindingsCodegen.apply(
cuir_node,
module_name=self.module_name,
backend=self.backend,
format_source=format_source,
)
return {
"computation": {"computation.hpp": implementation},
"bindings": {"bindings.cu": bindings},
}
def generate_dace_args(self, ir, sdfg):
oir = GTIRToOIR().visit(ir)
field_extents = compute_fields_extents(oir, add_k=True)
offset_dict: Dict[str, Tuple[int, int, int]] = {
k: (max(-v[0][0], 0), max(-v[1][0], 0), -v[2][0])
for k, v in field_extents.items()
}
k_origins = {
field_name: boundary[0]
for field_name, boundary in compute_k_boundary(ir).items()
}
for name, origin in k_origins.items():
offset_dict[name] = (offset_dict[name][0], offset_dict[name][1],
origin)
symbols = {f"__{var}": f"__{var}" for var in "IJK"}
for name, array in sdfg.arrays.items():
if not array.transient:
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 test_visit_While():
testee = WhileFactory()
GTIRToOIR().visit(testee)
def test_visit_Assign_VariableKOffset():
testee = ParAssignStmtFactory(right__offset=VariableKOffsetFactory())
assign_stmt = GTIRToOIR().visit(testee)
assert assign_stmt.iter_tree().if_isinstance(oir.VariableKOffset).to_list()
def test_visit_FieldIfStmt_no_else():
testee = (FieldIfStmtBuilder().cond(
FieldAccessBuilder("cond").dtype(DataType.BOOL).build()).build())
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_FieldIfStmt_nesting():
testee = FieldIfStmtFactory(true_branch=BlockStmtFactory(body=[FieldIfStmtFactory()]))
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_FieldIfStmt_no_else():
testee = FieldIfStmtFactory(false_branch=None)
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_FieldIfStmt():
testee = FieldIfStmtFactory(false_branch=BlockStmtFactory())
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_ScalarIfStmt():
testee = ScalarIfStmt(
cond=gtir_utils.DummyExpr(dtype=DataType.BOOL, kind=ExprKind.SCALAR),
true_branch=BlockStmt(body=[]),
)
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_HorizontalRestriction_HorizontalMask():
testee = HorizontalRestrictionFactory(mask=HorizontalMaskFactory())
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def test_visit_FieldIfStmt_nesting():
testee = FieldIfStmtFactory(true_branch__body__0=FieldIfStmtFactory())
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
def _make_npir(self) -> npir.Computation:
return OirToNpir().visit(
# TODO (ricoh) apply optimizations, skip only the ones that fail
OirPipeline(GTIRToOIR().visit(self.builder.gtir)).apply([]))
def test_visit_ScalarIfStmt():
testee = ScalarIfStmtFactory()
GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
