[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap) - #20651

Merged
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head
Jul 4, 2026
Merged

[ExecuTorch][WebGPU] 2D-fold mul + permute dispatch (lift 65535 1D cap)#20651
meta-codesync[bot] merged 6 commits into
gh/JulianCloudNTH/83/basefrom
gh/JulianCloudNTH/83/head

Conversation

@ghost

@ghostghost commented Jun 30, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

Lift the 65535 workgroup-per-dim cap for mul and permute so they run at any numel.

mul.Tensor and permute still used compute_1d_workgroup_count, which throws once numel / wg_size > 65535 — hit by a realistic Llama-3.2-1B LoRA layer (mul over [2048, 8192] = 262k workgroups; permute of [2048, 2048] = 65536). add/sub/div/fill/sdpa already use the 2D fold; this brings mul + permute in line.

Key changes:

  • mul/BinaryOp.cpp, permute/Permute.cppcompute_1d_workgroup_countcompute_2d_workgroup_count (returns utils::WgCount); dispatch + resize hook now set both workgroup_count_x and workgroup_count_y.
  • binary_mul.wgsl, permute.wgslmain takes @builtin(num_workgroups); flat index gid.x + gid.y * (num_workgroups.x * wg_size) (regenerated *_wgsl.h).

Mirrors the landed add op fold (runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110149677

Differential Revision: D110149677

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence June 30, 2026 21:15 — with GitHub Actions Inactive
@meta-clameta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026
@pytorch-bot

pytorch-botBot commented Jun 30, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20651

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 91b2d97 with merge base 124b677 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]
@ghost

ghost commented Jul 3, 2026

Copy link
Copy Markdown
Author

@claude Review thoroughly. Check for stale code / bugs.

@claude

claudeBot commented Jul 3, 2026

Copy link
Copy Markdown

Claude finished @JulianCloudNTH's task in 2m 47s —— View job


Review: WebGPU 2D-fold for mul + permute

  • Read PR diff and understand the change
  • Compare against the reference add op fold
  • Check compute_2d_workgroup_count / WgCount / add_dispatch signatures
  • Verify WGSL flat-index math is correct
  • Check for stale code / bugs
  • Post review

Reviewed against the landed add fold (ops/add/{BinaryOp.cpp,binary_add.wgsl}), WebGPUUtils.h (compute_2d_workgroup_count/fold_workgroup_count_2d/WgCount), and the dispatch path in WebGPUGraph.{h,cpp}. The change is correct and a faithful mirror of the reference — no bugs found. One stale comment worth fixing.

✅ Correctness verified

  • add_dispatch brace-init matches the struct.WebGPUDispatch field order is {pipeline, bind_group, workgroup_count_x, kernel_name, workgroup_count_y} (WebGPUGraph.h:46-51). Both call sites pass {pipeline, bind_group, workgroup_count.x, "mul"/"permute", workgroup_count.y} — positions line up exactly. add uses "" as the label; mul/permute supplying real bench labels is a minor improvement, not a bug.
  • Dispatch consumes .y.WebGPUGraph.cpp:818 and :898 both call wgpuComputePassEncoderDispatchWorkgroups(pass, workgroup_count_x, workgroup_count_y, 1), so the folded y is actually launched.
  • Flat-index math is exact.idx = gid.x + gid.y * (num_workgroups.x * wg_size). With grid (x, y, 1) and @workgroup_size(wg_size,1,1): gid.x ∈ [0, x·wg_size), gid.y ∈ [0, y), num_workgroups.x == x, so idx covers [0, x·y·wg_size), and x·y ≥ ceil(numel/wg_size). The idx >= out_meta.numel guard drops the tail. Byte-identical to binary_add.wgsl:17.
  • mul resize hook rewrites both workgroup_count_x and workgroup_count_y (BinaryOp.cpp:204-205), matching add:190-191. permute has no resize hook (it registers none), so nothing to update there — consistent with the op being static-shape only.
  • num_workgroups builtin is correctly added to both main signatures and both regenerated *_wgsl.h SHAs were updated.

🟡 Minor: stale comment

backends/webgpu/runtime/ops/mul/BinaryOp.cpp:37

// Rank guard (NCHW backend is <= 4 dims; 1D dispatch only).

The 1D dispatch only clause is now inaccurate — dispatch is 2D-folded. Suggest trimming to // Rank guard (NCHW backend is <= 4 dims).permute has no equivalent stale comment. Fix this →

💭 Observations (non-blocking, shared with the landed add)

  • gid.y * (num_workgroups.x * wg_size) is u32 arithmetic. In practice bounded by numel (a uint32_t), so no overflow for any realistic tensor; only a theoretical concern as numel → 2³². This is inherited verbatim from the reference add/sdpa fold, so it's not something this PR should change in isolation.
  • No test exercises the >65535 path.test_mul.py / test_permute.py exist but I didn't find a case at the numel that previously threw (mul over [2048, 8192], permute of [2048, 2048]). Adding one would lock in the fix and guard against regression — though if the add fold landed without such a case, matching that precedent is reasonable.

Verdict: LGTM. The one stale comment is optional cleanup; the logic is sound and consistent with the established pattern.
· branch gh/JulianCloudNTH/83/head

[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:28 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 20:52 — with GitHub Actions Inactive
[ghstack-poisoned]
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 21:37 — with GitHub Actions Inactive
@ghost
ghost temporarily deployed to cadence July 3, 2026 22:06 — with GitHub Actions Inactive
@meta-codesync
meta-codesyncBot merged commit 1055559 into gh/JulianCloudNTH/83/baseJul 4, 2026
181 of 183 checks passed
@meta-codesync
meta-codesyncBot deleted the gh/JulianCloudNTH/83/head branch July 4, 2026 17:06
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
ghost pushed a commit that referenced this pull request Jul 4, 2026
Pull Request resolved: #20651
**Lift the 65535 workgroup-per-dim cap for `mul` and `permute` so they run at any numel.**
`mul.Tensor` and `permute` still used `compute_1d_workgroup_count`, which throws once `numel / wg_size > 65535` — hit by a realistic Llama-3.2-1B LoRA layer (`mul` over `[2048, 8192]` = 262k workgroups; `permute` of `[2048, 2048]` = 65536). `add`/`sub`/`div`/`fill`/`sdpa` already use the 2D fold; this brings `mul` + `permute` in line.
Key changes:
- `mul/BinaryOp.cpp`, `permute/Permute.cpp` — `compute_1d_workgroup_count` → `compute_2d_workgroup_count` (returns `utils::WgCount`); dispatch + resize hook now set both `workgroup_count_x` and `workgroup_count_y`.
- `binary_mul.wgsl`, `permute.wgsl` — `main` takes `@builtin(num_workgroups)`; flat index `gid.x + gid.y * (num_workgroups.x * wg_size)` (regenerated `*_wgsl.h`).
Mirrors the landed `add` op fold (`runtime/ops/add/{BinaryOp.cpp,binary_add.wgsl}`).
Co-authored-with: Claude Code.
ghstack-source-id: 399812930
@exported-using-ghexport
Differential Revision: [D110149677](https://our.internmc.facebook.com/intern/diff/D110149677/)
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@psiddh@nil-is-all@JCNTH