update lower warp memory

Author: masahi

src/tir/transforms/lower_warp_memory.cc

@@ -281,6 +281,17 @@ class WarpAccessRewriter : protected StmtExprMutator {
       return GetRef<PrimExpr>(op);
     }
 
+    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);
+      // }
+      return GetRef<PrimExpr>(op);
+    }
+
     return StmtExprMutator::VisitExpr_(op);
   }
 
@@ -466,11 +477,13 @@ namespace transform {
 Pass LowerWarpMemory() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     auto* n = f.CopyOnWrite();
+    // LOG(INFO) << f;
     auto target = f->GetAttr<Target>(tvm::attr::kTarget);
-    int warp_size = target.value()->GetAttr<Integer>("thread_warp_size", 1).value();
+    int warp_size = 32;
     WarpMemoryRewriter warp_memory_rewriter(warp_size);
     auto stmt = warp_memory_rewriter.Rewrite(std::move(n->body));
     n->body = UpdatePointerStorageScope(warp_memory_rewriter.new_storage_scopes_)(stmt);
+    LOG(INFO) << f;
     return f;
   };
   return CreatePrimFuncPass(pass_func, 0, "tir.LowerWarpMemory", {});

tests/python/unittest/test_mma_16x8x8_4k_tune.py

@@ -184,8 +184,8 @@ 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:
-    C_warp = T.match_buffer(a, [32, 4], dtype="float32", scope="warp")
-    C = T.match_buffer(c, [16, 8], dtype="float32", scope="global")
+    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)
 
     with T.block("root"):
         T.reads(C_warp[0:32, 0:4])
@@ -351,9 +351,8 @@ def lambda_b(i, j):
     )
 
     if use_ldmatrix:
-        # sch.tensorize(loop_a, "mma.ldmatrix_a")
-        # sch.tensorize(loop_b, "mma.ldmatrix_b")
-        pass
+        sch.tensorize(loop_a, "mma.ldmatrix_a")
+        sch.tensorize(loop_b, "mma.ldmatrix_b")
     else:
         warp_loop1, warp_loop2 = sch.get_loops(A_warp)[-2:]
         f_0, f_1 = sch.split(warp_loop1, factors=[None, 8])