[Metal] Fix review findings: dedup warp partition, kernel_only, barri…

…er comment, immutable gemm ops - Extract shared ComputeSquareWarpPartition to src/backend/metal/op/utils.h, used by both gemm.cc (Square policy) and copy.cc (simdgroup store lowering). - Implement kernel_only parameter in MetalKernelAdapter.get_kernel_source. - Add comment documenting that Metal's simdgroup_barrier synchronizes the full threadgroup (no per-simdgroup barrier exists). - Replace module-level mutable global in metal_fragment_to_simdgroup.py with @lru_cache and remove CUDA-only gemm ops from the set.
LeiWang1999 · May 22, 2026 · Apr 26, 2026 · May 7, 2026 · May 7, 2026 · May 11, 2026
commit 7f7adfde43d259e4b28e288bfb4c5aef554bf6c0
diff --git a/src/backend/metal/op/copy.cc b/src/backend/metal/op/copy.cc
@@ -11,10 +11,6 @@
 
 #include <tvm/tirx/builtin.h>
 
-#include <algorithm>
-#include <cmath>
-#include <limits>
-
 namespace tvm {
 namespace tl {
 
@@ -68,28 +64,8 @@ Stmt LowerSIMDGroupCopy(const CopyNode &op, const LowerArgs &T,
 
   constexpr int kMPerWarp = 8;
   constexpr int kNPerWarp = 8;
-  int m_warp = 1, n_warp = num_warps;
-  int max_m = M / kMPerWarp;
-  int max_n = N / kNPerWarp;
-  float ideal = N > 0 ? static_cast<float>(M) / N : 1.f;
-  float best_score = std::numeric_limits<float>::max();
-  for (int m = 1; m <= std::min(num_warps, max_m); ++m) {
-    if (num_warps % m != 0) {
-      continue;
-    }
-    int n = num_warps / m;
-    if (n > max_n) {
-      continue;
-    }
-    float m_per = static_cast<float>(M) / (m * kMPerWarp);
-    float n_per = static_cast<float>(N) / (n * kNPerWarp);
-    float score = std::abs(m_per / n_per - ideal);
-    if (score < best_score) {
-      best_score = score;
-      m_warp = m;
-      n_warp = n;
-    }
-  }
+  auto [m_warp, n_warp] =
+      ComputeSquareWarpPartition(num_warps, M, N, kMPerWarp, kNPerWarp);
 
   TVM_FFI_ICHECK(M >= m_warp * kMPerWarp && N >= n_warp * kNPerWarp)
       << "Cannot partition " << M << "x" << N << " matrix across " << m_warp

diff --git a/src/backend/metal/op/gemm.cc b/src/backend/metal/op/gemm.cc
@@ -5,12 +5,11 @@
 
 #include "op/gemm.h"
 
+#include "backend/metal/op/utils.h"
 #include "target/utils.h"
 
 #include <tvm/runtime/logging.h>
 
-#include <cmath>
-#include <limits>
 #include <utility>
 
 namespace tvm {
@@ -58,34 +57,8 @@ std::pair<int, int> ComputeMetalWarpPartition(const GemmWarpPolicyNode &policy,
       }
     }
   } else if (policy.isSquare()) {
-    int max_m_warps = M / kMPerWarp;
-    float ideal_ratio = N > 0 ? static_cast<float>(M) / N : 1.0f;
-
-    int best_m = 1;
-    int best_n = 1;
-    float best_balance = std::numeric_limits<float>::max();
-    for (int m = 1; m <= max_m_warps && m <= num_warps; m++) {
-      int n = num_warps / m;
-
-      float m_per_warp = static_cast<float>(M) / (m * kMPerWarp);
-      float n_per_warp = static_cast<float>(N) / (n * kNPerWarp);
-      if (m_per_warp < 1 || n_per_warp < 1) {
-        continue;
-      }
-      if (m * n != num_warps) {
-        continue;
-      }
-
-      float balance = std::abs(m_per_warp / n_per_warp - ideal_ratio);
-      if (balance < best_balance) {
-        best_balance = balance;
-        best_m = m;
-        best_n = n;
-      }
-    }
-
-    m_warp = best_m;
-    n_warp = best_n;
+    std::tie(m_warp, n_warp) =
+        ComputeSquareWarpPartition(num_warps, M, N, kMPerWarp, kNPerWarp);
   } else {
     TVM_FFI_ICHECK(0) << "Unknown GemmWarpPolicy";
   }

diff --git a/src/backend/metal/op/utils.h b/src/backend/metal/op/utils.h
@@ -6,6 +6,10 @@
 #ifndef TVM_TL_BACKEND_METAL_OP_UTILS_H_
 #define TVM_TL_BACKEND_METAL_OP_UTILS_H_
 
+#include <cmath>
+#include <limits>
+#include <utility>
+
 #include "op/utils.h"
 
 namespace tvm {
@@ -20,6 +24,34 @@ inline bool IsRegisterBuffer(const Buffer &buffer) {
   return IsFragmentBuffer(buffer) || IsSIMDGroupBuffer(buffer);
 }
 
+inline std::pair<int, int> ComputeSquareWarpPartition(int num_warps, int M,
+                                                      int N, int kMPerWarp,
+                                                      int kNPerWarp) {
+  int max_m = M / kMPerWarp;
+  int max_n = N / kNPerWarp;
+  float ideal_ratio = N > 0 ? static_cast<float>(M) / N : 1.0f;
+
+  int best_m = 1, best_n = 1;
+  float best_balance = std::numeric_limits<float>::max();
+  for (int m = 1; m <= std::min(num_warps, max_m); ++m) {
+    if (num_warps % m != 0)
+      continue;
+    int n = num_warps / m;
+    if (n > max_n)
+      continue;
+
+    float m_per = static_cast<float>(M) / (m * kMPerWarp);
+    float n_per = static_cast<float>(N) / (n * kNPerWarp);
+    float balance = std::abs(m_per / n_per - ideal_ratio);
+    if (balance < best_balance) {
+      best_balance = balance;
+      best_m = m;
+      best_n = n;
+    }
+  }
+  return {best_m, best_n};
+}
+
 } // namespace metal
 } // namespace tl
 } // namespace tvm

diff --git a/src/target/codegen_metal.cc b/src/target/codegen_metal.cc
@@ -282,6 +282,10 @@ void CodeGenTileLangMetal::PrintType(DataType t,
 void CodeGenTileLangMetal::PrintStorageSync(const CallNode *op) {
   const std::string &sync = op->args[0].as<StringImmNode>()->value;
   if (sync == "warp") {
+    // Metal has no per-simdgroup barrier; simdgroup_barrier synchronizes
+    // the entire threadgroup (same as threadgroup_barrier).  We emit the
+    // narrower intrinsic so the source documents the intended scope even
+    // though the hardware effect is identical.
     this->PrintIndent();
     this->stream << "simdgroup_barrier(mem_flags::mem_threadgroup);\n";
   } else if (sync == "shared") {

diff --git a/tilelang/jit/adapter/torch/metal.py b/tilelang/jit/adapter/torch/metal.py
@@ -54,6 +54,12 @@ def __init__(
     _kernel = None
 
     def get_kernel_source(self, kernel_only: bool = True) -> str:
+        if kernel_only:
+            # Return just the kernel function body, stripping Metal
+            # module-level boilerplate (includes, structs, etc.).
+            idx = self.kernel_global_source.find("kernel void ")
+            if idx >= 0:
+                return self.kernel_global_source[idx:]
         return self.kernel_global_source
 
     def _convert_torch_func(self) -> Callable:

diff --git a/tilelang/transform/metal_fragment_to_simdgroup.py b/tilelang/transform/metal_fragment_to_simdgroup.py
@@ -7,24 +7,20 @@
 
 from __future__ import annotations
 
+from functools import lru_cache
+
 from tvm import tirx as tir
 from tvm import IRModule
 from tvm.ir import Op, PointerType
 from tvm.tirx import SBlock
 from tvm.tirx.transform import prim_func_pass
 
-_GEMM_OPS = None
-
 
+@lru_cache(maxsize=1)
 def _get_gemm_ops():
-    global _GEMM_OPS
-    if _GEMM_OPS is None:
-        _GEMM_OPS = {
-            Op.get("tl.tileop.gemm"),
-            Op.get("tl.tileop.wgmma_gemm"),
-            Op.get("tl.tileop.tcgen05_gemm"),
-        }
-    return _GEMM_OPS
+    return frozenset({
+        Op.get("tl.tileop.gemm"),
+    })
 
 
 def _extract_buffer_var_from_region(region_call):