From 5ef1500ed27add9adc3e58fed0e809219edea345 Mon Sep 17 00:00:00 2001 From: Zheng Yu Date: Mon, 27 Jul 2026 00:38:52 -0700 Subject: [PATCH] [BugFix] Handle strided global buffers correctly in 1D TMA copies (#2746) * Honor explicit strides in 1D TMA offsets * Reject non-contiguous strided regions for 1D TMA --------- Co-authored-by: Zheng Yu <17133261+cla7aye15I4nd@users.noreply.github.com> Co-authored-by: SiriusNEO --- src/cuda/op/copy.cc | 15 ++-- src/cuda/op/copy_analysis.cc | 57 +++++++++----- .../test_tilelang_language_tma_copy.py | 76 +++++++++++++++++++ 3 files changed, 125 insertions(+), 23 deletions(-) diff --git a/src/cuda/op/copy.cc b/src/cuda/op/copy.cc index 4580f9fb9..45656c9c1 100644 --- a/src/cuda/op/copy.cc +++ b/src/cuda/op/copy.cc @@ -2307,11 +2307,16 @@ Stmt Copy::LowerBulk1D(const CopyNode &op, const LowerArgs &lower_args, global_indices.push_back(r->min); } std::vector global_strides; - PrimExpr global_stride = 1; - for (size_t i = 0; i < global_tensor->shape.size(); i++) { - auto s = global_tensor->shape[global_tensor->shape.size() - i - 1]; - global_strides.insert(global_strides.begin(), global_stride); - global_stride *= s; + if (global_tensor->strides.empty()) { + PrimExpr global_stride = 1; + for (size_t i = 0; i < global_tensor->shape.size(); i++) { + auto s = global_tensor->shape[global_tensor->shape.size() - i - 1]; + global_strides.insert(global_strides.begin(), global_stride); + global_stride *= s; + } + } else { + global_strides.assign(global_tensor->strides.begin(), + global_tensor->strides.end()); } PrimExpr global_offset = 0; diff --git a/src/cuda/op/copy_analysis.cc b/src/cuda/op/copy_analysis.cc index 3c9a7a457..1e6af3aff 100644 --- a/src/cuda/op/copy_analysis.cc +++ b/src/cuda/op/copy_analysis.cc @@ -282,6 +282,43 @@ bool IsSemanticallyLinearLayout(const Layout &layout, linear_layout->GetLinearizedForwardIndex()); } +bool IsContiguousGlobalRegion(const Buffer &global_tensor, + const Array &global_range, + arith::Analyzer *analyzer) { + ICHECK_EQ(global_range.size(), global_tensor->shape.size()); + ICHECK(global_tensor->strides.empty() || + global_tensor->strides.size() == global_tensor->shape.size()); + + // Pointer-based 1D TMA can coalesce multiple logical dimensions only when + // every varying dimension is packed directly after its contiguous suffix. + bool non_full_dim_encountered = false; + PrimExpr expected_stride = 1; + for (int i = static_cast(global_range.size()) - 1; i >= 0; --i) { + bool is_singleton = analyzer->CanProve( + global_range[i]->extent == 1, arith::ProofStrength::kSymbolicBound); + if (!non_full_dim_encountered) { + if (!is_singleton && !global_tensor->strides.empty() && + !analyzer->CanProveEqual(global_tensor->strides[i], + expected_stride)) { + return false; + } + + bool is_full_dim = analyzer->CanProve( + global_range[i]->extent == global_tensor->shape[i] && + global_range[i]->min == 0, + arith::ProofStrength::kSymbolicBound); + if (is_full_dim) { + expected_stride *= global_tensor->shape[i]; + } else { + non_full_dim_encountered = true; + } + } else if (!is_singleton) { + return false; + } + } + return true; +} + bool CheckBulkCopy1D(const Buffer &global_tensor, const Buffer &shared_tensor, const Array &global_range, const Array &shared_range, @@ -294,24 +331,8 @@ bool CheckBulkCopy1D(const Buffer &global_tensor, const Buffer &shared_tensor, IsSemanticallyLinearLayout(existing, shared_tensor->shape, analyzer); } - bool global_is_contiguous = true; - bool global_not_full_dim_encounter = false; - for (int i = global_range.size() - 1; i >= 0; i--) { - if (!global_not_full_dim_encounter) { - if (!analyzer->CanProve(global_range[i]->extent == - global_tensor->shape[i] && - global_range[i]->min == 0, - arith::ProofStrength::kSymbolicBound)) { - global_not_full_dim_encounter = true; - } - } else { - if (!analyzer->CanProve(global_range[i]->extent == 1, - arith::ProofStrength::kSymbolicBound)) { - global_is_contiguous = false; - break; - } - } - } + bool global_is_contiguous = + IsContiguousGlobalRegion(global_tensor, global_range, analyzer); PrimExpr shared_elements = 1; for (size_t i = 0; i < shared_range.size(); i++) { diff --git a/testing/python/language/test_tilelang_language_tma_copy.py b/testing/python/language/test_tilelang_language_tma_copy.py index a87e61992..661f2761e 100644 --- a/testing/python/language/test_tilelang_language_tma_copy.py +++ b/testing/python/language/test_tilelang_language_tma_copy.py @@ -177,6 +177,82 @@ def test_tma_copy_default_block_leader_scope_codegen(): assert "tl_shuffle_elect<256>()" in source, "Expected block-wide elect<256>" +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(9, 0) +def test_tma_copy_uses_explicit_global_stride_for_row_offset(): + rows = 2 + cols = 32 + row_stride = 40 + + @T.prim_func + def main( + A: T.StridedTensor((rows, cols), (row_stride, 1), T.float32), + B: T.Tensor((cols,), T.float32), + ): + with T.Kernel(1, threads=32): + a_shared = T.alloc_shared((cols,), T.float32) + T.copy(A[1, 0:cols], a_shared, prefer_instruction="tma") + T.copy(a_shared, B) + + import torch + + kernel = tilelang.compile(main, out_idx=[1]) + assert "CUtensorMap" not in kernel.get_kernel_source() + padded = torch.arange(rows * row_stride, dtype=torch.float32, device="cuda").reshape(rows, row_stride) + source = padded[:, :cols] + assert source.stride() == (row_stride, 1) + output = kernel(source) + torch.testing.assert_close(output, source[1]) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(9, 0) +def test_tma_copy_uses_descriptor_for_padded_multirow_region(): + rows = 2 + cols = 32 + row_stride = 40 + + @T.prim_func + def load( + A: T.StridedTensor((rows, cols), (row_stride, 1), T.float32), + B: T.Tensor((rows, cols), T.float32), + ): + with T.Kernel(1, threads=32): + a_shared = T.alloc_shared((rows, cols), T.float32) + T.copy(A, a_shared, prefer_instruction="tma") + T.copy(a_shared, B, prefer_instruction="sync") + + @T.prim_func + def store( + A: T.Tensor((rows, cols), T.float32), + B: T.StridedTensor((rows, cols), (row_stride, 1), T.float32), + ): + with T.Kernel(1, threads=32): + a_shared = T.alloc_shared((rows, cols), T.float32) + T.copy(A, a_shared, prefer_instruction="sync") + T.copy(a_shared, B, prefer_instruction="tma") + + import torch + + source = torch.arange(rows * cols, dtype=torch.float32, device="cuda").reshape(rows, cols) + + load_kernel = tilelang.compile(load, out_idx=[1]) + assert "CUtensorMap" in load_kernel.get_kernel_source() + padded_source = torch.full((rows, row_stride), -1.0, dtype=torch.float32, device="cuda") + strided_source = padded_source[:, :cols] + strided_source.copy_(source) + output = load_kernel(strided_source) + torch.testing.assert_close(output, source) + + store_kernel = tilelang.compile(store) + assert "CUtensorMap" in store_kernel.get_kernel_source() + padded_destination = torch.full((rows, row_stride), -1.0, dtype=torch.float32, device="cuda") + strided_destination = padded_destination[:, :cols] + store_kernel(source, strided_destination) + torch.testing.assert_close(strided_destination, source) + torch.testing.assert_close(padded_destination[:, cols:], torch.full_like(padded_destination[:, cols:], -1.0)) + + def matmul_tma_copy_store( M, N,