Function test_crossthread_reduction2

tests/python/codegen/test_target_codegen_cuda.py:376–428

(target, dev)

@tvm.testing.parametrize_targets("cuda", "rocm")
def test_crossthread_reduction2(target, dev):
    def sched(nthdx, nthdy):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(var_A: T.handle, var_B: T.handle):
                T.func_attr({"tirx.noalias": True})
                n, k0, k1 = T.int32(), T.int32(), T.int32()
                A = T.match_buffer(var_A, (n, k0, k1))
                B = T.match_buffer(var_B, (n,))
                for i in T.thread_binding(n, thread="blockIdx.x"):
                    for k0_0 in T.thread_binding(nthdx, thread="threadIdx.x"):
                        for k0_1 in range((k0 + nthdx - 1) // nthdx):
                            for k1_0 in T.thread_binding(nthdy, thread="threadIdx.y"):
                                for k1_1 in range((k1 + nthdy - 1) // nthdy):
                                    with T.sblock("B"):
                                        v_i = T.axis.spatial(n, i)
                                        v_k0 = T.axis.reduce(
                                            k0, k0_0 * ((k0 + nthdx - 1) // nthdx) + k0_1
                                        )
                                        v_k1 = T.axis.reduce(
                                            k1, k1_0 * ((k1 + nthdy - 1) // nthdy) + k1_1
                                        )
                                        T.where(
                                            k0_0 * ((k0 + nthdx - 1) // nthdx) + k0_1 < k0
                                            and k1_0 * ((k1 + nthdy - 1) // nthdy) + k1_1 < k1
                                        )
                                        T.reads(A[v_i, v_k0, v_k1])
                                        T.writes(B[v_i])
                                        with T.init():
                                            B[v_i] = T.float32(0.0)
                                        B[v_i] = B[v_i] + A[v_i, v_k0, v_k1]

        func = tvm.compile(Module, target="cuda")
        return func

    def verify(nthdx, nthdy):
        func = sched(nthdx, nthdy)
        nn = 3
        # checks three typical cases
        vx = [nthdx - 1, nthdx, nthdx + 1]
        vy = [nthdy - 1, nthdy, nthdy + 1]
        for kk0, kk1 in [(x, y) for x in vx for y in vy]:
            size = (nn, kk0, kk1)
            a = tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
            b = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
            func(a, b)
            tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=(1, 2)), rtol=1e-3)

    verify(16, 16)
    verify(32, 32)
    verify(16, 32)
    verify(32, 16)


@pytest.mark.gpu