simplify blocking

tests/python/contrib/test_gemm_acc32_vnni.py

@@ -51,13 +51,13 @@ def verify(target="llvm -mcpu=cascadelake"):
         dev = tvm.device(target, 0)
         pc = dot_16x1x16_uint8_int8_int32_cascadelake()
         ak = te.reduce_axis((0, k), name="k")
-        packedW = te.placeholder((n // 16, 16 * (k // 4), 4), name="packedW", dtype="int8")
+        packedW = te.placeholder((n // 16, k // 4, 16, 4), name="packedW", dtype="int8")
 
         t_fc = te.compute(
             (m, n),
             lambda i, j: te.sum(
                 X[i, ak].astype("int32")
-                * packedW[tvm.tir.indexdiv(j, 16), tvm.tir.indexdiv(ak, 4) * 16 + j % 16, ak % 4].astype("int32"),
+                * packedW[tvm.tir.indexdiv(j, 16), tvm.tir.indexdiv(ak, 4),  j % 16, ak % 4].astype("int32"),
                 axis=ak,
             ),
             name="F",
@@ -69,28 +69,33 @@ def verify(target="llvm -mcpu=cascadelake"):
         a_yo, a_yi = t_sch[t_fc].split(a_y, factor=16)
         a_xo, a_xi = t_sch[t_fc].split(a_x, factor=32)
         a_ko, a_ki = t_sch[t_fc].split(a_k, factor=4)
+
         a_koo, a_koi = t_sch[t_fc].split(a_ko, factor=4)
         t_sch[t_fc].reorder(a_yo, a_xo, a_xi, a_koo, a_koi, a_yi, a_ki)
-
         t_sch[t_fc].unroll(a_koi)
+
+        # a_koo, a_koi = t_sch[t_fc].split(a_ko, factor=4)
+        # t_sch[t_fc].reorder(a_yo, a_xo, a_xi, a_ko, a_yi, a_ki)
+        # t_sch[t_fc].unroll(a_koi)
+
         t_sch[t_fc].tensorize(a_yi, pc)
+        print(tvm.lower(t_sch, [X, packedW, t_fc]))
 
         t_func = tvm.build(t_sch, [X, packedW, t_fc], target, name="intrinsic")
         t_evaluator = t_func.time_evaluator(t_func.entry_name, dev, number=10)
-        print(t_func.get_source("asm"))
+        # print(t_func.get_source("asm"))
 
         # generate the plain data
         a_ = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
         b_ = np.random.uniform(1, 10, size=(n, k)).astype("int8")
 
-        packW = np.random.uniform(1, 10, size=(n // 16, 16 * (k // 4), 4)).astype("int8")
+        packW = np.random.uniform(1, 10, size=(n // 16, (k // 4), 16, 4)).astype("int8")
         # This occurs in pre_compute stage
         for r_idx in range(n // 16):
-            for s_idx in range(16 * (k // 4)):
-                for t_idx in range(4):
-                    packW[r_idx][s_idx][t_idx] = b_[r_idx * 16 + s_idx % 16][
-                        (s_idx // 16) * 4 + t_idx
-                    ]
+            for ko in range(k // 4):
+                for s_idx in range(16):
+                    for t_idx in range(4):
+                        packW[r_idx][ko][s_idx][t_idx] = b_[r_idx * 16 + s_idx][ko * 4 + t_idx]
 
         x = tvm.nd.array(a_, dev)
         w = tvm.nd.array(packW, dev)