Skip to content

Commit ba7afe9

Browse files
committed
[ET-VK] Lower eligible conv1d as conv2d over a singleton height dim
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 #22329
1 parent 500849b commit ba7afe9

5 files changed

Lines changed: 203 additions & 0 deletions

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: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
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+
transposed, groups = node.args[6], node.args[8]
76+
if transposed or groups != 1:
77+
return False
78+
# conv2d does not support batched input at all, so a rewrite there
79+
# would turn a working conv1d into a throw.
80+
if in_shape[0] != 1:
81+
return False
82+
# Depthwise and pointwise 1-D convs keep their own shaders.
83+
if w_shape[0] < _IM2COL_MIN_C_OUT or w_shape[2] == 1:
84+
return False
85+
# im2col requires unit dilation.
86+
return list(node.args[5]) == [1]
87+
88+
def _rewrite(self, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
89+
assert self._exported_program is not None
90+
view = exir_ops.edge.aten.view_copy.default
91+
in_node, weight_node = node.args[0], node.args[1]
92+
93+
in_shape = list(in_node.meta["val"].shape)
94+
w_shape = list(weight_node.meta["val"].shape)
95+
out_shape = list(node.meta["val"].shape)
96+
in_4d_shape = [in_shape[0], in_shape[1], 1, in_shape[2]]
97+
w_4d_shape = [w_shape[0], w_shape[1], 1, w_shape[2]]
98+
out_4d_shape = [out_shape[0], out_shape[1], 1, out_shape[2]]
99+
100+
# The weight tensor is contiguous, so inserting a singleton dim is a
101+
# pure metadata change. A weight shared by several convs is only
102+
# reshaped once; the rank check in _eligible() skips it afterwards.
103+
weight = get_param_tensor(self._exported_program, weight_node)
104+
assert weight is not None
105+
set_param_tensor(
106+
self._exported_program, weight_node, weight.reshape(w_4d_shape)
107+
)
108+
weight_node.meta["val"] = _reshape_fake(
109+
weight_node.meta["val"], w_4d_shape
110+
)
111+
112+
with graph.inserting_before(node):
113+
in_4d = graph.call_function(view, args=(in_node, in_4d_shape))
114+
in_4d.meta = dict(in_node.meta)
115+
in_4d.meta["val"] = _reshape_fake(in_node.meta["val"], in_4d_shape)
116+
117+
def pad(arg: List[int], lead: int) -> List[int]:
118+
return [lead, arg[0]]
119+
120+
node.args = (
121+
in_4d,
122+
weight_node,
123+
node.args[2],
124+
pad(node.args[3], 1), # stride
125+
pad(node.args[4], 0), # padding
126+
pad(node.args[5], 1), # dilation
127+
node.args[6],
128+
pad(node.args[7], 0), # output_padding
129+
node.args[8],
130+
)
131+
132+
with graph.inserting_after(node):
133+
out_3d = graph.call_function(view, args=(node, out_shape))
134+
out_3d.meta = dict(node.meta)
135+
node.replace_all_uses_with(out_3d)
136+
out_3d.args = (node, out_shape)
137+
node.meta["val"] = _reshape_fake(node.meta["val"], out_4d_shape)
138+
139+
def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
140+
assert self._exported_program is not None
141+
142+
modified = False
143+
for node in list(graph_module.graph.nodes):
144+
if node.op != "call_function":
145+
continue
146+
if node.target != exir_ops.edge.aten.convolution.default:
147+
continue
148+
if not self._eligible(node):
149+
continue
150+
self._rewrite(graph_module.graph, node)
151+
modified = True
152+
153+
if not modified:
154+
return PassResult(graph_module, False)
155+
156+
graph_module.recompile()
157+
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/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)