test_tilelang_gemm_mfma_intrinsic.py 8.06 KB
Newer Older
1
2
3
4
5
6
7
8
9
import torch
import tilelang.testing
from tilelang import tvm as tvm
import tilelang.language as T
from tilelang.intrinsics import make_mfma_swizzle_layout as make_swizzle_layout
from tilelang.intrinsics.mfma_macro_generator import (
    MatrixCoreIntrinEmitter,)
from tilelang.transform import simplify_prim_func

10
tilelang.testing.set_random_seed(0)
11
12
13
14
15
16
17
18
19
20


@simplify_prim_func
def tl_matmul(
    M,
    N,
    K,
    in_dtype,
    out_dtype,
    accum_dtype,
21
22
23
    a_transposed=False,
    b_transposed=True,
    k_pack=1,
24
25
26
27
28
29
30
31
32
33
34
35
36
):
    assert in_dtype in [
        "float16",
        "int8",
    ], "Currently only float16 and int8 are supported"
    assert out_dtype in [
        "float16",
        "float32",
        "int32",
    ], "Currently only float16, float32 and int32 are supported"

    micro_size_x = micro_size_y = micro_size_k = 16

37
    if in_dtype in {"float8_e4m3fnuz", "int8"}:
38
39
        micro_size_k = 32

40
41
42
43
    block_row_warps = 2
    block_col_warps = 2
    warp_row_tiles = 32
    warp_col_tiles = 32
44
45
46

    chunk = 32 * k_pack

47
    shared_scope = "shared"
48
49
50
51
52
53
    cache_write_shared = False

    block_M = block_row_warps * warp_row_tiles
    block_N = block_col_warps * warp_col_tiles
    block_K = chunk

54
55
56
57
    A_shape = (K, M) if a_transposed else (M, K)
    B_shape = (N, K) if b_transposed else (K, N)
    A_shared_shape = (block_K, block_M) if a_transposed else (block_M, block_K)
    B_shared_shape = (block_N, block_K) if b_transposed else (block_K, block_N)
58
    C_shared_shape = (
59
        block_M // micro_size_x,
60
61
62
63
64
65
66
        block_N // micro_size_y,
        micro_size_x,
        micro_size_y,
    )

    warp_size = 64
    threads = warp_size * (block_row_warps * block_col_warps)
67
68
    local_size_a = (k_pack * micro_size_x * micro_size_k) // warp_size
    local_size_b = (k_pack * micro_size_y * micro_size_k) // warp_size
69
70
71
72
73
74
75
76
77
    local_size_c = (micro_size_x * micro_size_y) // warp_size
    warp_rows = warp_row_tiles // micro_size_x
    warp_cols = warp_col_tiles // micro_size_y

    # MMA Wrapper to Auto Generate Code for MMA
    mfma_emitter = MatrixCoreIntrinEmitter(
        a_dtype=in_dtype,
        b_dtype=in_dtype,
        accum_dtype=accum_dtype,
78
79
        a_transposed=a_transposed,
        b_transposed=b_transposed,
80
81
82
83
84
        block_row_warps=block_row_warps,
        block_col_warps=block_col_warps,
        warp_row_tiles=warp_row_tiles,
        warp_col_tiles=warp_col_tiles,
        chunk=chunk,
85
        k_pack=k_pack,
86
87
88
89
    )

    @T.prim_func
    def main(
90
91
92
            A: T.Tensor(A_shape, in_dtype),
            B: T.Tensor(B_shape, in_dtype),
            C: T.Tensor((M, N), out_dtype),
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    ):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):

            A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope)
            B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope)
            C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope)
            A_local = T.alloc_local((warp_rows * local_size_a), in_dtype)
            B_local = T.alloc_local((warp_cols * local_size_b), in_dtype)
            C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)

            T.annotate_layout({
                A_shared: make_swizzle_layout(A_shared),
                B_shared: make_swizzle_layout(B_shared),
            })

            # Improve L2 Cache
            T.use_swizzle(panel_size=10)

            T.clear(C_local)

            for ko in T.Pipelined((K // block_K), num_stages=0):

                # Load A into shared memory
116
117
118
119
                if a_transposed:
                    T.copy(A[ko * block_K, by * block_M], A_shared)
                else:
                    T.copy(A[by * block_M, ko * block_K], A_shared)
120
121

                # Load B into shared memory
122
123
124
125
                if b_transposed:
                    T.copy(B[bx * block_N, ko * block_K], B_shared)
                else:
                    T.copy(B[ko * block_K, bx * block_N], B_shared)
126

127
                for ki in T.serial(0, (block_K // (k_pack * micro_size_k))):
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

                    # Load A into fragment
                    mfma_emitter.ldmatrix_a(
                        A_local,
                        A_shared,
                        ki,
                    )

                    # Load B into fragment
                    mfma_emitter.ldmatrix_b(
                        B_local,
                        B_shared,
                        ki,
                    )

                    # Perform Matrix Multiplication
                    mfma_emitter.mfma(A_local, B_local, C_local)

            # Perform STMatrix
            if cache_write_shared:
                mfma_emitter.stmatrix(
                    C_local,
                    C_shared,
                )

                # Store shared into global
                for i, j in T.Parallel(block_M, block_N):
                    C[by * block_M + i, bx * block_N + j] = C_shared[
                        i // micro_size_x,
                        j // micro_size_y,
                        i % micro_size_x,
                        j % micro_size_y,
                    ]
            else:
                mfma_emitter.stmatrix(
                    C_local,
                    C,
                    pid_m=by,
                    pid_n=bx,
                )

    return main


172
173
174
175
176
177
178
179
180
181
182
183
def assert_tl_matmul_correctness(M,
                                 N,
                                 K,
                                 in_dtype,
                                 out_dtype,
                                 accum_dtype="float32",
                                 a_transposed=False,
                                 b_transposed=True,
                                 k_pack=1):
    matmul = tl_matmul(M, N, K, in_dtype, out_dtype, accum_dtype, a_transposed, b_transposed,
                       k_pack)
    print(matmul)
184
185
    kernel = tilelang.compile(matmul)
    src_code = kernel.get_kernel_source()
186
187
    # src_code is the generated cuda source
    assert src_code is not None
188
189
    A_shape = (K, M) if a_transposed else (M, K)
    B_shape = (N, K) if b_transposed else (K, N)
190
    if in_dtype == "int8":
191
192
        A = torch.randint(-128, 127, A_shape, device="cuda", dtype=torch.int8)
        B = torch.randint(-128, 127, B_shape, device="cuda", dtype=torch.int8)
193
    else:
194
195
        A = torch.rand(A_shape, device="cuda", dtype=getattr(torch, in_dtype))
        B = torch.rand(B_shape, device="cuda", dtype=getattr(torch, in_dtype))
196
197
    C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, out_dtype))

198
    kernel(A, B, C)
199
    print(kernel.get_kernel_source())
200

201
    profiler = kernel.get_profiler()
202

203
    latency = profiler.do_bench()
204
205
206
207

    # Ensure that the latency is not None
    assert latency is not None

208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
    if a_transposed and b_transposed:
        # Get Reference Result
        ref_c = torch.matmul(A.T.to(torch.float32),
                             B.T.to(torch.float32)).to(getattr(torch, out_dtype))
    elif a_transposed and not b_transposed:
        # Get Reference Result
        ref_c = torch.matmul(A.Tto(torch.float32),
                             B.to(torch.float32)).to(getattr(torch, out_dtype))
    elif not a_transposed and b_transposed:
        # Get Reference Result
        ref_c = torch.matmul(A.to(torch.float32),
                             B.T.to(torch.float32)).to(getattr(torch, out_dtype))
    else:
        # Get Reference Result
        ref_c = torch.matmul(A.to(torch.float32), B.to(torch.float32)).to(getattr(torch, out_dtype))

224
225
226
227
228
229
230
231
232
    print(C)
    print(ref_c)
    torch.testing.assert_close(C, ref_c, rtol=1e-2, atol=1e-2)


@tilelang.testing.requires_rocm
def test_assert_tl_matmul():
    assert_tl_matmul_correctness(128, 128, 128, "float16", "float16")
    assert_tl_matmul_correctness(128, 256, 256, "float16", "float32")
233
    assert_tl_matmul_correctness(128, 256, 256, "float16", "float32", k_pack=2)
234
235
236
    assert_tl_matmul_correctness(128, 128, 128, "int8", "int32", accum_dtype="int32")
    assert_tl_matmul_correctness(128, 256, 256, "int8", "int32", accum_dtype="int32")
    assert_tl_matmul_correctness(128, 256, 256, "int8", "int32", accum_dtype="int32", k_pack=2)
237
238
239
240


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