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 = 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_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 __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 _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_FieldIfStmt_no_else(): testee = (FieldIfStmtBuilder().cond( FieldAccessBuilder("cond").dtype(DataType.BOOL).build()).build()) 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 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_HorizontalRestriction_HorizontalMask(): testee = HorizontalRestrictionFactory(mask=HorizontalMaskFactory()) GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())
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_ScalarIfStmt(): testee = ScalarIfStmtFactory() 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 test_visit_FieldIfStmt(): testee = FieldIfStmtFactory(false_branch=BlockStmtFactory()) 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_nesting(): testee = FieldIfStmtFactory(true_branch=BlockStmtFactory(body=[FieldIfStmtFactory()])) GTIRToOIR().visit(testee, ctx=GTIRToOIR.Context())