int8 4k tensorize works

Author: masahi

tests/python/unittest/test_mma_16x8x32_4k_tune.py

@@ -22,7 +22,9 @@ def ldmatrix_a_desc(a: T.handle, c: T.handle) -> None:
                 v0, v1 = T.axis.remap("SS", [ax0, ax1])
                 T.reads(A_shared[v0, v1])
                 T.writes(A_warp[v0 % 8 * 4 + v1 % 16 // 4, v1 // 16 * 8 + v0 // 8 * 4 + v1 % 4])
-                A_warp[v0 % 8 * 4 + v1 % 16 // 4, v1 // 16 * 8 + v0 // 8 * 4 + v1 % 4] = A_shared[v0, v1]
+                A_warp[v0 % 8 * 4 + v1 % 16 // 4, v1 // 16 * 8 + v0 // 8 * 4 + v1 % 4] = A_shared[
+                    v0, v1
+                ]
 
 
 @T.prim_func
@@ -122,9 +124,19 @@ def mma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
         for i, j, k in T.grid(16, 16, 32):
             with T.block("C"):
                 i, j, k = T.axis.remap("SSR", [i, j, k])
-                T.reads(C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2], A[i % 8 * 4 + k % 16 // 4, k // 16 * 8 + i // 8 * 4 + k % 4], B[j % 8 * 4 + k % 4, j // 8 * 8 + k // 16 * 4 + k % 4])
-                T.writes(C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2])
-                C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2] = C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2] + T.cast(A[i % 8 * 4 + k % 16 // 4, k // 16 * 8 + i // 8 * 4 + k % 4], "int32") * T.cast(B[j % 8 * 4 + k % 4, j // 8 * 8 + k // 16 * 4 + k % 4], "int32")
+                T.reads(
+                    C[i % 8 * 4 + j % 8 // 2, j % 16 // 8 * 4 + i % 16 // 8 * 2 + j % 2],
+                    A[i % 8 * 4 + k % 16 // 4, k % 32 // 16 * 8 + i % 16 // 8 * 4 + k % 4],
+                    B[j % 8 * 4 + k % 4, j % 16 // 8 * 8 + k % 32 // 16 * 4 + k % 4],
+                )
+                T.writes(C[i % 8 * 4 + j % 8 // 2, j % 16 // 8 * 4 + i % 16 // 8 * 2 + j % 2])
+                C[i % 8 * 4 + j % 8 // 2, j % 16 // 8 * 4 + i % 16 // 8 * 2 + j % 8 % 2] = C[
+                    i % 8 * 4 + j % 8 // 2, j % 16 // 8 * 4 + i % 16 // 8 * 2 + j % 8 % 2
+                ] + T.cast(
+                    A[i % 8 * 4 + k % 16 // 4, k % 32 // 16 * 8 + i % 16 // 8 * 4 + k % 4], "int32"
+                ) * T.cast(
+                    B[j % 8 * 4 + k % 4, j % 16 // 8 * 8 + k % 32 // 16 * 4 + k % 4], "int32"
+                )
 
 
 @T.prim_func
@@ -266,6 +278,7 @@ def mma_fill_impl(a: T.handle) -> None:
 M = 4096
 K = 4096
 
+
 def matmul_int8(n, m, k):
     a = te.placeholder((n, k), name="A", dtype="int8")
     b = te.placeholder((k, m), name="B", dtype="int8")
@@ -289,8 +302,8 @@ def schedule(sch: tir.Schedule):
     block = sch.get_block("C")
     i, j, k = sch.get_loops(block)
     i, i_tc = sch.split(i, factors=[None, 16])
-    j, j_tc = sch.split(j, factors=[None, 32])
-    k, k_tc = sch.split(k, factors=[None, 16])
+    j, j_tc = sch.split(j, factors=[None, 16])
+    k, k_tc = sch.split(k, factors=[None, 32])
 
     sch.reorder(
         i,
@@ -311,8 +324,8 @@ def schedule(sch: tir.Schedule):
         num_ty = sch.get(i_factors[2]) * sch.get(j_factors[2])
     else:
         i_factors = [4, 8, 2, 4, 1]
-        j_factors = [1, 32, 2, 1, 2]
-        k_factors = [128, 2, 1]
+        j_factors = [1, 64, 2, 1, 2]
+        k_factors = [64, 2, 1]
 
         num_ty = i_factors[2] * j_factors[2]
 
@@ -381,13 +394,10 @@ def fetch_to_shared(block, idx, ndim):
 
     block_init_c = sch.decompose_reduction(block_outer, sch.get_loops(block_outer)[3])
 
-    def tile_wmma_fragment(block_read, height, is_b=False):
+    def tile_wmma_fragment(block_read, height, width):
         i, j = sch.get_loops(block_read)[-2:]
         i0, i1 = sch.split(i, factors=[None, height])
-        if is_b:
-            j0, j1 = sch.split(j, factors=[32, None])
-        else:
-            j0, j1 = sch.split(j, factors=[None, 32])
+        j0, j1 = sch.split(j, factors=[None, width])
         sch.reorder(i0, j0, i1, j1)
         return i1
 
@@ -411,7 +421,6 @@ def shared_16x32_to_ldmatrix_32x16_layout(i, j):
         thread_id = 4 * (i % 8) + (j % 16) // 4
         return i_0, j_0, thread_id, 8 * (j // 16) + (i // 8) * 4 + j % 4
 
-
     def shared_32x16_to_ldmatrix_32x16_layout(i, j):
         i_0 = i // 32
         j_0 = j // 16
@@ -422,8 +431,8 @@ def shared_32x16_to_ldmatrix_32x16_layout(i, j):
         thread_id = (i % 4) + 4 * (j % 8)
         return i_0, j_0, thread_id, 8 * (j // 8) + (i // 16) * 4 + i % 4
 
-    loop_a = tile_wmma_fragment(A_warp, 16)
-    loop_b = tile_wmma_fragment(B_warp, 16, True)
+    loop_a = tile_wmma_fragment(A_warp, 16, 32)
+    loop_b = tile_wmma_fragment(B_warp, 32, 16)
 
     sch.transform_layout(A_warp, 0, "write", index_map=shared_16x32_to_ldmatrix_32x16_layout)
     sch.transform_layout(B_warp, 0, "write", index_map=shared_32x16_to_ldmatrix_32x16_layout)
@@ -460,44 +469,44 @@ def shared_32x16_to_ldmatrix_32x16_layout(i, j):
 schedule(sch)
 print(sch.mod.script())
 
-# if tune:
-#     with tempfile.TemporaryDirectory() as work_dir:
-#         sch = ms.tune_tir(
-#             mod=workload,
-#             target=tvm.target.Target("nvidia/geforce-rtx-3070"),
-#             config=ms.TuneConfig(
-#                 strategy="evolutionary",
-#                 num_trials_per_iter=32,
-#                 max_trials_per_task=128,
-#                 max_trials_global=128,
-#             ),
-#             work_dir=work_dir,
-#             space=ms.space_generator.ScheduleFn(schedule),
-#         )
-#         if sch is None:
-#             print("No valid schedule found!")
-#         else:
-#             print(sch.mod.script())
-#             print(sch.trace)
-# else:
-#     target = "cuda"
-#     f = tvm.build(sch.mod["main"], target=target, name="dense")
-
-# dev = tvm.device("cuda", 0)
-# a_np = np.random.uniform(size=(N, K)).astype("int8")
-# b_np = np.random.uniform(size=(K, M)).astype("int8")
-# c_np = np.dot(a_np.astype("int32"), b_np.astype("int32"))
-# a = tvm.nd.array(a_np, dev)
-# b = tvm.nd.array(b_np, dev)
-# c = tvm.nd.array(np.zeros((M, N), dtype="int32"), dev)
-# f = tvm.build(sch.mod["main"], target="cuda", name="dense")
-
-# print(f.imported_modules[0].get_source())
-# f(a, b, c)
-# tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
-# print("ok")
-
-# evaluator = f.time_evaluator(f.entry_name, dev, number=1000)
-# gflops = (N * M * K) * 2 / 1e9
-# time_ms = evaluator(a, b, c).mean * 1e3
-# print("matmul with tensor core: %f ms, %f GFLOPS" % (time_ms, gflops / (time_ms / 1e3)))
+if tune:
+    with tempfile.TemporaryDirectory() as work_dir:
+        sch = ms.tune_tir(
+            mod=workload,
+            target=tvm.target.Target("nvidia/geforce-rtx-3070"),
+            config=ms.TuneConfig(
+                strategy="evolutionary",
+                num_trials_per_iter=32,
+                max_trials_per_task=128,
+                max_trials_global=128,
+            ),
+            work_dir=work_dir,
+            space=ms.space_generator.ScheduleFn(schedule),
+        )
+        if sch is None:
+            print("No valid schedule found!")
+        else:
+            print(sch.mod.script())
+            print(sch.trace)
+else:
+    target = "cuda"
+    f = tvm.build(sch.mod["main"], target=target, name="dense")
+
+dev = tvm.device("cuda", 0)
+a_np = np.random.uniform(size=(N, K)).astype("int8")
+b_np = np.random.uniform(size=(K, M)).astype("int8")
+c_np = np.dot(a_np.astype("int32"), b_np.astype("int32"))
+a = tvm.nd.array(a_np, dev)
+b = tvm.nd.array(b_np, dev)
+c = tvm.nd.array(np.zeros((M, N), dtype="int32"), dev)
+f = tvm.build(sch.mod["main"], target="cuda", name="dense")
+
+print(f.imported_modules[0].get_source())
+f(a, b, c)
+tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
+print("ok")
+
+evaluator = f.time_evaluator(f.entry_name, dev, number=1000)
+gflops = (N * M * K) * 2 / 1e9
+time_ms = evaluator(a, b, c).mean * 1e3
+print("matmul with tensor core: %f ms, %f GFLOPS" % (time_ms, gflops / (time_ms / 1e3)))