Repository navigation
[ET-VK] Lower eligible conv1d as conv2d over a singleton height dim - #22330
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22330
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 5851f5a with merge base 500849b ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
conv1d.glsl has no tiling and no reuse between invocations and packs its texels along the batch dim, so it runs at roughly 3 percent of the throughput the same GPU reaches on a matmul of equal MAC count. Route eligible 1-D convolutions through the existing conv2d im2col + GEMM path instead. Also backports set_param_tensor from upstream, which the pass uses to reshape the weight in place. Upstream: pytorch/executorch#22329, pytorch/executorch#22330
SS-JIA
left a comment
SS-JIA
left a comment
There was a problem hiding this comment.
AI-assisted review. Inline findings below.
| return False | ||
| # conv2d does not support batched input at all, so a rewrite there | ||
| # would turn a working conv1d into a throw. | ||
| if in_shape[0] != 1: |
There was a problem hiding this comment.
[High] For a symbolic batch dimension, this comparison can specialize from the traced value 1 without proving every allowed runtime value is 1. The rewritten conv2d path rejects batched input, so a valid larger runtime batch can fail. Require a shape-range proof with lower and upper bounds both equal to 1; otherwise skip the rewrite.
AI reviewer note: this comment was generated by an AI reviewer.
There was a problem hiding this comment.
Fixed: the pass now bails unless the batch dim is a concrete 1, so a symbolic batch stays on the conv1d path. Covered by test_conv1d_as_conv2d_rejects_symbolic_batch.
| # reshaped once; the rank check in _eligible() skips it afterwards. | ||
| weight = get_param_tensor(self._exported_program, weight_node) | ||
| assert weight is not None | ||
| set_param_tensor( |
There was a problem hiding this comment.
[High] This mutates the shared parameter before validating all consumers. If one eligible conv and one ineligible op share the weight, the first rewrite changes the weight to rank 4 and leaves the later consumer using the wrong shape. Clone per rewritten consumer or require all consumers to be safely rewritten before mutating shared state.
AI reviewer note: this comment was generated by an AI reviewer.
There was a problem hiding this comment.
Fixed: the rewrite is skipped when the weight has more than one user, so a shared parameter is never reshaped out from under another consumer. Covered by test_conv1d_as_conv2d_rejects_shared_weight.
| conv1d_module = Conv1dModule() | ||
| sample_inputs = (torch.randn(size=(1, 32, 64), dtype=torch.float32),) | ||
|
|
||
| self.lower_module_and_test_output( |
There was a problem hiding this comment.
[Medium] Output equality does not prove Conv1dAsConv2dPass ran; the unchanged conv1d path can satisfy this test. Assert the transformed graph contains the intended 4-D convolution/views and add shared-weight plus symbolic-batch rejection cases.
AI reviewer note: this comment was generated by an AI reviewer.
There was a problem hiding this comment.
Added graph-level tests in test_vulkan_passes.py: test_conv1d_as_conv2d_rewrites_eligible_conv asserts the convolution input is rank 4 and two view_copy nodes were inserted, plus the shared-weight and symbolic-batch rejection cases.
2406ac3 to
74d6429
Compare
74d6429 to
8ec2a34
Compare
|
@msluszniak looks good, but lintrunner is failing |
|
Pushed a ufmt fix for |
|
@pytorchbot label "module: vulkan" |
pytorch-bot
Bot
added
the
module: vulkan
conv1d.glsl computes one output element per invocation with no tiling and no reuse between invocations, and its work grid packs texels along the batch dim, wasting three of every four lanes at batch 1. On an Adreno 840 it reaches about 30 GFLOP/s where the tiled linear reaches 900 GFLOP/s in the same graph. The conv2d im2col + GEMM path already solves this, so rewrite 1-D convolutions that are certain to reach it into a 2-D convolution over a singleton height dim. Only convs with groups 1, unit dilation, kernel > 1, batch 1 and out_channels >= kIm2colMinCOut are rewritten, so nothing moves onto a shader that has not been compared against conv1d. The Whisper-tiny encoder goes from 153.2 ms to 88.8 ms, from 0.84x XNNPACK to 1.44x. Fixes pytorch#22329
… use Two ways the rewrite could reach a graph it does not own, both from review. The weight is reshaped in place, so a weight with a second consumer is left behind at rank 4 while that consumer still expects rank 3. The comment claimed the rank check in _eligible() covered this; it does not. It stops the second conv1d being rewritten, which is precisely the problem: the conv keeps its 1-D arguments and loses its 1-D weight. Only rewrite when this convolution is the weight's sole consumer. A symbolic batch was compared against 1, which reads the value it was traced with rather than the values the graph accepts. A dimension free to be 4 at runtime could pass a test for 1 and reach a conv2d that rejects batched input. Require a batch that is already a plain 1. Tests move to test_vulkan_passes.py, where the graph can be inspected without a device. The delegate test only compared outputs, which an unrewritten conv1d satisfies just as well, so it could not have failed if the pass stopped firing. The shared-weight test fails against the previous pass. The symbolic-batch test does not: export specializes a size-1 dimension, so a batch that is both symbolic and hinted 1 is not constructible here, and a batch hinted higher is rejected either way. It pins the guard rather than reproducing the defect.
ed4b665 to
5851f5a
Compare
Fixes #22329.
conv1d.glslcomputes one output element per invocation with no tiling, no shared memory and no register blocking, and every tap re-fetches both a weight texel and an input texel, so nothing is reused between invocations. On top of thatconv1d_gwgsetsz = div_up_4(batch), so the texel's four lanes span the batch dimension and three of every four are dead at batch 1.Measured inside a single Whisper-tiny encoder execution on an Adreno 840, using the shader query pool:
conv1d_float, 80 -> 384, k=3, L=3000conv1d_float, 384 -> 384, k=3, s=2, L=1500linear_vec_bias, M=1500, N=1536, K=384linear_vec_bias, M=1500, N=384, K=1536Same GPU, same run. Those two conv1d dispatches were 43% of the encoder's runtime.
conv2d_im2col.glsl+conv2d_gemm.glslalready solve this for 2-D convolutions. This pass rewrites 1-D convolutions into a 2-D convolution over a singleton height dim so the existing conv2d machinery (im2col selection, memory layout tagging, weight prepacking) handles them with no runtime changes.The weight is reshaped in place through
set_param_tensorrather than with aview_copynode, becauseadd_conv2d_nodereads it as a constantTensorRefand a runtime view would hand it aTENSORinstead.Eligibility
Only convolutions that are certain to reach the im2col path are rewritten, so nothing silently moves onto a shader that has not been compared against conv1d:
groups == 1and unit dilation, matchingshould_use_conv2d_im2colout_channels >= kIm2colMinCOut(128), so the im2col path is selected on every vendor rather than only on Malikernel_size > 1, leaving pointwise convs onconv1d_pwResults
Whisper-tiny encoder, Samsung Galaxy S26 Ultra (Adreno 840), fp32, interleaved A/B with the order reversed each round, warm-up iteration discarded, 60 timed iterations per arm. Both Vulkan arms include #22328, without which the encoder is not deterministic on this GPU at all.
Cosine against the CPU reference is 0.99999684 after the change.
Isolated, just the encoder's two convs plus gelu, GPU time from the query pool:
20x on the frontend, and the output is unchanged at cosine 1.00000000, max abs diff 4.9e-04 against the CPU reference in both.
Routing and correctness matrix
Eight conv1d configurations exported and run on device, with the dispatched shader read back from the query pool:
conv2d_im2col+conv2d_gemmconv2d_im2col+conv2d_gemmconv2d_im2col+conv2d_gemmconv1d_floatconv1d_pw_bias_textureconv1d_dw_bias_textureconv1d_floatconv1d_floatEvery ineligible case keeps its original shader and every case matches the CPU reference.
test_vulkan_backend_conv1d_as_conv2dcovers the rewritten path in CI; the two existing conv1d tests are grouped and depthwise, so they only exercise the untouched path.Depends on #22328
Moving the conv frontend onto the im2col + GEMM path exposes a pre-existing Adreno defect in
conv2d_gemm, which dispatches with the same near-square local work group aslinear. Distinct outputs over 60 executions on the Adreno 840, withLinear.cppalready fixed:This is not caused by the rewrite.
conv2d_gemmis already reached today by any conv2d without_channels >= 128, and it is nondeterministic there for the same reason. #22328 now fixes the sharedpick_xy_square_lwghelper rather than one caller, and with it every model above is 1 / 60.The 88.8 ms figure above was measured with only
Linear.cppfixed. With #22328 applied in full the encoder is 91.4 ms, still 1.34x XNNPACK, and correct on every run.Why rewrite rather than write a tiled conv1d shader
The 384 -> 384 dispatch moves 21.2 GB of texture traffic for 8.7 MB of compulsory traffic, so the shader is texture bound with essentially no reuse, and the batch-in-lanes waste is a 4x term next to that. A native tiled conv1d would fix the reuse the same way the GEMM does, and would additionally save the im2col materialisation, which is 0.25 ms of the 2.64 ms frontend above, roughly 10%. This pass gets the rest of it by reusing kernels that are already written, tuned and tested, and leaves a dedicated conv1d kernel available later as an optimisation rather than a prerequisite.
cc @SS-JIA @manuelcandales @digantdesai @cbilgin