[ET-VK][conv2d] Use a 2x2x8 tile and tuned local sizes for pointwise convolution by SS-JIA · Pull Request #24433 · pytorch/executorch

SS-JIA · 2026-11-02T15:35:45Z

Stack from ghstack (oldest at bottom):

-> [ET-VK][conv2d] Use a 2x2x8 tile and tuned local sizes for pointwise convolution #24433

Ulterior Motive

Closing the latency gap between ET-VK and LiteRT (TFLite GPU) on the DSR denoising model on mobile GPUs. The model's 80 1x1 convolutions were the largest per-operator gap on Mali.

Rationale

What: Reworks conv2d_pw_tiled around a 2x2-pixel x 8-channel tile with buffer weights, branch-free bounds handling and per-GPU local sizes. The 2x1-pixel x 4-channel tile stays for small Half grids and for Float on Mali, and Float on Mali uses texture weights.

Why: Pointwise convolutions are load/store bound on Mali. The old tile issued 6 texel loads per 8 vec4 FMAs, read weights through a texture, and guarded every load with a bounds check. The 2x2 x 8 tile issues half as many loads per FMA.

Details

Tile
- Before: one tile, 2 consecutive output positions x 4 output channels per invocation.
- After: a 2x2 output pixel block x 8 output channels (_4x8 variants) for Float on every GPU except Mali, and for Half when the grid has at least 4096 invocations. Smaller Half grids keep the 2x1 x 4 tile, which fills the GPU better there, and so does Float on Mali.
Weight storage
- Before: texture2d, unless the packed weights exceed texture limits.
- After: buffer (_bufw variants) for Float on every GPU except Mali, and for Half when the device supports 16-bit storage buffers; texture2d otherwise. On a Mali-G715 in CI, Float buffer weights produced wrong results for padded and strided convolutions with either tile. The same cases passed with the parent's texture weights.
Bounds handling
- Before: every input and weight load is guarded.
- After: pixels and channel tiles past the edge read the nearest in-range address and are discarded at the store, so weight loads have no branches. One redundant K check stays on the input load; without any branch in the K loop, the Adreno compiler reads the weight buffer through a much slower path.
Local size
- Before: 64-invocation square for every dispatch.
- After: wide tile work groups span the channel tiles of a row of pixel tiles, with 512 threads on Mali, 128 on Adreno and 64 elsewhere, halved until they fit the device limits. The narrow tile uses 1x64 along pixel tiles on Mali when there are at least 64 of them, and the 64-invocation square otherwise.
Tests
- The pointwise reference supports Half. pw_correctness runs Half too, and adds partial-tile shapes below and above the Half tile switch.
- New timing-only dsr_pw_set with the DSR model's 1x1 convolution shapes.

The tile switch and local sizes come from a sweep over 12 local sizes and both tiles on 57 shapes (see Test Plan). Float on Mali keeps the 2x1 x 4 tile: in the Mali offline compiler the 2x2 x 8 Float tile spills 88 bytes per thread while the 2x1 x 4 Float tile does not spill, and the wide Float tile was 2-31% slower in the sweep.

Differential Revision: D128786678

[ghstack-poisoned]

pytorch-bot · 2026-11-02T15:35:49Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 16 Pending, 2 Unclassified Failures

As of commit 8c9548a with merge base 822cb60 ():

NEW FAILURES - The following jobs have failed:

Test ARM Backend / test-arm-vgf / test-backend-linux (arm_vgf_fp, models) / linux-job (gh)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 56
Test ARM Backend / test-arm-vgf / test-backend-linux (arm_vgf_fp, operators) / linux-job (gh)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 56
Test QNN Windows MSVC build / build-qnn-windows-arm64 (gh)

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

Build Aarch64 Linux Native Wheels / executorch-native / build-wheel-py3_11-cpu-aarch64 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
ModuleNotFoundError: No module named 'setuptools'
Build Aarch64 Linux Native Wheels / executorch-native / upload / upload-wheel-py3_11-cpu-aarch64 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
Unable to download artifact(s): Artifact not found

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

github-actions · 2026-11-02T15:36:36Z

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.

shoumikhin · 2026-11-02T16:46:16Z

@claude review effort=max

claude · 2026-11-02T16:47:11Z

Claude finished @shoumikhin's task in 31m 7s —— View job

PR Review: #24433

This PR adds a 2x2-pixel × 8-channel tile, buffer-weight variants and vendor-tuned work groups to conv2d_pw_tiled; nothing blocks, but Half now accumulates in fp16 with no accuracy check above C_in = 33, and the weight-storage condition keys on the declared dtype, so two of the eight generated shader variants can never be selected.

Detailed review

Summary

Adds _4x8 (2x2 pixels × 8 channels) and _bufw (buffer weights) variants of conv2d_pw_tiled, chosen by dtype, GPU vendor and grid size, plus vendor-tuned local work groups. Nothing blocks. The issues:

Half numerics change without being mentioned or tested at realistic K.
The weight-storage condition reads the wrong dtype.
The dispatch decision is duplicated across the three pickers.
A few test-helper and comment nits.

Code Quality

Is the Mali buffer-weight bug really Float-only? (low confidence)
- Where: Conv2dPW.cpp:199-207.
- Buffer-weight indexing does not depend on padding or stride.
- What does differ for padded or strided convs is the input coordinate math and the per-pixel coords.valid branch in the K loop (conv2d_pw_tiled.glsl:121-122). Half _bufw runs that same code on Mali.
- If the Mali-G715 failure is a compiler issue around that branch, Half may be affected too.
- Please confirm that the padded and strided Half op_tests pass on Mali-G715, and link a tracking issue so the workaround can be revisited.
Each picker recomputes the tile choice instead of reading it from the selected shader.
- Where: Conv2dPW.cpp:118 and :150.
- pick_conv2d_pw_tiled_gwg and pick_conv2d_pw_tiled_lwg call use_wide_pw_tile again and ignore the shader that DynamicDispatchNode passes them ((void)shader; at :110).
- If the shader condition ever changes (for example, to depend on weight storage), the grid layout and work group would silently stop matching the compiled tile.
- Fix: derive it with shader.kernel_name.find("_4x8") != std::string::npos, as Conv2dDW.cpp:119-136 and Preprocess.cpp:41,60 do. Fix this →
load_packed_weight_tile_clamped does not clamp k4 when TILE_K4 == 1.
- Where: linear_fp_packed_weight_tile_load.glslh:86.
- The comment says out-of-range blocks load the nearest in-range block. With TILE_K4 == 1, though, an out-of-range k4_start is used as-is and reads past the packed weights.
- conv2d_pw_tiled only passes k4 < K4, but this is a shared gemm/tile_utils header.
- Fix: state the k4_start < K4 precondition in the comment. Fix this →
The new Half skip copies the PERF predicate instead of using conv2d_is_perf_shape.
- Where: conv2d.cpp:664-666.
- It copies the strict > from create_conv2d_pw_test_case (conv2d.cpp:319-321).
- conv2d_is_perf_shape (conv2d.cpp:32-42) uses >= because 64-wide FP16 shapes "can exceed the half tolerance and false-fail". Its comment also asks that both call sites stay on the helper "to avoid the two predicates drifting apart".
- As written, a Half pw config with C_in = 64 and small spatial dims would be labeled ACCU and reference-checked.
- Fix: use the helper at both pw call sites. This changes no current label. Fix this →

Backward Compatibility

Half switches from fp32 to fp16 accumulation, untested at the channel counts this PR targets.
- Where: conv2d_pw_tiled.glsl:12.
- What changed: VEC4_T moves from texel_load_type(DTYPE, STORAGE) (vec4 for Half textures) to buffer_gvec_type(DTYPE, 4) (f16vec4). LINEAR_FP_OUTPUT_TILE_VEC4_T defaults to VEC4_T.
- Effect: on GPUs where Half is not aliased to Float, every Half model on this path now accumulates the K loop in fp16.
- Context:
  - This matches conv2d_tiled and conv2d_pw_s1p0.
  - The PR description does not mention it.
  - The new reference skip (conv2d.cpp:661-663) itself says FP16 accumulation error "would exceed the half tolerance" at large K.
- Coverage: Half accuracy is checked only up to C_in = 33 (pw_correctness) and C_in = 16 (op_tests/cases.py test_cases_pw). The PR's Half benchmarks reach C_in = 768 (pw_performance) and 512 (dsr_pw_set, whose reference is nullptr).
- Fix:
  - State the precision change in the PR description.
  - Add a Half case at realistic K with bounded operands, as set_half_accuracy_operands does in dsr_convs_sets.cpp:54.
  - Alternatively, if the speedup does not depend on fp16 math, keep an fp32 accumulator via LINEAR_FP_OUTPUT_TILE_VEC4_T.
Fix this →

Performance

Weight storage keys on the declared dtype, not the effective one.
- Where: Conv2dPW.cpp:205-207.
- Cause: dtype_of(weight_data) returns the TensorRef's declared dtype. Tensors, however, fall back from Half to Float when has_full_float16_buffers_support() is false (Tensor.cpp:158-161).
- Effect on a non-Mali GPU without fp16 buffers: a Half model runs conv2d_pw_tiled_4x8_float with texture weights. The same model exported as Float gets _4x8_bufw_float, which the comment says is the faster path.
- Dead variants: an effective-Half output implies float16_buffers_enabled(), so conv2d_pw_tiled_half and conv2d_pw_tiled_4x8_half can never be selected. They only add binary size.
- Fix:
```
const bool buffer_weights = !graph.device_is_mali() ||
    (graph.dtype_of(weight_data) == vkapi::kHalf &&
     graph.float16_buffers_enabled());
```
  Then restrict the two texture-weight variants to float with a per-variant generate_variant_forall, as binary_op_buffer.yaml does. Fix this →
Wide work groups are not shrunk to fit small grids (minor).
- Where: Conv2dPW.cpp:162-166.
- Float on non-Mali takes the wide path at every size, and the returned LocalWorkGroup(nthreads / y, y, 1u) ignores gwg[0].
- Example: a 1x1-spatial conv such as a squeeze-excite projection (C_out = 64, gwg = {1, 8, 1}). On Adreno that is a 16x8 group with 8 of 128 invocations active. Other non-Mali GPUs get an 8x8 group with 8 of 64 active.
- Calling fit_to_global will not help, because it does nothing for explicitly sized groups (DispatchGrid.cpp:245-248).
- Fix: cap x at gwg[0], or fall back to pick_xy_square_lwg when gwg[0] < min_x. Fix this →

Specific Comments

conv2d_pw_tiled.glsl:137-138: stale comment. The coordinates now come from dividing by tiles_x, not W_out.
conv2d.cpp:661-663: the comment says "large-K half shapes", but the predicate also skips on C_out, H and W. None of those change the accumulation length of a 1x1 conv.

Recommendation

Approve. Nothing blocks. Before landing:

Call out the fp16 accumulation change in the PR description and add a Half accuracy case at realistic K.
Fold in the weight-storage fix, which is small.

The rest can be follow-ups.

[ghstack-poisoned]

…convolution Pull Request resolved: #24433 ## Ulterior Motive Closing the latency gap between ET-VK and LiteRT (TFLite GPU) on the DSR denoising model on mobile GPUs. The model's 80 1x1 convolutions were the largest per-operator gap on Mali. ## Rationale **What**: Reworks `conv2d_pw_tiled` around a 2x2-pixel x 8-channel tile with buffer weights, branch-free bounds handling and per-GPU local sizes. The 2x1-pixel x 4-channel tile stays for small Half grids and for Float on Mali, and Float on Mali uses texture weights. **Why**: Pointwise convolutions are load/store bound on Mali. The old tile issued 6 texel loads per 8 vec4 FMAs, read weights through a texture, and guarded every load with a bounds check. The 2x2 x 8 tile issues half as many loads per FMA. ## Details - Tile - **Before**: one tile, 2 consecutive output positions x 4 output channels per invocation. - **After**: a 2x2 output pixel block x 8 output channels (`_4x8` variants) for Float on every GPU except Mali, and for Half when the grid has at least 4096 invocations. Smaller Half grids keep the 2x1 x 4 tile, which fills the GPU better there, and so does Float on Mali. - Weight storage - **Before**: texture2d, unless the packed weights exceed texture limits. - **After**: buffer (`_bufw` variants) for Float on every GPU except Mali, and for Half when the device supports 16-bit storage buffers; texture2d otherwise. On a Mali-G715 in CI, Float buffer weights produced wrong results for padded and strided convolutions with either tile. The same cases passed with the parent's texture weights. - Bounds handling - **Before**: every input and weight load is guarded. - **After**: pixels and channel tiles past the edge read the nearest in-range address and are discarded at the store, so weight loads have no branches. One redundant K check stays on the input load; without any branch in the K loop, the Adreno compiler reads the weight buffer through a much slower path. - Local size - **Before**: 64-invocation square for every dispatch. - **After**: wide tile work groups span the channel tiles of a row of pixel tiles, with 512 threads on Mali, 128 on Adreno and 64 elsewhere, halved until they fit the device limits. The narrow tile uses 1x64 along pixel tiles on Mali when there are at least 64 of them, and the 64-invocation square otherwise. - Tests - The pointwise reference supports Half. `pw_correctness` runs Half too, and adds partial-tile shapes below and above the Half tile switch. - New timing-only `dsr_pw_set` with the DSR model's 1x1 convolution shapes. The tile switch and local sizes come from a sweep over 12 local sizes and both tiles on 57 shapes (see Test Plan). Float on Mali keeps the 2x1 x 4 tile: in the Mali offline compiler the 2x2 x 8 Float tile spills 88 bytes per thread while the 2x1 x 4 Float tile does not spill, and the wide Float tile was 2-31% slower in the sweep. ghstack-source-id: 462666670 @exported-using-ghexport Differential Revision: [D128786678](https://our.internmc.facebook.com/intern/diff/D128786678/)

[ghstack-poisoned]

…convolution Pull Request resolved: #24433 ## Ulterior Motive Closing the latency gap between ET-VK and LiteRT (TFLite GPU) on the DSR denoising model on mobile GPUs. The model's 80 1x1 convolutions were the largest per-operator gap on Mali. ## Rationale **What**: Reworks `conv2d_pw_tiled` around a 2x2-pixel x 8-channel tile with buffer weights, branch-free bounds handling and per-GPU local sizes. The 2x1-pixel x 4-channel tile stays for small Half grids and for Float on Mali, and Float on Mali uses texture weights. **Why**: Pointwise convolutions are load/store bound on Mali. The old tile issued 6 texel loads per 8 vec4 FMAs, read weights through a texture, and guarded every load with a bounds check. The 2x2 x 8 tile issues half as many loads per FMA. ## Details - Tile - **Before**: one tile, 2 consecutive output positions x 4 output channels per invocation. - **After**: a 2x2 output pixel block x 8 output channels (`_4x8` variants) for Float on every GPU except Mali, and for Half when the grid has at least 4096 invocations. Smaller Half grids keep the 2x1 x 4 tile, which fills the GPU better there, and so does Float on Mali. - Weight storage - **Before**: texture2d, unless the packed weights exceed texture limits. - **After**: buffer (`_bufw` variants) for Float on every GPU except Mali, and for Half when the device supports 16-bit storage buffers; texture2d otherwise. On a Mali-G715 in CI, Float buffer weights produced wrong results for padded and strided convolutions with either tile. The same cases passed with the parent's texture weights. - Bounds handling - **Before**: every input and weight load is guarded. - **After**: pixels and channel tiles past the edge read the nearest in-range address and are discarded at the store, so weight loads have no branches. One redundant K check stays on the input load; without any branch in the K loop, the Adreno compiler reads the weight buffer through a much slower path. - Local size - **Before**: 64-invocation square for every dispatch. - **After**: wide tile work groups span the channel tiles of a row of pixel tiles, with 512 threads on Mali, 128 on Adreno and 64 elsewhere, halved until they fit the device limits. The narrow tile uses 1x64 along pixel tiles on Mali when there are at least 64 of them, and the 64-invocation square otherwise. - Tests - The pointwise reference supports Half. `pw_correctness` runs Half too, and adds partial-tile shapes below and above the Half tile switch. - New timing-only `dsr_pw_set` with the DSR model's 1x1 convolution shapes. The tile switch and local sizes come from a sweep over 12 local sizes and both tiles on 57 shapes (see Test Plan). Float on Mali keeps the 2x1 x 4 tile: in the Mali offline compiler the 2x2 x 8 Float tile spills 88 bytes per thread while the 2x1 x 4 Float tile does not spill, and the wide Float tile was 2-31% slower in the sweep. ghstack-source-id: 462948223 @exported-using-ghexport Differential Revision: [D128786678](https://our.internmc.facebook.com/intern/diff/D128786678/)

Update

a90027f

[ghstack-poisoned]

meta-cla Bot 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 Nov 2, 2026

meta-codesync Bot added the meta-exported label Nov 2, 2026

SS-JIA temporarily deployed to cadence November 2, 2026 15:36 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence November 2, 2026 15:51 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence November 2, 2026 15:53 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence November 2, 2026 15:54 — with GitHub Actions Inactive

shoumikhin approved these changes Nov 2, 2026

View reviewed changes

Update

aeb7e9e

[ghstack-poisoned]

SS-JIA temporarily deployed to cadence November 2, 2026 17:05 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence November 2, 2026 17:34 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence November 2, 2026 17:39 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence November 2, 2026 17:46 — with GitHub Actions Inactive

Update

8c9548a

[ghstack-poisoned]

SS-JIA temporarily deployed to cadence November 2, 2026 21:30 — with GitHub Actions Inactive

SS-JIA requested a deployment to cadence November 2, 2026 22:14 — with GitHub Actions In progress

SS-JIA requested a deployment to cadence November 2, 2026 22:17 — with GitHub Actions In progress

SS-JIA requested a deployment to cadence November 2, 2026 22:19 — with GitHub Actions In progress

Navigation Menu

Repository navigation

[ET-VK][conv2d] Use a 2x2x8 tile and tuned local sizes for pointwise convolution - #24433

[ET-VK][conv2d] Use a 2x2x8 tile and tuned local sizes for pointwise convolution#24433
SS-JIA wants to merge 3 commits into
gh/SS-JIA/679/basefrom
gh/SS-JIA/679/head

SS-JIA commented Nov 2, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Nov 2, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Nov 2, 2026

Uh oh!

shoumikhin commented Nov 2, 2026

Uh oh!

claude Bot commented Nov 2, 2026 •

edited

Loading

Summary

Code Quality

Backward Compatibility

Performance

Specific Comments

Recommendation

Uh oh!

This branch is being deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

SS-JIA commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Ulterior Motive

Rationale

Details

Uh oh!

pytorch-bot Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 3 New Failures, 16 Pending, 2 Unclassified Failures

Uh oh!

github-actions Bot commented Nov 2, 2026

This PR needs a release notes: label

Uh oh!

shoumikhin commented Nov 2, 2026

Uh oh!

claude Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24433

Summary

Code Quality

Backward Compatibility

Performance

Specific Comments

Recommendation

Uh oh!

This branch is being deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

SS-JIA commented Nov 2, 2026 •

edited

Loading

pytorch-bot Bot commented Nov 2, 2026 •

edited

Loading

This PR needs a `release notes:` label

claude Bot commented Nov 2, 2026 •

edited

Loading