tilelang/testing/python/cpu/test_tilelang_cpu_gemm.py at main · tile-ai/tilelang · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import tilelang
import tilelang.testing
from tilelang import tvm as tvm
import tilelang.cpu.language as T
import torch


def test_matmul_compile():
    def matmul_jit_test(M, N, K, block_M, block_N, block_K, dtype=T.float16, accum_dtype=T.float32):
        # a simple kernel just for jit test
        @T.prim_func
        def matmul(
            A: T.Tensor((M, K), dtype),
            B: T.Tensor((K, N), dtype),
            C: T.Tensor((M, N), dtype),
        ):
            with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M)) as (bx, by):
                A_local = T.alloc_local((block_M, block_K), dtype)
                B_local = T.alloc_local((block_K, block_N), dtype)
                C_local = T.alloc_local((block_M, block_N), accum_dtype)

                for p in T.serial(block_M):
                    for w in T.serial(block_N):
                        C_local[p, w] = 0
                for ko in T.serial(K // block_K):
                    for i in T.serial(block_M):
                        for k in T.serial(block_K):
                            A_local[i, k] = A[by * block_M + i, ko * block_K + k]

                    for k in T.serial(block_K):
                        for j in T.serial(block_N):
                            B_local[k, j] = B[ko * block_K + k, bx * block_N + j]

                    for i in T.serial(block_M):
                        for j in T.serial(block_N):
                            for k in T.serial(block_K):
                                C_local[i, j] += A_local[i, k] * B_local[k, j]

                for i in T.serial(block_M):
                    for j in T.serial(block_N):
                        C[by * block_M + i, bx * block_N + j] = C_local[i, j]

        return matmul

    M, N, K = 1024, 512, 512
    block_M, block_N, block_K = M // 4, N // 4, K // 4
    cpu_func = matmul_jit_test(M, N, K, block_M, block_N, block_K)
    with tvm.target.Target("c"):
        complied_fun = tilelang.compile(cpu_func, -1, execution_backend="cython")

    in_dtype = T.float16
    A = torch.randn(M, K, dtype=torch.__getattribute__(in_dtype))
    B = torch.randn(K, N, dtype=torch.__getattribute__(in_dtype))

    C = complied_fun(A, B)
    C_torch = torch.matmul(A, B)

    tilelang.testing.torch_assert_close(C, C_torch, atol=1e-2, rtol=1e-2, max_mismatched_ratio=0.05)


def test_matmul_with_copy_cython():
    """CPU kernel using T.copy with cython backend.

    Verifies that T.copy works end-to-end on CPU: the vectorized copy
    uses vector types (e.g. float4) defined in common.h, and the
    wrapper correctly skips redundant re-lowering.
    """
    M, N, K = 128, 128, 128
    block_M, block_N, block_K = 32, 32, 32

    @T.prim_func
    def matmul(
        A: T.Tensor((M, K), "float32"),
        B: T.Tensor((K, N), "float32"),
        C: T.Tensor((M, N), "float32"),
    ):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M)) as (bx, by):
            A_local = T.alloc_local((block_M, block_K), "float32")
            B_local = T.alloc_local((block_K, block_N), "float32")
            C_local = T.alloc_local((block_M, block_N), "float32")

            T.clear(C_local)
            for ko in T.serial(K // block_K):
                T.copy(A[by * block_M, ko * block_K], A_local)
                T.copy(B[ko * block_K, bx * block_N], B_local)
                for i, j, k in T.grid(block_M, block_N, block_K):
                    C_local[i, j] += A_local[i, k] * B_local[k, j]
            T.copy(C_local, C[by * block_M, bx * block_N])

    compiled = tilelang.compile(matmul, target="c", out_idx=-1, execution_backend="cython")

    a = torch.randn(M, K, dtype=torch.float32)
    b = torch.randn(K, N, dtype=torch.float32)
    c = compiled(a, b)
    ref = a @ b
    torch.testing.assert_close(c, ref, rtol=1e-5, atol=1e-5)


if __name__ == "__main__":
    tilelang.testing.main()