Beispiel #1
0
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
Beispiel #2
0
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 + "]"
Beispiel #3
0
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}]"