Skip to content

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

Open
akaashrp wants to merge 1 commit into
apache:mainfrom
akaashrp:upstream/webgpu-cumsum-noncontiguous
Open

[Fix][Relax] Lower non-contiguous WebGPU cumsum#20133
akaashrp wants to merge 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 from c866382 to 190f40c Compare August 22, 2026 03:29
@akaashrp
akaashrp force-pushed the upstream/webgpu-cumsum-noncontiguous branch from 190f40c to d7d4320 Compare August 22, 2026 06:11
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.

2 participants