Function test_llvm_div

tests/python/codegen/test_target_codegen_llvm.py:562–698

Check that the semantics of div and mod is correct

()

@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_llvm_div():
    """Check that the semantics of div and mod is correct"""

    def check(start, end, dstart, dend, dtype, floor_div=False):
        a_size = end - start + 1
        b_size = dend - dstart + 1

        div_fn = tvm.tirx.floordiv if floor_div else tvm.tirx.truncdiv
        mod_fn = tvm.tirx.floormod if floor_div else tvm.tirx.truncmod

        # Build clipping helpers — capture TIR const values from env
        _start = tvm.tirx.const(start, dtype)
        _end = tvm.tirx.const(end, dtype)
        _dstart = tvm.tirx.const(dstart, dtype)
        _dend = tvm.tirx.const(dend, dtype)

        if start == end:
            clipa = lambda x: _start
        else:
            clipa = lambda x: T.min(_end, T.max(_start, x))

        if dstart == dend:
            clipb = lambda x: _dstart
        else:
            clipb = lambda x: T.min(_dend, T.max(_dstart, x))

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(
                A: T.Buffer((a_size,), dtype),
                B: T.Buffer((b_size,), dtype),
                D: T.Buffer((a_size, b_size), dtype),
                M: T.Buffer((a_size, b_size), dtype),
            ):
                T.func_attr({"tirx.noalias": True})
                for i, j in T.grid(a_size, b_size):
                    with T.sblock("D"):
                        v_i, v_j = T.axis.remap("SS", [i, j])
                        T.reads(A[v_i], B[v_j])
                        T.writes(D[v_i, v_j])
                        D[v_i, v_j] = div_fn(clipa(A[v_i]), clipb(B[v_j]))
                    with T.sblock("M"):
                        v_i, v_j = T.axis.remap("SS", [i, j])
                        T.reads(A[v_i], B[v_j])
                        T.writes(M[v_i, v_j])
                        M[v_i, v_j] = mod_fn(clipa(A[v_i]), clipb(B[v_j]))

        f = tvm.compile(Module, target="llvm")

        # Fill input arrays with values
        A_arr = tvm.runtime.empty((a_size,), dtype)
        B_arr = tvm.runtime.empty((b_size,), dtype)
        A_arr.copyfrom(np.arange(start, end + 1, dtype=dtype))
        B_np = np.arange(dstart, dend + 1, dtype=dtype)
        # If the range of the divisor contains 0, replace it with 1 to avoid division by zero
        if dend >= 0 and dstart <= 0:
            B_np[-dstart] = 1