[Backend] Support TMA lowering for arbitrary (swizzled) SMEM layout (#2380)

Author: Yongqi-Zhuo

examples/gemm/example_gemm_intrinsics.py

@@ -17,8 +17,7 @@ def make_swizzle_layout(shared_buf):
         return T.Layout(shape, lambda *args: args)
 
     def transform_func(i, j):
-        new_warp_i, new_warp_j = get_swizzle_layout(i, j, shape[-1], dtype)
-        return [new_warp_i, new_warp_j]
+        return list(get_swizzle_layout(i, j, shape[-1], dtype))
 
     return T.Layout(shape, transform_func)
 

src/cuda/op/copy.cc

@@ -34,6 +34,7 @@
 
 #include "backend/common/target_utils.h"
 #include "cuda/op/copy.h"
+#include "layout/cute_layout.h"
 #include "multi_version_buffer_rewriter.h"
 #include "op/builtin.h"
 #include "op/copy.h"
@@ -2620,15 +2621,14 @@ class ManualWSDetector : public StmtExprVisitor {
 /// swizzle modes (32B / 64B / 128B).  Any other layout (e.g. padded,
 /// Volta-style) cannot be used with TMA.
 static bool IsTmaCompatibleLayout(const Layout &layout, const Buffer &buffer) {
-  // Recognised swizzle → TMA with swizzle.
-  if (DetectSwizzleMode(layout, buffer) != SwizzleMode::kNone) {
-    return true;
-  }
-  // Identity / row-major linear → TMA without swizzle.
-  if (StructuralEqual()(layout, makeLinearLayout(buffer->shape))) {
-    return true;
-  }
-  return false;
+  Optional<cute::ComposedLayout> composed =
+      cute::ComposedLayoutFromTileLang(layout);
+  if (!composed.defined())
+    return false;
+  // Recast to byte space (the swizzle atom is defined on byte addresses).
+  cute::ComposedLayout composed_bytes =
+      composed.value().Recast(buffer->dtype.bits(), /*new_bits=*/8);
+  return composed_bytes->swizzle->IsTMACompatible();
 }
 
 class TiledWSCandidate : public StmtExprVisitor {