Skip to content

Commit 43bdb7b

Browse files
robellrobell
authored
Arm backend: Fix grid sampler dispatch shape (#20961)
Arm backend: Fix grid sampler dispatch shape Revised rules for Arm custom ops and the dispatch shape, also applied to grid sample implementation Signed-off-by: Rob Elliott <Robert.Elliott@arm.com> Change-Id: I4728474acfc882a7cc13a2478e03c397b44ffbf2 cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @rascani --------- Signed-off-by: Rob Elliott <Robert.Elliott@arm.com>
1 parent 0f02e85 commit 43bdb7b

9 files changed

Lines changed: 218 additions & 38 deletions

‎backends/arm/test/misc/test_custom_shader_payload.py‎

Lines changed: 57 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
GRID_SAMPLER_2D_SHADER_LANGUAGE,
2828
GRID_SAMPLER_2D_SHADER_SOURCE,
2929
GRID_SAMPLER_2D_VK_FORMAT,
30-
GRID_SAMPLER_2D_WORKGROUP_SIZES,
3130
)
3231

3332

@@ -45,11 +44,12 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_round_trip():
4544
interpolation_mode=0,
4645
padding_mode=2,
4746
align_corners=True,
47+
output_shape=(1, 4, 8, 8),
4848
)
4949
decoded = decode_payload(encode_payload(payload))
5050

5151
assert decoded["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT
52-
assert decoded["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES
52+
assert decoded["workgroup_sizes"] == [1, 1, 1]
5353
assert decoded["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
5454
assert base64.b64decode(decoded["shader_code"])[:4] == b"\x03\x02\x23\x07"
5555
assert decoded["input_0_type"] == "Tensor"
@@ -58,7 +58,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_round_trip():
5858
assert decoded["input_0_binding"] == 0
5959
assert decoded["input_1_type"] == "Tensor"
6060
assert decoded["input_1_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
61-
assert decoded["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
61+
assert decoded["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM"
6262
assert decoded["input_1_binding"] == 1
6363
assert decoded["output_0_type"] == "Tensor"
6464
assert decoded["output_0_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
@@ -72,10 +72,12 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_uses_sampler_for_c4():
7272
padding_mode=0,
7373
align_corners=False,
7474
input_shape=(1, 4, 8, 8),
75+
output_shape=(1, 4, 4, 4),
7576
input_dtype=torch.float32,
7677
)
7778

7879
assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
80+
assert payload["workgroup_sizes"] == [1, 1, 1]
7981
assert base64.b64decode(payload["shader_code"])[:4] == b"\x03\x02\x23\x07"
8082
assert payload["input_0_type"] == "Image"
8183
assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT
@@ -104,11 +106,13 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_uses_int8_sampler_for_c
104106
padding_mode=0,
105107
align_corners=False,
106108
input_shape=(1, 4, 8, 8),
109+
output_shape=(1, 4, 4, 4),
107110
input_dtype=torch.int8,
108111
output_dtype=torch.int8,
109112
)
110113

111114
assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
115+
assert payload["workgroup_sizes"] == [1, 1, 1]
112116
assert base64.b64decode(payload["shader_code"])[:4] == b"\x03\x02\x23\x07"
113117
assert payload["input_0_type"] == "Image"
114118
assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_INT8_VK_FORMAT
@@ -130,23 +134,58 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_keeps_c3_on_buffer():
130134
padding_mode=0,
131135
align_corners=False,
132136
input_shape=(1, 3, 8, 8),
137+
output_shape=(1, 3, 4, 4),
133138
input_dtype=torch.float32,
134139
)
135140

136141
assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
142+
assert payload["workgroup_sizes"] == [1, 1, 1]
137143
assert payload["input_0_type"] == "Tensor"
138144
assert payload["input_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
145+
assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM"
139146
assert payload["output_0_type"] == "Tensor"
140147
assert payload["output_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
141148
assert "input_0_sampler" not in payload
142149

143150

151+
def test_grid_sampler_2d_custom_shader_payload_sampler_dispatch_rounds_up_output():
152+
payload = build_grid_sampler_2d_payload(
153+
interpolation_mode=0,
154+
padding_mode=0,
155+
align_corners=False,
156+
input_shape=(1, 4, 32, 32),
157+
output_shape=(1, 4, 17, 9),
158+
input_dtype=torch.float32,
159+
)
160+
161+
assert payload["input_0_type"] == "Image"
162+
assert payload["output_0_type"] == "Image"
163+
assert payload["workgroup_sizes"] == [2, 3, 1]
164+
165+
166+
def test_grid_sampler_2d_custom_shader_payload_buffer_dispatch_rounds_up_output():
167+
payload = build_grid_sampler_2d_payload(
168+
interpolation_mode=2,
169+
padding_mode=0,
170+
align_corners=False,
171+
input_shape=(1, 4, 32, 32),
172+
output_shape=(1, 4, 17, 9),
173+
input_dtype=torch.float32,
174+
)
175+
176+
assert payload["input_0_type"] == "Tensor"
177+
assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM"
178+
assert payload["output_0_type"] == "Tensor"
179+
assert payload["workgroup_sizes"] == [2, 3, 1]
180+
181+
144182
def test_grid_sampler_2d_custom_shader_payload_no_target_align_corners_sampler():
145183
payload = build_grid_sampler_2d_payload(
146184
interpolation_mode=0,
147185
padding_mode=0,
148186
align_corners=True,
149187
input_shape=(1, 4, 8, 8),
188+
output_shape=(1, 4, 8, 8),
150189
input_dtype=torch.float32,
151190
)
152191

@@ -170,6 +209,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_int8_align_corners_samp
170209
padding_mode=0,
171210
align_corners=True,
172211
input_shape=(1, 4, 8, 8),
212+
output_shape=(1, 4, 8, 8),
173213
input_dtype=torch.int8,
174214
output_dtype=torch.int8,
175215
)
@@ -194,13 +234,15 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_bicubic_buffer():
194234
padding_mode=0,
195235
align_corners=False,
196236
input_shape=(1, 4, 8, 8),
237+
output_shape=(1, 4, 8, 8),
197238
input_dtype=torch.float32,
198239
)
199240

200241
assert payload["input_0_type"] == "Tensor"
201242
assert payload["input_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
202243
assert payload["output_0_type"] == "Tensor"
203244
assert payload["output_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
245+
assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM"
204246
assert "input_0_sampler" not in payload
205247

206248

@@ -209,6 +251,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_uses_spirv():
209251
interpolation_mode=0,
210252
padding_mode=0,
211253
align_corners=False,
254+
output_shape=(1, 4, 8, 8),
212255
)
213256

214257
shader_binary = base64.b64decode(payload["shader_code"])
@@ -253,11 +296,22 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_rejects_bad_modes():
253296
interpolation_mode=99,
254297
padding_mode=0,
255298
align_corners=False,
299+
output_shape=(1, 4, 8, 8),
256300
)
257301

258302
with pytest.raises(ValueError, match="Unsupported padding_mode"):
259303
build_grid_sampler_2d_payload(
260304
interpolation_mode=0,
261305
padding_mode=99,
262306
align_corners=False,
307+
output_shape=(1, 4, 8, 8),
308+
)
309+
310+
311+
def test_grid_sampler_2d_custom_shader_payload_requires_output_shape():
312+
with pytest.raises(ValueError, match="requires output_shape for dispatch"):
313+
build_grid_sampler_2d_payload(
314+
interpolation_mode=0,
315+
padding_mode=0,
316+
align_corners=False,
263317
)

‎backends/arm/test/misc/test_custom_shader_payloads.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
GRID_SAMPLER_2D_SHADER_LANGUAGE,
3434
GRID_SAMPLER_2D_SHADER_SOURCE,
3535
GRID_SAMPLER_2D_VK_FORMAT,
36-
GRID_SAMPLER_2D_WORKGROUP_SIZES,
3736
)
3837
from torch.export import export
3938

@@ -107,12 +106,13 @@ def test_buffer_shader_payload_vgf_encodes_bindings_and_formats():
107106
interpolation_mode=0,
108107
padding_mode=0,
109108
align_corners=False,
109+
output_shape=(1, 4, 8, 8),
110110
)
111111
)
112112
)
113113

114114
assert payload["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT
115-
assert payload["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES
115+
assert payload["workgroup_sizes"] == [1, 1, 1]
116116
assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
117117
assert payload["input_0_binding"] == 0
118118
assert payload["input_1_binding"] == 1
@@ -150,6 +150,7 @@ def test_shader_payload_vgf_uses_expected_glsl_and_spirv_asset():
150150
interpolation_mode=0,
151151
padding_mode=0,
152152
align_corners=False,
153+
output_shape=(1, 4, 8, 8),
153154
)
154155

155156
assert GRID_SAMPLER_2D_SHADER_SOURCE == "grid_sampler.glsl"

‎backends/arm/test/ops/test_custom_shader_lowering.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
RewriteGridSamplerToTosaCustomPass,
3838
)
3939
from executorch.backends.arm.vgf.shaders.grid_sampler import (
40+
_dispatch_shape_for_output_shape,
4041
decode_payload,
4142
grid_sampler_2d_operator_name,
4243
)
@@ -256,3 +257,31 @@ def test_shader_lowering_vgf_decodes_expected_implementation_attrs():
256257
assert payload["input_0_binding"] == 0
257258
assert payload["input_1_binding"] == 1
258259
assert payload["output_0_binding"] == 2
260+
261+
262+
def test_grid_sampler_dispatch_shape_includes_batch_dimension():
263+
assert _dispatch_shape_for_output_shape((2, 3, 5, 9)) == [2, 1, 2]
264+
265+
266+
def test_in_tree_grid_sampler_vgf_payload_uses_batched_dispatch_shape():
267+
edge_model = to_edge(
268+
export(
269+
_GridSampleModule(),
270+
(
271+
torch.randn(2, 3, 11, 13),
272+
torch.randn(2, 5, 9, 2),
273+
),
274+
)
275+
)
276+
277+
with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
278+
transformed = edge_model.transform([RewriteGridSamplerToTosaCustomPass()])
279+
280+
custom_node = next(
281+
node
282+
for node in transformed.exported_program().graph.nodes
283+
if node.target == exir_ops.backend.tosa.CUSTOM.default
284+
)
285+
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
286+
287+
assert payload["workgroup_sizes"] == [2, 1, 2]

‎backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py‎

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
GRID_SAMPLER_2D_SHADER_ENTRY_POINT,
2929
GRID_SAMPLER_2D_SHADER_LANGUAGE,
3030
GRID_SAMPLER_2D_VK_FORMAT,
31-
GRID_SAMPLER_2D_WORKGROUP_SIZES,
3231
)
3332
from executorch.exir import to_edge
3433
from executorch.exir.dialects._ops import ops as exir_ops
@@ -88,7 +87,7 @@ def test_rewrite_grid_sampler_to_tosa_custom_vgf_no_target():
8887

8988
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
9089
assert payload["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT
91-
assert payload["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES
90+
assert payload["workgroup_sizes"] == [1, 1, 1]
9291
assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
9392
assert payload["input_0_type"] == "Image"
9493
assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT
@@ -111,6 +110,28 @@ def test_rewrite_grid_sampler_to_tosa_custom_vgf_no_target():
111110
assert any(node.target == exir_ops.edge.aten.slice_copy.Tensor for node in nodes)
112111

113112

113+
def test_rewrite_grid_sampler_to_tosa_custom_sampler_dispatch_rounds_up_output():
114+
model = GridSampler2d()
115+
example_inputs = (
116+
torch.randn(1, 4, 32, 32),
117+
torch.randn(1, 17, 9, 2),
118+
)
119+
120+
edge_model = to_edge(export(model, example_inputs))
121+
with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
122+
edge_model = edge_model.transform([RewriteGridSamplerToTosaCustomPass()])
123+
nodes = list(edge_model.exported_program().graph.nodes)
124+
125+
custom_node = next(
126+
node for node in nodes if node.target == exir_ops.backend.tosa.CUSTOM.default
127+
)
128+
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
129+
130+
assert payload["input_0_type"] == "Image"
131+
assert payload["output_0_type"] == "Image"
132+
assert payload["workgroup_sizes"] == [2, 3, 1]
133+
134+
114135
def test_rewrite_grid_sampler_to_tosa_custom_no_target_uses_sampler_for_c4():
115136
model = GridSampler2d()
116137
example_inputs = (
@@ -129,6 +150,7 @@ def test_rewrite_grid_sampler_to_tosa_custom_no_target_uses_sampler_for_c4():
129150
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
130151

131152
assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
153+
assert payload["workgroup_sizes"] == [1, 1, 1]
132154
assert (
133155
payload["input_0_vkdescriptortype"]
134156
== "VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER"
@@ -202,6 +224,7 @@ def test_rewrite_grid_sampler_to_tosa_custom_c3_pad_for_align_corners():
202224
)
203225
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
204226

227+
assert payload["workgroup_sizes"] == [1, 1, 1]
205228
assert payload["input_0_type"] == "Image"
206229
assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT
207230
assert (
@@ -233,7 +256,33 @@ def test_rewrite_grid_sampler_to_tosa_custom_no_c3_pad_for_bicubic():
233256
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
234257

235258
assert payload["input_0_type"] == "Tensor"
259+
assert payload["workgroup_sizes"] == [1, 1, 1]
260+
assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM"
236261
assert not any(node.target == exir_ops.edge.aten.cat.default for node in nodes)
237262
assert not any(
238263
node.target == exir_ops.edge.aten.slice_copy.Tensor for node in nodes
239264
)
265+
266+
267+
def test_rewrite_grid_sampler_to_tosa_custom_buffer_dispatch_rounds_up_output():
268+
model = GridSampler2d()
269+
model.interpolation_mode_ = 2
270+
example_inputs = (
271+
torch.randn(1, 4, 32, 32),
272+
torch.randn(1, 17, 9, 2),
273+
)
274+
275+
edge_model = to_edge(export(model, example_inputs))
276+
with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
277+
edge_model = edge_model.transform([RewriteGridSamplerToTosaCustomPass()])
278+
nodes = list(edge_model.exported_program().graph.nodes)
279+
280+
custom_node = next(
281+
node for node in nodes if node.target == exir_ops.backend.tosa.CUSTOM.default
282+
)
283+
payload = decode_payload(custom_node.kwargs["implementation_attrs"])
284+
285+
assert payload["input_0_type"] == "Tensor"
286+
assert payload["output_0_type"] == "Tensor"
287+
assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM"
288+
assert payload["workgroup_sizes"] == [2, 3, 1]

‎backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,16 +158,21 @@ def _encode_payload(
158158
padding_mode: int,
159159
align_corners: bool,
160160
input_tensor: torch.fx.Node,
161+
output_tensor: torch.fx.Node,
161162
output_dtype: torch.dtype | None = None,
162163
) -> list[int]:
163164
input_val = input_tensor.meta.get("val")
164165
if input_val is None:
165166
raise RuntimeError("grid_sampler_2d input is missing tensor metadata")
167+
output_val = output_tensor.meta.get("val")
168+
if output_val is None:
169+
raise RuntimeError("grid_sampler_2d output is missing tensor metadata")
166170
payload = build_grid_sampler_2d_payload(
167171
interpolation_mode=interpolation_mode,
168172
padding_mode=padding_mode,
169173
align_corners=align_corners,
170174
input_shape=tuple(input_val.shape),
175+
output_shape=tuple(output_val.shape),
171176
input_dtype=input_val.dtype,
172177
output_dtype=output_dtype,
173178
)
@@ -258,6 +263,7 @@ def call(self, graph_module):
258263
padding_mode=padding_mode,
259264
align_corners=align_corners,
260265
input_tensor=custom_input,
266+
output_tensor=node,
261267
output_dtype=output_dtype,
262268
)
263269
nhwc_input = create_node(

0 commit comments

Comments
 (0)