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()