Skip to content

Commit b0883bd

Browse files
authored
[ET-VK] Lower eligible conv1d as conv2d over a singleton height dim (#22330)
Fixes #22329. `conv1d.glsl` computes 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 that `conv1d_gwg` sets `z = 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: | dispatch | useful throughput | | --- | --- | | `conv1d_float`, 80 -> 384, k=3, L=3000 | 31.3 GFLOP/s | | `conv1d_float`, 384 -> 384, k=3, s=2, L=1500 | 27.6 GFLOP/s | | `linear_vec_bias`, M=1500, N=1536, K=384 | 921 GFLOP/s | | `linear_vec_bias`, M=1500, N=384, K=1536 | 900 GFLOP/s | Same GPU, same run. Those two conv1d dispatches were 43% of the encoder's runtime. `conv2d_im2col.glsl` + `conv2d_gemm.glsl` already 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_tensor` rather than with a `view_copy` node, because `add_conv2d_node` reads it as a constant `TensorRef` and a runtime view would hand it a `TENSOR` instead. ### 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 == 1` and unit dilation, matching `should_use_conv2d_im2col` - `out_channels >= kIm2colMinCOut` (128), so the im2col path is selected on every vendor rather than only on Mali - `kernel_size > 1`, leaving pointwise convs on `conv1d_pw` - batch 1, since conv2d throws on batched input and a rewrite would turn a working conv1d into a hard error - a constant weight, so it can be reshaped in place ### Results 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. | arm | median | vs XNNPACK | | --- | --- | --- | | Vulkan before | 153.2 ms | 0.84x | | XNNPACK | 128.0 ms | 1.00x | | Vulkan after | **88.8 ms** | **1.44x** | 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: | | before | after | | --- | --- | --- | | conv dispatches | conv1d 17.7 + 47.7 ms | im2col 0.05 + gemm 0.64 + im2col 0.19 + gemm 1.47 ms | | total | 66.5 ms | 3.28 ms | 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: | case | shader after the pass | rewritten | cosine | max abs diff | | --- | --- | --- | --- | --- | | 80 -> 384, k=3 | `conv2d_im2col` + `conv2d_gemm` | yes | 1.00000000 | 2.0e-06 | | 384 -> 384, k=3, s=2 | `conv2d_im2col` + `conv2d_gemm` | yes | 1.00000000 | 4.1e-06 | | 128 -> 256, k=5, p=0 | `conv2d_im2col` + `conv2d_gemm` | yes | 1.00000012 | 3.1e-06 | | 64 -> 64, k=3 | `conv1d_float` | no | 1.00000000 | 8.3e-07 | | 256 -> 256, k=1 | `conv1d_pw_bias_texture` | no | 1.00000000 | 9.5e-07 | | 256 -> 256, k=3, depthwise | `conv1d_dw_bias_texture` | no | 1.00000012 | 4.8e-07 | | 128 -> 256, k=3, dilation 2 | `conv1d_float` | no | 1.00000000 | 1.2e-06 | | 256 -> 256, k=3, groups 4 | `conv1d_float` | no | 1.00000000 | 9.5e-07 | Every ineligible case keeps its original shader and every case matches the CPU reference. `test_vulkan_backend_conv1d_as_conv2d` covers 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 as `linear`. Distinct outputs over 60 executions on the Adreno 840, with `Linear.cpp` already fixed: | | distinct / 60 | | --- | --- | | Whisper encoder, conv1d frontend | 1 | | Whisper encoder, this PR | 9 | | Whisper conv frontend alone, this PR | 13 | This is not caused by the rewrite. `conv2d_gemm` is already reached today by any conv2d with `out_channels >= 128`, and it is nondeterministic there for the same reason. #22328 now fixes the shared `pick_xy_square_lwg` helper rather than one caller, and with it every model above is 1 / 60. The 88.8 ms figure above was measured with only `Linear.cpp` fixed. 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
1 parent 2512896 commit b0883bd

6 files changed

Lines changed: 314 additions & 1 deletion

File tree

‎backends/vulkan/_passes/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
# pyre-strict
88

9+
from executorch.backends.vulkan._passes.conv1d_as_conv2d import Conv1dAsConv2dPass
910
from executorch.backends.vulkan._passes.fold_qdq import FoldQDQPass
1011
from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass
1112
from executorch.backends.vulkan._passes.fuse_quantized_ops import (
@@ -28,6 +29,7 @@
2829
from executorch.backends.vulkan._passes.tag_memory_meta_pass import TagMemoryMetaPass
2930

3031
__all__ = [
32+
"Conv1dAsConv2dPass",
3133
"FoldQDQPass",
3234
"FusePatternsPass",
3335
"FuseQuantizedOpsTransform",
Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# pyre-strict
8+
9+
from typing import List, Optional
10+
11+
import torch
12+
13+
from executorch.backends.transforms.utils import (
14+
get_param_tensor,
15+
is_param_node,
16+
set_param_tensor,
17+
)
18+
from executorch.exir.dialects._ops import ops as exir_ops
19+
from executorch.exir.pass_base import ExportPass, PassResult
20+
21+
from torch.export import ExportedProgram
22+
23+
# Mirrors kIm2colMinCOut in Convolution.cpp. At or above this output channel
24+
# count, should_use_conv2d_im2col() selects the im2col + GEMM path on every
25+
# vendor, so a rewritten conv is guaranteed to land there rather than on the
26+
# direct conv2d shader, which has not been compared against conv1d.
27+
_IM2COL_MIN_C_OUT = 128
28+
29+
30+
def _reshape_fake(val: torch.Tensor, sizes: List[int]) -> torch.Tensor:
31+
fake_mode = getattr(val, "fake_mode", None)
32+
if fake_mode is None:
33+
return val.reshape(sizes)
34+
with fake_mode:
35+
return val.reshape(sizes)
36+
37+
38+
class Conv1dAsConv2dPass(ExportPass):
39+
"""
40+
Rewrite eligible 1-D convolutions into a 2-D convolution over a singleton
41+
height dimension.
42+
43+
conv1d.glsl computes one output element per invocation with no tiling and no
44+
reuse between invocations, and its global work grid packs texels along the
45+
batch dim, wasting three of every four lanes at batch 1. It measures around
46+
30x slower than the conv2d im2col + GEMM path for the same MAC count.
47+
Expressing the convolution as 2-D lets the existing conv2d machinery
48+
(im2col selection, memory layout tagging, weight prepacking) handle it with
49+
no runtime changes.
50+
51+
Only convolutions that are certain to reach the im2col path are rewritten,
52+
so nothing silently moves onto a shader that has not been compared against
53+
conv1d. The weight is reshaped in place rather than through a view node
54+
because the conv2d implementation reads it as a constant TensorRef.
55+
"""
56+
57+
def __init__(self) -> None:
58+
super().__init__()
59+
self._exported_program: Optional[ExportedProgram] = None
60+
61+
def _eligible(self, node: torch.fx.Node) -> bool:
62+
assert self._exported_program is not None
63+
in_node, weight_node = node.args[0], node.args[1]
64+
if not isinstance(in_node, torch.fx.Node):
65+
return False
66+
if not isinstance(weight_node, torch.fx.Node):
67+
return False
68+
in_shape = in_node.meta["val"].shape
69+
w_shape = weight_node.meta["val"].shape
70+
if len(in_shape) != 3 or len(w_shape) != 3:
71+
return False
72+
# The weight has to be a constant so it can be reshaped in place.
73+
if not is_param_node(self._exported_program, weight_node):
74+
return False
75+
# ...and it has to belong to this convolution alone. The reshape is a
76+
# mutation of shared state: a second consumer would keep reading the
77+
# weight it was built against and find it rank 4, whether that consumer
78+
# is another conv1d this pass then declines to rewrite, or an op that
79+
# has nothing to do with convolution.
80+
if len(weight_node.users) != 1:
81+
return False
82+
transposed, groups = node.args[6], node.args[8]
83+
if transposed or groups != 1:
84+
return False
85+
# conv2d does not support batched input at all, so a rewrite there
86+
# would turn a working conv1d into a throw.
87+
#
88+
# A symbolic batch is rejected rather than compared. Comparing it would
89+
# read the value it happened to be traced with, so a dimension free to
90+
# be 4 at runtime would pass a test for 1 and reach a conv2d that
91+
# cannot serve it. Only a batch that is already a plain 1 is provably 1
92+
# for every input the graph accepts.
93+
if not isinstance(in_shape[0], int) or in_shape[0] != 1:
94+
return False
95+
# Depthwise and pointwise 1-D convs keep their own shaders.
96+
if w_shape[0] < _IM2COL_MIN_C_OUT or w_shape[2] == 1:
97+
return False
98+
# im2col requires unit dilation.
99+
return list(node.args[5]) == [1]
100+
101+
def _rewrite(self, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
102+
assert self._exported_program is not None
103+
view = exir_ops.edge.aten.view_copy.default
104+
in_node, weight_node = node.args[0], node.args[1]
105+
106+
in_shape = list(in_node.meta["val"].shape)
107+
w_shape = list(weight_node.meta["val"].shape)
108+
out_shape = list(node.meta["val"].shape)
109+
in_4d_shape = [in_shape[0], in_shape[1], 1, in_shape[2]]
110+
w_4d_shape = [w_shape[0], w_shape[1], 1, w_shape[2]]
111+
out_4d_shape = [out_shape[0], out_shape[1], 1, out_shape[2]]
112+
113+
# The weight tensor is contiguous, so inserting a singleton dim is a
114+
# pure metadata change. It is safe to do in place only because
115+
# _eligible() has established this convolution is the weight's sole
116+
# consumer.
117+
weight = get_param_tensor(self._exported_program, weight_node)
118+
assert weight is not None
119+
set_param_tensor(
120+
self._exported_program, weight_node, weight.reshape(w_4d_shape)
121+
)
122+
weight_node.meta["val"] = _reshape_fake(weight_node.meta["val"], w_4d_shape)
123+
124+
with graph.inserting_before(node):
125+
in_4d = graph.call_function(view, args=(in_node, in_4d_shape))
126+
in_4d.meta = dict(in_node.meta)
127+
in_4d.meta["val"] = _reshape_fake(in_node.meta["val"], in_4d_shape)
128+
129+
def pad(arg: List[int], lead: int) -> List[int]:
130+
return [lead, arg[0]]
131+
132+
node.args = (
133+
in_4d,
134+
weight_node,
135+
node.args[2],
136+
pad(node.args[3], 1), # stride
137+
pad(node.args[4], 0), # padding
138+
pad(node.args[5], 1), # dilation
139+
node.args[6],
140+
pad(node.args[7], 0), # output_padding
141+
node.args[8],
142+
)
143+
144+
with graph.inserting_after(node):
145+
out_3d = graph.call_function(view, args=(node, out_shape))
146+
out_3d.meta = dict(node.meta)
147+
node.replace_all_uses_with(out_3d)
148+
out_3d.args = (node, out_shape)
149+
node.meta["val"] = _reshape_fake(node.meta["val"], out_4d_shape)
150+
151+
def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
152+
assert self._exported_program is not None
153+
154+
modified = False
155+
for node in list(graph_module.graph.nodes):
156+
if node.op != "call_function":
157+
continue
158+
if node.target != exir_ops.edge.aten.convolution.default:
159+
continue
160+
if not self._eligible(node):
161+
continue
162+
self._rewrite(graph_module.graph, node)
163+
modified = True
164+
165+
if not modified:
166+
return PassResult(graph_module, False)
167+
168+
graph_module.recompile()
169+
return PassResult(graph_module, True)

‎backends/vulkan/_passes/targets.bzl‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,20 @@ def define_common_targets(is_fbcode = False):
9393
],
9494
)
9595

96+
runtime.python_library(
97+
name = "conv1d_as_conv2d",
98+
srcs = ["conv1d_as_conv2d.py"],
99+
visibility = [
100+
"//executorch/backends/...",
101+
],
102+
deps = [
103+
"//caffe2:torch",
104+
"//executorch/backends/transforms:utils",
105+
"//executorch/exir:pass_base",
106+
"//executorch/exir/dialects:lib",
107+
],
108+
)
109+
96110
runtime.python_library(
97111
name = "fold_qdq",
98112
srcs = ["fold_qdq.py"],
@@ -144,6 +158,7 @@ def define_common_targets(is_fbcode = False):
144158
"//executorch/examples/...",
145159
],
146160
deps = [
161+
":conv1d_as_conv2d",
147162
":fold_qdq",
148163
":fuse_patterns",
149164
":fuse_quantized_ops",

‎backends/vulkan/test/test_vulkan_delegate.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1094,6 +1094,33 @@ def forward(self, x):
10941094
(torch.randn(size=(64,), dtype=torch.float32),),
10951095
)
10961096

1097+
def test_vulkan_backend_conv1d_as_conv2d(self):
1098+
# Eligible for Conv1dAsConv2dPass: groups 1, unit dilation, kernel > 1,
1099+
# batch 1 and out_channels >= the im2col threshold, so this lowers via
1100+
# the conv2d im2col + GEMM path rather than conv1d.glsl.
1101+
class Conv1dModule(torch.nn.Module):
1102+
def __init__(self):
1103+
super().__init__()
1104+
self.conv = torch.nn.Conv1d(
1105+
in_channels=32,
1106+
out_channels=128,
1107+
kernel_size=3,
1108+
stride=2,
1109+
padding=1,
1110+
bias=True,
1111+
)
1112+
1113+
def forward(self, x):
1114+
return self.conv(x)
1115+
1116+
conv1d_module = Conv1dModule()
1117+
sample_inputs = (torch.randn(size=(1, 32, 64), dtype=torch.float32),)
1118+
1119+
self.lower_module_and_test_output(
1120+
conv1d_module,
1121+
sample_inputs,
1122+
)
1123+
10971124
@disable_test("layer norm compute shader not working with swiftshader")
10981125
def test_vulkan_backend_native_layer_norm(self):
10991126
class NativeLayerNormModule(torch.nn.Module):

‎backends/vulkan/test/test_vulkan_passes.py‎

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import unittest
2-
from typing import Optional, Tuple
2+
from typing import List, Optional, Tuple
33

44
import torch
55

6+
from executorch.backends.vulkan._passes.conv1d_as_conv2d import Conv1dAsConv2dPass
67
from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass
78

89
from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge
@@ -86,7 +87,104 @@ def op_node_count(graph_module: torch.fx.GraphModule, canonical_op_name: str) ->
8687
return count
8788

8889

90+
def run_conv1d_as_conv2d(
91+
model: torch.nn.Module,
92+
sample_inputs: Tuple[torch.Tensor],
93+
dynamic_shapes=None,
94+
) -> torch.fx.GraphModule:
95+
"""Exports `model` to edge and runs Conv1dAsConv2dPass over it."""
96+
program = torch.export.export(
97+
model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True
98+
)
99+
edge_program = to_edge(
100+
program,
101+
compile_config=EdgeCompileConfig(
102+
_skip_dim_order=False, _check_ir_validity=False
103+
),
104+
)
105+
exported = edge_program.exported_program()
106+
107+
conv_pass = Conv1dAsConv2dPass()
108+
conv_pass._exported_program = exported
109+
return conv_pass(exported.graph_module).graph_module
110+
111+
112+
def conv_input_ranks(graph_module: torch.fx.GraphModule) -> List[int]:
113+
"""The rank of the input to every convolution in the graph."""
114+
ranks = []
115+
for node in graph_module.graph.nodes:
116+
if get_target_canonical_name(node) == "convolution.default":
117+
ranks.append(len(node.args[0].meta["val"].shape))
118+
return ranks
119+
120+
89121
class TestVulkanPasses(unittest.TestCase):
122+
def test_conv1d_as_conv2d_rewrites_eligible_conv(self):
123+
# Output equality cannot show this: an unrewritten conv1d computes the
124+
# same numbers. What distinguishes the two is the rank the convolution
125+
# is asked to work at, so that is what is asserted.
126+
class Conv1dModule(torch.nn.Module):
127+
def __init__(self):
128+
super().__init__()
129+
self.conv = torch.nn.Conv1d(32, 128, kernel_size=3, padding=1)
130+
131+
def forward(self, x):
132+
return self.conv(x)
133+
134+
gm = run_conv1d_as_conv2d(Conv1dModule(), (torch.randn(1, 32, 64),))
135+
136+
self.assertEqual(conv_input_ranks(gm), [4])
137+
# Two views: one lifting the input to 4-D, one lowering the output back.
138+
self.assertEqual(op_node_count(gm, "view_copy.default"), 2)
139+
140+
def test_conv1d_as_conv2d_rejects_shared_weight(self):
141+
# The rewrite reshapes the weight in place. A second consumer would go
142+
# on reading a tensor that has silently become rank 4, so a shared
143+
# weight has to be left alone.
144+
class SharedWeightModule(torch.nn.Module):
145+
def __init__(self):
146+
super().__init__()
147+
self.weight = torch.nn.Parameter(torch.randn(128, 32, 3))
148+
149+
def forward(self, x, y):
150+
return (
151+
torch.nn.functional.conv1d(x, self.weight, padding=1),
152+
torch.nn.functional.conv1d(y, self.weight, padding=1),
153+
)
154+
155+
gm = run_conv1d_as_conv2d(
156+
SharedWeightModule(), (torch.randn(1, 32, 64), torch.randn(1, 32, 64))
157+
)
158+
159+
self.assertEqual(conv_input_ranks(gm), [3, 3])
160+
self.assertEqual(op_node_count(gm, "view_copy.default"), 0)
161+
162+
def test_conv1d_as_conv2d_rejects_symbolic_batch(self):
163+
# A batch free to vary at runtime is not a batch of 1, whatever value it
164+
# was traced with, and conv2d refuses batched input. The guard is on the
165+
# dimension being a plain integer rather than on its traced value, so
166+
# that a comparison never decides this from the hint alone.
167+
#
168+
# Traced at 2 because export specializes a size-1 dimension: the batch
169+
# has to be genuinely symbolic here for the graph to be the one the
170+
# guard exists for.
171+
class Conv1dModule(torch.nn.Module):
172+
def __init__(self):
173+
super().__init__()
174+
self.conv = torch.nn.Conv1d(32, 128, kernel_size=3, padding=1)
175+
176+
def forward(self, x):
177+
return self.conv(x)
178+
179+
gm = run_conv1d_as_conv2d(
180+
Conv1dModule(),
181+
(torch.randn(2, 32, 64),),
182+
dynamic_shapes={"x": {0: torch.export.Dim.AUTO}},
183+
)
184+
185+
self.assertEqual(conv_input_ranks(gm), [3])
186+
self.assertEqual(op_node_count(gm, "view_copy.default"), 0)
187+
90188
def test_fuse_torchao_quantized_embedding(self):
91189
"""A torchao-dialect 4-bit weight-only quantized embedding
92190
(torchao.dequantize_affine -> aten.embedding) should fuse into a single

‎backends/vulkan/vulkan_preprocess.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
ViewCopyToSqueezeUnsqueezePass,
1818
)
1919
from executorch.backends.vulkan._passes import (
20+
Conv1dAsConv2dPass,
2021
FoldQDQPass,
2122
FuseQuantizedOpsTransform,
2223
insert_prepack_nodes,
@@ -190,6 +191,7 @@ def preprocess( # noqa: C901
190191
AddmmToLinearTransform(),
191192
InsertDtypePromotionPass(),
192193
FusePatternsPass(),
194+
Conv1dAsConv2dPass(),
193195
FuseClampPass(),
194196
RemoveRedundantOpsTransform(),
195197
FuseQuantizedOpsTransform(),

0 commit comments

Comments
 (0)