Refactor GEMM layout and Python integration for improved functionality

- Updated `gemm_layouts.cc` to correct the order of layout replication and repetition for transposed and non-transposed GEMM operations.
- Enhanced `gemm_py.cc` to handle block realization more robustly, ensuring correct assignment of global symbols and block attributes.
- Refactored `inject_pipeline.cc` to streamline buffer read/write region handling, improving clarity and maintainability.
- Cleaned up test cases in `test_tilelang_tilelibrary_gemm.py` by removing unnecessary print statements and adjusting function calls for better test execution flow.

These changes enhance the clarity, functionality, and robustness of GEMM operations and their testing in the TileLang framework.

Author: LeiWang1999

src/layout/gemm_layouts.cc

@@ -205,14 +205,14 @@ Fragment makeGemmFragmentB(const int block_m, const int block_n,
   ICHECK(block_k % 16 == 0);
   if (transposed) {
     auto base_layout = makeGemmFragment8x8()->Repeat({1, 2}, false, false);
-    auto warp_layout = base_layout->Repeat({block_n / warp_n, 1}, true, false)->Replicate(block_m / warp_m);
+    auto warp_layout = base_layout->Replicate(block_m / warp_m)->Repeat({block_n / warp_n, 1}, true, false);
     auto block_layout =
         warp_layout->Repeat({warp_n / 8, block_k / 16}, false, false);
     return block_layout;
   } else {
     auto base_layout =
         makeGemmFragment8x8Transposed()->Repeat({2, 1}, false, false);
-    auto warp_layout = base_layout->Repeat({1, block_n / warp_n}, true)->Replicate(block_m / warp_m);
+    auto warp_layout = base_layout->Replicate(block_m / warp_m)->Repeat({1, block_n / warp_n}, true);
     auto block_layout =
         warp_layout->Repeat({block_k / 16, warp_n / 8}, false, true);
     return block_layout;

src/op/gemm_py.cc

@@ -227,16 +227,25 @@ Stmt GemmPyNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
   if (const auto f = ffi::Function::GetGlobal("tl.gemm_py.lower")) {
     auto prim_func = Downcast<PrimFunc>(
         (*f)(GetRef<GemmPy>(this), T.target, T.thread_bounds, T.thread_var));
-    BlockRealize block_realize = Downcast<BlockRealize>(prim_func->body);
     ICHECK(prim_func->attrs.defined());
     auto global_symbol = prim_func->attrs.GetAttr<String>("global_symbol");
     ICHECK(global_symbol.defined());
-    auto block = block_realize->block;
-    {
-      BlockNode* n = block.CopyOnWrite();
-      n->name_hint = global_symbol.value();
+    if (prim_func->body.as<BlockRealizeNode>()) {
+      BlockRealize block_realize = Downcast<BlockRealize>(prim_func->body);
+      auto block = block_realize->block;
+      {
+        BlockNode* n = block.CopyOnWrite();
+        n->name_hint = global_symbol.value();
+      }
+      return BlockRealize(block_realize->iter_values, block_realize->predicate, block);
     }
-    return BlockRealize(block_realize->iter_values, block_realize->predicate, block);
+    // warp with block realize node
+    return BlockRealize(
+      /*iter_values=*/Array<PrimExpr>(),
+      /*predicate=*/const_true(),
+      /*block=*/
+      Block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
+            /*name_hint=*/global_symbol.value(), prim_func->body)); 
   } else {
     LOG(FATAL) << "No lower function found for gemm_py";
   }

src/transform/inject_pipeline.cc

@@ -248,7 +248,6 @@ class PipelineRewriter : public StmtExprMutator {
         buffer_remap_.Set(buffer, RewriteAllocBuffer(buffer, num_versions));
       }
     }
-
     ordered_stmts_.resize(pipeline_info_.size());
     for (const auto &[block, anno] : pipeline_info_) {
       ordered_stmts_.Set(anno.order, block);
@@ -676,11 +675,6 @@ class PipelineRewriter : public StmtExprMutator {
       new_block = Downcast<Block>(Substitute(
           new_block, {{pipeline_loop_->loop_var, normalized_access_index}}));
 
-      Array<Array<BufferRegion>> access = GetBlockReadWriteRegion(block, buffer_data_to_buffer_);
-      BlockNode* n = new_block.CopyOnWrite();
-      n->reads = access[0];
-      n->writes = access[1];
-
       if (pipeline_info_[block].async) {
         auto &local_state = async_states_local[stage];
         local_state.producer_head = normalized_access_index;
@@ -957,6 +951,11 @@ class PipelineInjector : private StmtExprMutator {
 
     Block block = Downcast<Block>(StmtExprMutator::VisitStmt_(op));
 
+    Array<Array<BufferRegion>> access = GetBlockReadWriteRegion(block, buffer_data_to_buffer_);
+    BlockNode* n = block.CopyOnWrite();
+    n->reads = access[0];
+    n->writes = access[1];
+
     for (const auto &buffer : op->alloc_buffers) {
       buffer_data_to_buffer_.erase(buffer->data);
     }

testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py

@@ -90,7 +90,6 @@ def run_gemm_ss(
             tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
             tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
         })
-    print(kernel.get_kernel_source())
     profiler = kernel.get_profiler()
 
     def ref_program(A, B):
@@ -209,7 +208,6 @@ def run_gemm_rs(
             tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
             tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
         })
-    print(kernel.get_kernel_source())
     profiler = kernel.get_profiler()
 
     def ref_program(A, B):
@@ -280,10 +278,7 @@ def main(
                 else:
                     T.copy(B[k * block_K, bx * block_N], B_shared)
                 T.copy(B_shared, B_frag)
-                # for i, j in T.Parallel(block_N, block_K):
-                #     B_frag[i, j] = B_shared[j, i]
-                # T.gemm_v2(A_shared, B_frag, C_local, trans_A, trans_B)
-                T.gemm(A_shared, B_frag, C_local, trans_A, trans_B)
+                T.gemm_v2(A_shared, B_frag, C_local, trans_A, trans_B)
             T.copy(C_local, C[by * block_M, bx * block_N])
 
     return main
@@ -327,7 +322,6 @@ def run_gemm_sr(
             tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
             tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
         })
-    print(kernel.get_kernel_source())
     profiler = kernel.get_profiler()
 
     def ref_program(A, B):
@@ -448,7 +442,6 @@ def run_gemm_rr(
             tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
             tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
         })
-    print(kernel.get_kernel_source())
     profiler = kernel.get_profiler()
 
     def ref_program(A, B):
@@ -478,9 +471,13 @@ def test_gemm_rr():
     # tilelang.testing.main()
     tilelang.disable_cache()
     # test_gemm_ss()
-    run_gemm_sr(128, 128, 128, False, False, "float16", "float16", "float16", 128, 128, 32, 2)
+    # test_gemm_sr()
+    # test_gemm_rs()
+    # test_gemm_rr()
+    
+    # run_gemm_sr(128, 128, 128, False, False, "float16", "float16", "float16", 128, 128, 32, 2)
     # tilelang.testing.set_random_seed(42)
-    # run_gemm_ss(128, 128, 128, False, True, "float16", "float16", "float16", 128, 128, 32, 1)
+    run_gemm_ss(128, 128, 128, False, True, "float16", "float16", "float16", 128, 128, 32, 1)
     # print("gemm fp16 nt ss done")
     # exit()
 

tilelang/engine/phase.py

@@ -140,11 +140,7 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
         mod = tilelang.transform.IfStmtBinding()(mod)
         mod = tir.transform.PlanAndUpdateBufferAllocationLocation()(mod)
         mod = tilelang.transform.PipelinePlanning()(mod)
-        print("after pipeline planning")
-        print(mod)
         mod = tilelang.transform.InjectSoftwarePipeline()(mod)
-        print("after inject software pipeline")
-        print(mod)
         mod = tilelang.transform.MergeIfStmt()(mod)
         if allow_fence_proxy(target=target):
             # in hopper device, wgmma is an async proxy