diff --git a/src/cuda/transform/expand_simt_im2col.cc b/src/cuda/transform/expand_simt_im2col.cc new file mode 100644 index 0000000000..bb74cbea11 --- /dev/null +++ b/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 +#include +#include + +#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(node); + } + return arith::IRMutatorWithAnalyzer::VisitStmt_(node); + } + + Stmt VisitStmt_(const EvaluateNode *node) final { + Stmt stmt = arith::IRMutatorWithAnalyzer::VisitStmt_(node); + const auto *call = stmt.as()->value.as(); + if (!call || !call->op.same_as(Im2ColOp::Get()) || + !call->annotations.empty()) { + return stmt; + } + Im2ColOp op = Downcast(ParseOperator(GetRef(call))); + const Buffer &src = op->src_; + const Buffer &dst = op->dst_; + const auto ®ion = 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{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(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 diff --git a/src/op/copy.cc b/src/op/copy.cc index e679f73567..ea85a15852 100644 --- a/src/op/copy.cc +++ b/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); diff --git a/src/op/copy.h b/src/op/copy.h index a9fab87694..96f5407f53 100644 --- a/src/op/copy.h +++ b/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, diff --git a/testing/python/transform/test_tilelang_transform_expand_simt_im2col.py b/testing/python/transform/test_tilelang_transform_expand_simt_im2col.py new file mode 100644 index 0000000000..2b7d438d2e --- /dev/null +++ b/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) diff --git a/testing/python/transform/test_tilelang_transform_im2col_fallback.py b/testing/python/transform/test_tilelang_transform_im2col_fallback.py index ebf77ba65b..b1213b057a 100644 --- a/testing/python/transform/test_tilelang_transform_im2col_fallback.py +++ b/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: diff --git a/tilelang/cuda/pipeline.py b/tilelang/cuda/pipeline.py index 85f2e49046..a1d4c33f42 100644 --- a/tilelang/cuda/pipeline.py +++ b/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) diff --git a/tilelang/cuda/transform/__init__.py b/tilelang/cuda/transform/__init__.py index 703ea43fdb..16648b198a 100644 --- a/tilelang/cuda/transform/__init__.py +++ b/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 diff --git a/tilelang/transform/pass_config.py b/tilelang/transform/pass_config.py index b9e2d8b906..00e11ad6fc 100644 --- a/tilelang/transform/pass_config.py +++ b/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.