Skip to content

[Fix][Relax] Lower non-contiguous WebGPU cumsum - #20133

Merged
tlopex merged 1 commit into
apache:mainfrom
akaashrp:upstream/webgpu-cumsum-noncontiguous
Aug 25, 2026
Merged

[Fix][Relax] Lower non-contiguous WebGPU cumsum#20133
tlopex merged 1 commit into
apache:mainfrom
akaashrp:upstream/webgpu-cumsum-noncontiguous

Conversation

@akaashrp

@akaashrp akaashrp commented Aug 16, 2026

Copy link
Copy Markdown
Contributor

This PR adds direct WebGPU lowering for inclusive cumsum along non-innermost axes. Previously, these operations used TOPI's transpose-based GPU scan, which materializes a full-size transposed tensor and requires multiple WebGPU dispatches. This is particularly problematic when the scan extent is symbolic.

The new lowering:

  • Normalizes the scan axis and reshapes known-rank inputs to [outer, scan, inner] without transposing data
  • Uses a correctness fallback where each GPU thread scans one (outer, inner) row serially
  • Preserves input/output dtypes and virtual-device information
  • Retains the existing parallel 2D kernel for innermost-axis scans
  • Clamps the existing kernel's reverse-round count so short scans do not generate negative shifts during WebGPU index narrowing

@akaashrp akaashrp changed the title [FIX][RELAX] Lower non-contiguous WebGPU cumsum [Fix][Relax] Lower non-contiguous WebGPU cumsum Aug 16, 2026
@yongwww yongwww closed this Aug 17, 2026
@yongwww yongwww reopened this Aug 17, 2026
@akaashrp

Copy link
Copy Markdown
Contributor Author

@tvm-bot rerun

@akaashrp
akaashrp force-pushed the upstream/webgpu-cumsum-noncontiguous branch 4 times, most recently from 3222ef5 to ba18511 Compare August 24, 2026 09:38
@akaashrp
akaashrp force-pushed the upstream/webgpu-cumsum-noncontiguous branch from ba18511 to a4204b3 Compare August 24, 2026 10:44
@akaashrp
akaashrp requested a review from tlopex August 25, 2026 02:41
@tlopex
tlopex merged commit 1f352bf into apache:main Aug 25, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants