[BugFix] Pass buffer row stride to 2D scan kernel to fix silent miscomputation

[ColmaLiu]

Summary

T.cumsum and T.cummax silently return wrong results when operating on a 2‑D slice of a shared buffer whose row pitch differs from the slice width. The kernel compiles and runs without errors or crashes, but the numbers are deterministically wrong.

Root cause: LowerSharedScan passes only the slice's logical extents (H, W) to the device kernel, without any row‑pitch argument. The device template InclusiveScan2D then assumes the row stride equals W. When the backing buffer is wider than the slice (e.g. a (8, 64) buffer scanned over only the first 40 columns), every row after the first is mis‑addressed by (NBIG − N) elements, corrupting the scan.

The bug is architecture‑independent and affects both cumsum and cummax since they share the same InclusiveScan2D template. When the region is contiguous (buffer width == slice width) results are bit‑exact.

Fixes #2523.

Changes

• src/backend/common/op/scan.h — In the 2‑D branch of LowerSharedScan, compute the buffer's true row stride from op.src->shape[last_dim] and op.dst->shape[last_dim], and pass them as a (src_stride, dst_stride) pair after the logical extents. When src and dst are different buffers with different row pitches, each gets its own stride — avoiding silent mis‑addressing.

• src/tl_templates/cuda/scan.h — Split the single stride parameter into src_stride / dst_stride throughout the chain:

  • InclusiveScanLine: int stride → int src_stride, int dst_stride. All src[…] indexing uses src_stride; all dst[…] indexing
    uses dst_stride.
  • InclusiveScan1D::run: passes (1, 1) for the two strides (the contiguous in‑lane case is unchanged).
  • InclusiveScan2D::run: int src_stride, int dst_stride replace the single int stride. For Axis == 1 the row address is src + row * src_stride / dst + row * dst_stride; for Axis == 0 the column stride is forwarded as src_stride / dst_stride to Inclusive ScanLine.
  • CumSum2D::run, CumMax2D::run: accept and pass through the stride pair.

• src/tl_templates/hip/scan.h — Identical changes for the ROCm/HIP template.

Testing

• testing/python/language/test_tilelang_language_scan.py — Add two test functions (test_cumsum_strided_region, test_cummax_strided_region) following the existing convention. Each iterates over multiple (M, N, NBIG, dim, reverse) combinations that exercise the non‑contiguous path.

Summary

• Fixed incorrect 2-D T.cumsum and T.cummax results on non-contiguous shared-buffer slices.
• Propagated the backing buffer’s row stride from LowerSharedScan through CUDA and HIP scan templates.
• Updated row- and column-wise addressing while preserving the slice’s logical width and correct behavior for contiguous regions.
• Added tests covering varied shapes, scan dimensions, reverse scans, and strided regions across shared-buffer sub-regions.

C++ style / lint notes

• The PR changes C++ scan templates but does not modify docs/developer_guide/cpp_style.md.
• The C++ API Style Audit (warning only) CI step is relevant since the PR touches C++ template code; any TLCPP003/TLCPP004 findings are advisory and should be treated separately from correctness/build/test issues (which this PR addresses via updated scan logic and new Python tests).

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: b1311af8-cfcd-46bf-ab4b-0252aad7eb5f

📥 Commits

Reviewing files that changed from the base of the PR and between b2110b3 and 396d4d0.

📒 Files selected for processing (1)

• src/tl_templates/cuda/scan.h

🚧 Files skipped from review as they are similar to previous changes (1)

• src/tl_templates/cuda/scan.h

---

📝 Walkthrough

Walkthrough

2D shared-buffer scan lowering now propagates separate source and destination strides through CUDA and HIP scan templates. New tests validate cumsum and cummax on non-contiguous shared-memory regions across dimensions and reverse modes.

Changes

Strided 2D scan support

Layer / File(s)	Summary
Lowering propagates source and destination strides src/backend/common/op/scan.h	The 2D shared-scan lowering derives source and destination strides and includes both in the external call arguments.
CUDA and HIP scan addressing src/tl_templates/cuda/scan.h, src/tl_templates/hip/scan.h	Inclusive scans, cumsum, and cummax accept separate source and destination strides and use them for loads, stores, and row or column pointer arithmetic.
Strided scan validation testing/python/language/test_tilelang_language_scan.py	Tests run cumsum and cummax over larger shared buffers with non-contiguous 2D subregions, including reverse and axis variants.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SharedScanLowering
  participant ScanExtern
  participant GPUScanTemplate
  participant StridedScanTest
  SharedScanLowering->>ScanExtern: pass extents and source/destination strides
  ScanExtern->>GPUScanTemplate: invoke 2D cumsum or cummax
  GPUScanTemplate->>GPUScanTemplate: address rows or columns using separate strides
  StridedScanTest->>GPUScanTemplate: execute non-contiguous shared-region case
  GPUScanTemplate-->>StridedScanTest: return scanned output

Loading

Possibly related issues

• [BUG][Fuzzer][wrong-code] T.cumsum/T.cummax over a row-offset 2-D shared sub-region silently scans the wrong rows instead of the requested ones #2536: Related to 2D scan addressing; this change updates scan stride propagation, while the issue concerns a separate row-offset bug in MakeAccessPtrFromRegion.

Possibly related PRs

• tile-ai/tilelang#2664: Both changes modify CUDA InclusiveScanLine; this PR adds separate source and destination stride arithmetic, while that PR changes out-of-range lane padding for pipelining.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main bug fix: passing row stride to 2-D scans to prevent incorrect results.
Linked Issues check	✅ Passed	The PR fixes the non-contiguous 2-D scan bug, propagates true row pitch through lowering, and adds strided coverage.
Out of Scope Changes check	✅ Passed	The changes stay focused on the scan-stride fix and related tests; the HIP updates are consistent with the same fix.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/backend/common/op/scan.h`:
- Around line 57-63: Update the 2D scan lowering branch in the scan operation to
account for destination row stride separately from the source stride. Pass the
stride derived from op.dst alongside the existing op.src stride, and update the
corresponding tl::<symbol_prefix>2D kernel run signature and indexing to use
each buffer’s own stride; alternatively validate and reject mismatched layouts
before lowering.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: c090f8a3-eb29-457b-b0ec-8435984ea650

📥 Commits

Reviewing files that changed from the base of the PR and between bd76473 and 30c9615.

📒 Files selected for processing (4)

• src/backend/common/op/scan.h
• src/tl_templates/cuda/scan.h
• src/tl_templates/hip/scan.h
• testing/python/language/test_tilelang_language_scan.py

[coderabbitai]

Caution

Review failed

The head commit changed during the review from ebf4673 to 30c9615.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}