def generate_computation(self) -> Dict[str, Union[str, Dict]]: computation_name = (self.builder.caching.module_prefix + "computation" + self.builder.caching.module_postfix + ".py") source = NpirCodegen.apply(self.npir) if self.builder.options.format_source: source = format_source("python", source) return {computation_name: source}
def generate_computation(self) -> Dict[str, Union[str, Dict]]: computation_name = (self.builder.caching.module_prefix + "computation" + self.builder.caching.module_postfix + ".py") ignore_np_errstate = self.builder.options.backend_opts.get( "ignore_np_errstate", True) source = NpirCodegen.apply(self.npir, ignore_np_errstate=ignore_np_errstate) if self.builder.options.format_source: source = format_source("python", source) return {computation_name: source}
def test_full_computation_valid(tmp_path) -> None: result = NpirCodegen.apply( npir.Computation( params=["f1", "f2", "f3", "s1"], field_params=["f1", "f2", "f3"], field_decls=[ FieldDeclFactory(name="f1"), FieldDeclFactory(name="f2"), FieldDeclFactory(name="f3"), ], vertical_passes=[ VerticalPassFactory( temp_defs=[], body=[ npir.HorizontalBlock(body=[ VectorAssignFactory( left=FieldSliceFactory(name="f1", parallel_k=True), right=npir.VectorArithmetic( op=common.ArithmeticOperator.MUL, left=FieldSliceFactory(name="f2", parallel_k=True, offsets=(-2, -2, 0)), right=FieldSliceFactory(name="f3", parallel_k=True, offsets=(0, 3, 1)), ), ), ], ), ], ), VerticalPassFactory( lower=common.AxisBound.from_start(offset=1), upper=common.AxisBound.from_end(offset=-3), direction=common.LoopOrder.BACKWARD, temp_defs=[], body=[ npir.HorizontalBlock(body=[ VectorAssignFactory( left__name="f2", right=npir.VectorArithmetic( op=common.ArithmeticOperator.ADD, left=FieldSliceFactory(name="f2", parallel_k=False), right=FieldSliceFactory(name="f2", parallel_k=False, offsets=(0, 0, 1)), ), ), ], ) ], ), ], ), ) print(result) mod_path = tmp_path / "npir_codegen_1.py" mod_path.write_text(result) sys.path.append(str(tmp_path)) import npir_codegen_1 as mod f1 = np.zeros((10, 10, 10)) f2 = np.ones_like(f1) * 3 f3 = np.ones_like(f1) * 2 s1 = 5 mod.run( f1=f1, f2=f2, f3=f3, s1=s1, _domain_=(8, 5, 9), _origin_={ "f1": (2, 2, 0), "f2": (2, 2, 0), "f3": (2, 2, 0) }, ) assert (f1[2:, 2:-3, 0:-1] == 6).all() assert (f1[0:2, :, :] == 0).all() assert (f1[:, 0:2, :] == 0).all() assert (f1[:, -3:, :] == 0).all() assert (f1[:, :, -1:] == 0).all() exp_f2 = np.ones((10)) * 3 # Remember that reversed ranges still include the first (higher) argument and exclude the # second. Thus range(-4, 0, -1) contains the same indices as range(1, -3). exp_f2[-4:0:-1] = np.cumsum(exp_f2[1:-3]) assert (f2[3, 3, :] == exp_f2[:]).all()