def test_assign_with_mask_local() -> None:
result = NpirCodegen().visit(
VectorAssignFactory(
left=LocalScalarAccessFactory(name="tmp"),
mask=FieldSliceFactory(name="mask1", dtype=common.DataType.BOOL),
),
ctx=NpirCodegen.BlockContext(),
symtable={"tmp": ScalarDeclFactory(name="tmp", dtype=common.DataType.INT32)},
)
print(result)
assert re.match(r"tmp = np.where\(mask1.*, np.int32\(\)\)", result) is not None
def test_vector_assign(left, is_serial: bool) -> None:
result = NpirCodegen().visit(
VectorAssignFactory(left=left, right=FieldSliceFactory(name="right")),
ctx=NpirCodegen.BlockContext(),
is_serial=is_serial,
)
left_str, right_str = result.split(" = ")
k_str = "k_:k_+1" if is_serial else "k:K"
if isinstance(left, npir.FieldSlice):
assert left_str == "left[i:I, j:J, " + k_str + "]"
else:
assert left_str == "left"
assert right_str == "right[i:I, j:J, " + k_str + "]"
def test_vector_assign(left, is_serial: bool) -> None:
result = NpirCodegen().visit(
VectorAssignFactory(left=left,
right=FieldSliceFactory(name="right",
k_offset=-1)),
ctx=NpirCodegen.BlockContext(),
is_serial=is_serial,
)
left_str, right_str = result.split(" = ")
if isinstance(left, npir.FieldSlice):
k_str_left = "k_:k_ + 1" if is_serial else "k:K"
else:
k_str_left = ":" if is_serial else "k:K"
k_str_right = "k_ - 1:k_" if is_serial else "k - 1:K - 1"
assert left_str == f"left[i:I, j:J, {k_str_left}]"
assert right_str == f"right[i:I, j:J, {k_str_right}]"