mma store codegen working

Author: masahi

src/target/source/codegen_cuda.cc

@@ -821,6 +821,17 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) {
     std::string smem_elem_offset = this->PrintExpr(op->args[6]);
     this->stream << PrintLoadMatrixAssembly(trans, num, type, local_ptr, local_elem_offset,
                                             smem_ptr, smem_elem_offset);
+  } else if (op->op.same_as(builtin::mma_store())) {
+    std::string dst = this->PrintExpr(op->args[1]);
+    std::string src = this->PrintExpr(op->args[2]);
+    std::string src_offset = this->PrintExpr(op->args[3]);
+    std::string stride = this->PrintExpr(op->args[4]);
+
+    os << "for (int i = 0; i < 4; ++i) {\n";
+    os << dst << "[(i / 2 * 8 + threadIdx.x / 4) * " << stride
+       << " + (threadIdx.x % 4) * 2 + i % 2]"
+       << " = " << src << "[" << src_offset << " + i];\n";
+    os << "}\n";
   } else {
     CodeGenC::VisitExpr_(op, os);
   }

src/tir/transforms/lower_warp_memory.cc

@@ -282,13 +282,13 @@ class WarpAccessRewriter : protected StmtExprMutator {
     }
 
     if (op->op.same_as(builtin::mma_store())) {
-      // Array<PrimExpr> new_args = op->args;
-      // PrimExpr local_index, group;
-      // if (op->args[3].get() == buffer_) {
-      //   std::tie(local_index, group) = SplitIndexByGroup(op->args[4]);
-      //   new_args.Set(4, local_index);
-      //   return Call(op->dtype, op->op, new_args);
-      // }
+      Array<PrimExpr> new_args = op->args;
+      PrimExpr local_offset, group;
+      if (op->args[2].get() == buffer_) {
+        std::tie(local_offset, group) = SplitIndexByGroup(op->args[3]);
+        new_args.Set(3, local_offset);
+        return Call(op->dtype, op->op, new_args);
+      }
       return GetRef<PrimExpr>(op);
     }
 

tests/python/unittest/test_mma_16x8x8_4k_tune.py

@@ -184,16 +184,19 @@ def mma_store_desc(a: T.handle, c: T.handle) -> None:
 
 @T.prim_func
 def mma_store_impl(a: T.handle, c: T.handle) -> None:
+    s1 = T.var("int32")
+    s0 = T.var("int32")
+
     C_warp = T.match_buffer(a, [32, 4], dtype="float32", scope="warp", offset_factor=1)
-    C = T.match_buffer(c, [16, 8], dtype="float32", scope="global",offset_factor=1)
+    C = T.match_buffer(c, [16, 8], dtype="float32", scope="global",offset_factor=1, strides=[s1, s0])
 
     with T.block("root"):
         T.reads(C_warp[0:32, 0:4])
         T.writes(C[0:16, 0:8])
         tx = T.env_thread("threadIdx.x")
         T.launch_thread(tx, 32)
 
-        T.evaluate(T.mma_store("m16n8", C.data, C.elem_offset, C_warp.access_ptr("r"), tx, dtype="float32"))
+        T.evaluate(T.mma_store("m16n8", C.access_ptr("w"), C_warp.data, C_warp.elem_offset, s1, dtype="float32"))
 
 
 tir.TensorIntrin.register("mma.ldmatrix_a", ldmatrix_a_desc, ldmatrix_a_impl)
@@ -388,7 +391,6 @@ def lambda_b(i, j):
     fused_2 = sch.fuse(f_0, f_3)
 
     # print(sch.mod.script())
-
     # return
 
     sch.tensorize(fused_1, "mma_store")
@@ -423,20 +425,20 @@ def lambda_b(i, j):
     print(sch.mod.script())
     target = "cuda"
     f = tvm.build(sch.mod["main"], target=target, name="dense")
-    print(f.imported_modules[0].get_source())
-
-# dev = tvm.device("cuda", 0)
-# a_np = np.random.uniform(size=(N, K)).astype("float16")
-# b_np = np.random.uniform(size=(K, M)).astype("float16")
-# c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
-# a = tvm.nd.array(a_np, dev)
-# b = tvm.nd.array(b_np, dev)
-# c = tvm.nd.array(np.zeros((M, N), dtype="float32"), dev)
-# f = tvm.build(sch.mod["main"], target="cuda", name="dense")
-
-# f(a, b, c)
-# tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
-# print("ok")
+
+dev = tvm.device("cuda", 0)
+a_np = np.random.uniform(size=(N, K)).astype("float16")
+b_np = np.random.uniform(size=(K, M)).astype("float16")
+c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
+a = tvm.nd.array(a_np, dev)
+b = tvm.nd.array(b_np, dev)
+c = tvm.nd.array(np.zeros((M, N), dtype="float32"), 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