[CUDA][Transform] Expose SIMT im2col loads before scheduling

src/cuda/transform/expand_simt_im2col.cc

@@ -0,0 +1,145 @@
+/*!
+ * \file expand_simt_im2col.cc
+ * \brief Early exposure of CUDA SIMT im2col loads before scheduling.
+ */
+
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/transform.h>
+
+#include "arith/ir_mutator_with_analyzer.h"
+#include "cuda/target_utils.h"
+#include "op/copy.h"
+#include "op/operator.h"
+#include "op/utils.h"
+
+namespace tvm {
+namespace tl {
+
+using namespace tirx;
+using namespace ffi;
+
+TVM_REGISTER_PASS_CONFIG_OPTION("tl.enable_early_simt_im2col", Bool);
+
+namespace {
+
+// Keep specialized implementations (such as Hopper TMA) and explicit operation
+// annotations intact. Only expand the selected generic SIMT implementation.
+class SIMTIm2ColExpander : public arith::IRMutatorWithAnalyzer {
+public:
+  explicit SIMTIm2ColExpander(arith::Analyzer *analyzer)
+      : arith::IRMutatorWithAnalyzer(analyzer) {}
+
+private:
+  Stmt VisitStmt_(const SBlockNode *node) final {
+    // Explicit WS schedules refer to the original operations, not new loops.
+    if (node->annotations.count("tl.ws_schedule")) {
+      return GetRef<SBlock>(node);
+    }
+    return arith::IRMutatorWithAnalyzer::VisitStmt_(node);
+  }
+
+  Stmt VisitStmt_(const EvaluateNode *node) final {
+    Stmt stmt = arith::IRMutatorWithAnalyzer::VisitStmt_(node);
+    const auto *call = stmt.as<EvaluateNode>()->value.as<CallNode>();
+    if (!call || !call->op.same_as(Im2ColOp::Get()) ||
+        !call->annotations.empty()) {
+      return stmt;
+    }
+    Im2ColOp op = Downcast<Im2ColOp>(ParseOperator(GetRef<Call>(call)));
+    const Buffer &src = op->src_;
+    const Buffer &dst = op->dst_;
+    const auto &region = op->dstRegion_->region;
+    if (!IsGlobalBuffer(src) || !IsSharedBuffer(dst) ||
+        src->shape.size() != 4 || region.size() != 2 ||
+        src->dtype != dst->dtype || op->kernel_ <= 0 || op->stride_ <= 0 ||
+        op->dilation_ <= 0) {
+      return stmt;
+    }
+    // Do not reinterpret cropped source regions: the fallback currently uses
+    // whole-buffer NHWC coordinates. Leave unsupported metadata untouched.
+    for (size_t dim = 0; dim < 4; ++dim) {
+      const Range &range = op->srcRegion_->region[dim];
+      if (!analyzer_->CanProveEqual(range->min, 0) ||
+          !analyzer_->CanProveEqual(range->extent, src->shape[dim])) {
+        return stmt;
+      }
+    }
+    PrimExpr h = src->shape[1], w = src->shape[2], c = src->shape[3];
+    PrimExpr oh =
+        FloorDiv(h + 2 * op->padding_ - (op->kernel_ - 1) * op->dilation_ - 1,
+                 op->stride_) +
+        1;
+    PrimExpr ow =
+        FloorDiv(w + 2 * op->padding_ - (op->kernel_ - 1) * op->dilation_ - 1,
+                 op->stride_) +
+        1;
+    if (!analyzer_->CanProve(oh > 0) || !analyzer_->CanProve(ow > 0) ||
+        !analyzer_->CanProve(c > 0)) {
+      return stmt;
+    }
+
+    PrimExpr block_m = region[0]->extent, block_k = region[1]->extent;
+    Var i("im2col_m", block_m.dtype()), j("im2col_k", block_k.dtype());
+    PrimExpr m = op->nhw_step_ * block_m + i;
+    PrimExpr k = op->c_step_ * block_k + j;
+    PrimExpr batch = FloorDiv(m, oh * ow);
+    PrimExpr ih = FloorDiv(FloorMod(m, oh * ow), ow) * op->stride_ +
+                  FloorDiv(k, op->kernel_ * c) * op->dilation_ - op->padding_;
+    PrimExpr iw = FloorMod(m, ow) * op->stride_ +
+                  FloorMod(FloorDiv(k, c), op->kernel_) * op->dilation_ -
+                  op->padding_;
+    PrimExpr channel = FloorMod(k, c);
+    // Preserve zero fill for spatial/batch tails and incomplete reduction
+    // tiles. No divisibility assumption on channels or block_K is needed.
+    // These predicates only inspect coordinates. Eager boolean conjunction
+    // preserves the explicit T.Parallel form without short-circuit branches;
+    // the actual memory access remains guarded by if_then_else below.
+    PrimExpr valid = const_true();
+    for (const PrimExpr &condition :
+         Array<PrimExpr>{m >= 0, batch < src->shape[0], k >= 0,
+                         k < op->kernel_ * op->kernel_ * c, ih >= 0, ih < h,
+                         iw >= 0, iw < w}) {
+      valid = bitwise_and(valid, condition);
+    }
+    PrimExpr value =
+        if_then_else(valid, BufferLoad(src, {batch, ih, iw, channel}),
+                     make_zero(src->dtype));
+    Stmt body =
+        BufferStore(dst, value, {region[0]->min + i, region[1]->min + j});
+    body = For(j, make_zero(j.dtype()), block_k, ForKind::kParallel, body);
+    return For(i, make_zero(i.dtype()), block_m, ForKind::kParallel, body,
+               std::nullopt, {}, std::nullopt, node->span);
+  }
+};
+
+} // namespace
+
+tvm::transform::Pass ExpandSIMTIm2Col() {
+  auto pass_func = [](PrimFunc f, const IRModule &,
+                      tvm::transform::PassContext) {
+    auto target = f->GetAttr<Target>(tvm::attr::kTarget);
+    if (!target.defined() || !TargetIsCuda(target.value()) ||
+        !Im2ColUsesSIMT(target.value())) {
+      return f;
+    }
+    arith::Analyzer analyzer;
+    SIMTIm2ColExpander expander(&analyzer);
+    f.CopyOnWrite()->body = expander(f->body);
+    return f;
+  };
+  return tirx::transform::CreatePrimFuncPass(pass_func, 0,
+                                             "tl.ExpandSIMTIm2Col", {});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+  reflection::GlobalDef().def(
+      "tl.cuda.transform.UsesSIMTIm2Col", [](Target target) {
+        return TargetIsCuda(target) && Im2ColUsesSIMT(target);
+      });
+  reflection::GlobalDef().def("tl.cuda.transform.ExpandSIMTIm2Col",
+                              ExpandSIMTIm2Col);
+}
+
+} // namespace tl
+} // namespace tvm

src/op/copy.cc

@@ -278,6 +278,10 @@ const bool default_im2col_registered = RegisterDefaultIm2Col();
 
 } // namespace
 
+bool Im2ColUsesSIMT(Target target) {
+  return ResolveIm2ColImpl(target).lower == LowerIm2ColSIMT;
+}
+
 void RegisterCopyImpl(CopyImpl impl) {
   ICHECK(impl.name != nullptr);
   ICHECK(impl.match_target != nullptr);

src/op/copy.h

@@ -236,6 +236,9 @@ struct Im2ColImpl {
 
 void RegisterIm2ColImpl(Im2ColImpl impl);
 
+// Query the selected implementation rather than assuming an architecture range.
+bool Im2ColUsesSIMT(Target target);
+
 class Im2ColOp : public TileOperator {
 public:
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Im2ColOp, TileOperator,

testing/python/transform/test_tilelang_transform_expand_simt_im2col.py

@@ -0,0 +1,152 @@
+"""Early SIMT im2col: implementation selection, guards, and CUDA execution."""
+
+import pytest
+import torch
+
+import tilelang
+from tilelang import tvm
+import tilelang.language as T
+import tilelang.testing
+from test_tilelang_transform_im2col_fallback import _make_im2col_kernel
+
+
+def _runtime_simt_target():
+    major, minor = torch.cuda.get_device_capability()
+    if major < 8:
+        pytest.skip("Execution coverage requires SM80 or newer")
+    arch = f"sm_{major}{minor}" + ("a" if major >= 9 else "")
+    target = tvm.target.Target({"kind": "cuda", "arch": arch})
+    if not tilelang.cuda.transform.UsesSIMTIm2Col(target):
+        pytest.skip("This GPU selects a specialized im2col implementation")
+    return target
+
+
+def _expand(func, arch):
+    func = func.with_attr("target", tvm.target.Target({"kind": "cuda", "arch": arch}))
+    return tilelang.cuda.transform.ExpandSIMTIm2Col()(tvm.IRModule({"main": func}))
+
+
+@pytest.mark.parametrize("arch", ["sm_80", "sm_90", "sm_100", "sm_120", "sm_120a"])
+@tilelang.testing.requires_cuda
+def test_target_gating_and_idempotence(arch):
+    func = _make_im2col_kernel()
+    mod = _expand(func, arch)
+    assert tilelang.cuda.transform.UsesSIMTIm2Col(tvm.target.Target({"kind": "cuda", "arch": arch})) == (arch != "sm_90")
+    if arch != "sm_90":
+        assert "tl.tileop.im2col" not in mod.script()
+        assert "im2col_m" in mod.script()
+    else:
+        tvm.ir.assert_structural_equal(mod["main"].without_attr("target"), func)
+    tvm.ir.assert_structural_equal(tilelang.cuda.transform.ExpandSIMTIm2Col()(mod), mod)
+
+
+def _extract(channels, stride, dilation, padding, offset, annotations=None, crop=False):
+    n, h, w, kernel, bm, bk = 2, 7, 9, 3, 16, 32
+    oh = (h + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1
+    ow = (w + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1
+    mt, kt = (n * oh * ow + bm - 1) // bm, (kernel * kernel * channels + bk - 1) // bk
+
+    @T.prim_func
+    def extract(data: T.Tensor((n, h, w, channels), "float16"), out: T.Tensor((mt, kt, bm, bk), "float16")):
+        with T.Kernel(mt, kt, threads=128) as (mi, ki):
+            shared = T.alloc_shared((bm + offset, bk), "float16")
+            if crop:
+                T.im2col(
+                    data[:, 1:, :, :], shared[offset : offset + bm, :], mi, ki, kernel, stride, dilation, padding, annotations=annotations
+                )
+            else:
+                T.im2col(data, shared[offset : offset + bm, :], mi, ki, kernel, stride, dilation, padding, annotations=annotations)
+            T.copy(shared[offset : offset + bm, :], out[mi, ki, :, :])
+
+    return extract, (n, h, w, channels), (oh, ow, mt, kt, bm, bk)
+
+
+@pytest.mark.parametrize("annotations,crop", [({"test_preserve": 1}, False), (None, True)])
+@tilelang.testing.requires_cuda
+def test_unsupported_metadata_preserved(annotations, crop):
+    func, _, _ = _extract(32, 1, 1, 1, 0, annotations, crop)
+    mod = _expand(func, "sm_120a")
+    tvm.ir.assert_structural_equal(mod["main"].without_attr("target"), func)
+
+
+@tilelang.testing.requires_cuda
+def test_manual_ws_schedule_preserved():
+    func = _make_im2col_kernel()
+
+    def annotate(node):
+        if isinstance(node, tvm.tirx.SBlock):
+            return tvm.tirx.SBlock(
+                node.iter_vars,
+                node.reads,
+                node.writes,
+                node.name_hint,
+                node.body,
+                node.init,
+                node.alloc_buffers,
+                node.match_buffers,
+                {**node.annotations, "tl.ws_schedule": "preserve-sentinel"},
+            )
+
+    func = func.with_body(tvm.tirx.stmt_functor.ir_transform(func.body, annotate, None))
+    mod = _expand(func, "sm_120a")
+    tvm.ir.assert_structural_equal(mod["main"].without_attr("target"), func)
+
+
+@pytest.mark.parametrize("arch", ["sm_90"])
+@tilelang.testing.requires_cuda
+def test_other_target_source_unchanged(arch):
+    func = _make_im2col_kernel()
+    target = {"kind": "cuda", "arch": arch}
+    sources = []
+    for enabled in (False, True):
+        with tvm.target.Target(target), tvm.transform.PassContext(config={"tl.enable_early_simt_im2col": enabled}):
+            sources.append(tilelang.lower(func, target=target).kernel_source)
+    assert sources[0] == sources[1]
+    if arch == "sm_90":
+        assert "tma_load_im2col" in sources[1]
+
+
+@pytest.mark.parametrize(
+    "channels,stride,dilation,padding,offset", [(32, 1, 1, 1, 0), (48, 2, 2, 2, 0), (33, 1, 1, 1, 0), (65, 2, 1, 1, 3)]
+)
+@tilelang.testing.requires_cuda
+@tilelang.testing.requires_cuda_compute_version_ge(8, 0)
+def test_simt_tail_and_destination_region(channels, stride, dilation, padding, offset):
+    target = _runtime_simt_target()
+    func, shape, dims = _extract(channels, stride, dilation, padding, offset)
+    oh, ow, mt, kt, bm, bk = dims
+    data = torch.randn(shape, device="cuda", dtype=torch.float16)
+    cols = torch.nn.functional.unfold(data.permute(0, 3, 1, 2), 3, dilation=dilation, padding=padding, stride=stride)
+    # torch unfolds C, KH, KW; TileLang im2col uses KH, KW, C.
+    cols = cols.reshape(shape[0], channels, 3, 3, oh * ow).permute(0, 4, 2, 3, 1).reshape(-1, 9 * channels)
+    padded = torch.zeros((mt * bm, kt * bk), device="cuda", dtype=torch.float16)
+    padded[: cols.shape[0], : cols.shape[1]] = cols
+    expected = padded.reshape(mt, bm, kt, bk).permute(0, 2, 1, 3).contiguous()
+    kernel = tilelang.compile(
+        func,
+        out_idx=[1],
+        target=target,
+        pass_configs={"tl.enable_early_simt_im2col": True, "tl.disable_warp_specialized": True},
+    )
+    torch.testing.assert_close(kernel(data), expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("ws", [False, True])
+@pytest.mark.parametrize("num_stages", [0, 1, 2, 3])
+@tilelang.testing.requires_cuda
+@tilelang.testing.requires_cuda_compute_version_ge(8, 0)
+def test_simt_convolution(ws, num_stages):
+    target = _runtime_simt_target()
+    if ws and torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("WS requires TMA support")
+    func = _make_im2col_kernel(channels=48, num_stages=num_stages)
+    data = torch.randn((1, 8, 8, 48), device="cuda", dtype=torch.float16)
+    weight = torch.randn((3, 3, 48, 32), device="cuda", dtype=torch.float16)
+    expected = torch.nn.functional.conv2d(data.permute(0, 3, 1, 2), weight.permute(3, 2, 0, 1), padding=1).permute(0, 2, 3, 1)
+    kernel = tilelang.compile(
+        func,
+        out_idx=[2],
+        target=target,
+        pass_configs={"tl.enable_early_simt_im2col": True, "tl.disable_warp_specialized": not ws},
+    )
+    torch.testing.assert_close(kernel(data, weight), expected, rtol=1e-2, atol=1e-2)

testing/python/transform/test_tilelang_transform_im2col_fallback.py

@@ -6,7 +6,7 @@
 import tilelang.testing
 
 
-def _make_im2col_kernel(use_deprecated_alias=False, channels=32, block_K=32, block_M=16, hw=8):
+def _make_im2col_kernel(use_deprecated_alias=False, channels=32, block_K=32, block_M=16, hw=8, num_stages=3):
     N, C, H, W, F, K = 1, channels, hw, hw, 32, 3
     S, D, P = 1, 1, 1
     block_N = 32
@@ -30,7 +30,7 @@ def conv(
             out_flat = T.Tensor((N * OH * OW, F), T.float16, out.data)
 
             T.clear(out_local)
-            for k_iter in T.Pipelined(T.ceildiv(KH * KW * C, block_K), num_stages=3):
+            for k_iter in T.Pipelined(T.ceildiv(KH * KW * C, block_K), num_stages=num_stages):
                 if use_deprecated_alias:
                     T.c2d_im2col(data, data_shared, by, k_iter, KH, S, D, P)
                 else:

tilelang/cuda/pipeline.py

@@ -86,6 +86,11 @@ def CUDAPassPipelineBodyPrologue(mod: IRModule, target: Target) -> IRModule:
     mod = tilelang.transform.InjectAssumes()(mod)
     # Simplify the IR expressions
     mod = tilelang.transform.Simplify()(mod)
+    if tilelang.transform.get_pass_context().config.get("tl.enable_early_simt_im2col", False) and tilelang.cuda.transform.UsesSIMTIm2Col(
+        target
+    ):
+        mod = tilelang.cuda.transform.ExpandSIMTIm2Col()(mod)
+        mod = tilelang.transform.Simplify()(mod)
     # Verify reducer v2 epoch lifecycle and access rules (early, so
     # diagnostics point at user-written code)
     mod = tilelang.transform.CanonicalizeLegacyReducer()(mod)

tilelang/cuda/transform/__init__.py

@@ -3,6 +3,20 @@
 from .. import _ffi_api
 
 
+def UsesSIMTIm2Col(target):
+    """Whether this CUDA target selects the generic SIMT im2col implementation."""
+    return _ffi_api.UsesSIMTIm2Col(target)
+
+
+def ExpandSIMTIm2Col():
+    """Opt-in SIMT im2col expansion before WS and pipeline planning.
+
+    Preserve annotated operations and specialized implementations unchanged.
+    The CUDA pipeline enables this only with ``tl.enable_early_simt_im2col``.
+    """
+    return _ffi_api.ExpandSIMTIm2Col()
+
+
 def AutoWarpSpecialization():
     """Derive a warp-specialization schedule with the scheduler named by
     the ``tl.enable_auto_warp_specialization`` pass config (currently

tilelang/transform/pass_config.py

@@ -12,6 +12,9 @@ class PassConfigKey(str, Enum):
 
     # TileLang specific configs: TL_XX
 
+    TL_ENABLE_EARLY_SIMT_IM2COL = "tl.enable_early_simt_im2col"
+    """Expose CUDA SIMT im2col loads before scheduling. Default: False."""
+
     TL_SIMPLIFY = "tl.Simplify"
     """Configuration for TileLang simplification passes.