Skip to content

Commit 8fbc42c

Browse files
ArmRyanArmRyan
authored
Arm backend: Unsqueeze rank 0 tensor at vgf runtime (pytorch#14856)
Rank 0 tensors are not supported in SPV_ARM_tensor. We need to symbolically unsqueeze scalar IOs at runtime. * Remove xfails related to MLETORCH-1410 Signed-off-by: Ryan O'Shea <[email protected]>
1 parent 2eb8994 commit 8fbc42c

File tree

9 files changed

+22
-42
lines changed

9 files changed

+22
-42
lines changed

‎backends/arm/runtime/VGFSetup.cpp‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,13 @@ namespace vgf {
2424
/* static function to map format to byte count */
2525
static uint32_t get_format_size(VkFormat format);
2626

27+
// SPV_ARM_tensor does not support rank-0 representations according to the spec.
28+
// Use an unsqueezed dimension when the resource table contains an empty
29+
// shape. Tensors are output as rank 0 when copied back from the vgf backend.
30+
namespace {
31+
constexpr int64_t kScalarSentinelDimension = 1;
32+
}
33+
2734
// Debug function to inspect memory properties
2835
static string memory_flags_to_string(VkMemoryPropertyFlags flags) {
2936
if (flags == 0)
@@ -264,7 +271,11 @@ static void debug_print_resources(
264271
the_shape.size(),
265272
the_stride.size());
266273
for (int j = 0; j < the_shape.size(); j++) {
267-
ET_LOG(Info, " %d: dim %ld", j, the_shape[j]);
274+
ET_LOG(
275+
Info,
276+
" %d: dim %lld",
277+
j,
278+
static_cast<long long>(the_shape[j]));
268279
}
269280
// Allocate a tensor with bound memory
270281
break;
@@ -387,6 +398,7 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
387398
// Get tensor shape and strides
388399
auto shape = resource_decoder->getTensorShape(i);
389400
auto stride = resource_decoder->getTensorStride(i);
401+
const auto shape_size = shape.size();
390402

391403
switch (resource_decoder->getCategory(i)) {
392404
case vgflib::ResourceCategory::INPUT:
@@ -409,9 +421,9 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
409421
result = allocate_tensor(
410422
vk_physical,
411423
vk_device,
412-
vgflib::ToVkFormat(resource_decoder->getVkFormat(i)),
413-
static_cast<uint32_t>(shape.size()),
414-
shape.begin(),
424+
resource_format,
425+
shape_size == 0 ? 1 : static_cast<uint32_t>(shape_size),
426+
shape_size == 0 ? &kScalarSentinelDimension : shape.begin(),
415427
static_cast<uint32_t>(stride.size()),
416428
stride.begin(),
417429
&tensor_description,
@@ -422,8 +434,7 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
422434
ET_LOG(Error, "Failed to allocate tensor for VGF resource %d", i);
423435
return false;
424436
}
425-
size_t e_size = get_format_size(
426-
vgflib::ToVkFormat(resource_decoder->getVkFormat(i)));
437+
size_t e_size = get_format_size(resource_format);
427438
if (0 == e_size) {
428439
ET_LOG(Error, "failed to get element size of VkFormat");
429440
return false;
@@ -449,9 +460,11 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
449460
.sType = VK_STRUCTURE_TYPE_TENSOR_DESCRIPTION_ARM,
450461
.pNext = nullptr,
451462
.tiling = VK_TENSOR_TILING_LINEAR_ARM,
452-
.format = vgflib::ToVkFormat(resource_decoder->getVkFormat(i)),
453-
.dimensionCount = static_cast<uint32_t>(shape.size()),
454-
.pDimensions = shape.begin(),
463+
.format = resource_format,
464+
.dimensionCount =
465+
shape_size == 0 ? 1 : static_cast<uint32_t>(shape_size),
466+
.pDimensions =
467+
shape_size == 0 ? &kScalarSentinelDimension : shape.begin(),
455468
// Note: stride_data of 0's causes size==0, null means stride==size
456469
.pStrides = (0 == stride.size() ? nullptr : stride.begin()),
457470
.usage = VK_TENSOR_USAGE_DATA_GRAPH_BIT_ARM,

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,6 @@ def test_addmm_u85_INT(test_data: Tuple):
167167

168168
@common.parametrize("test_data", test_data_suite)
169169
@common.SkipIfNoModelConverter
170-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
171170
def test_addmm_vgf_FP(test_data: input_t1):
172171
pipeline = VgfPipeline[input_t1](
173172
Addmm(),
@@ -181,7 +180,6 @@ def test_addmm_vgf_FP(test_data: input_t1):
181180

182181
@common.parametrize("test_data", test_data_suite)
183182
@common.SkipIfNoModelConverter
184-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
185183
def test_addmm_vgf_INT(test_data: input_t1):
186184
pipeline = VgfPipeline[input_t1](
187185
Addmm(),

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,6 @@ def test_max_dim_tosa_FP_not_delegated():
139139

140140
@common.parametrize("test_data", Amax.test_data)
141141
@common.SkipIfNoModelConverter
142-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
143142
def test_amax_vgf_FP(test_data: Amax.input_t):
144143
data, dim, keep_dims = test_data()
145144
module = Amax(dim, keep_dims)
@@ -154,7 +153,6 @@ def test_amax_vgf_FP(test_data: Amax.input_t):
154153

155154
@common.parametrize("test_data", Amax.test_data)
156155
@common.SkipIfNoModelConverter
157-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
158156
def test_amax_vgf_INT(test_data: Amax.input_t):
159157
data, dim, keep_dims = test_data()
160158
module = Amax(dim, keep_dims)
@@ -169,7 +167,6 @@ def test_amax_vgf_INT(test_data: Amax.input_t):
169167

170168
@common.parametrize("test_data", Max.test_data)
171169
@common.SkipIfNoModelConverter
172-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
173170
def test_max_dim_vgf_FP_to_amax(test_data: Max.input_t):
174171
data, dim = test_data()
175172
pipeline = VgfPipeline[Max.input_t](
@@ -183,7 +180,6 @@ def test_max_dim_vgf_FP_to_amax(test_data: Max.input_t):
183180

184181
@common.parametrize("test_data", Max.test_data)
185182
@common.SkipIfNoModelConverter
186-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
187183
def test_max_dim_vgf_INT_to_amax(test_data: Max.input_t):
188184
data, dim = test_data()
189185
pipeline = VgfPipeline[Max.input_t](

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,6 @@ def test_min_dim_tosa_FP_not_delegated():
155155

156156
@common.parametrize("test_data", Amin.test_data)
157157
@common.SkipIfNoModelConverter
158-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
159158
def test_amin_vgf_FP(test_data: Amin.input_t):
160159
data, dim, keep_dims = test_data()
161160
pipeline = VgfPipeline[Amin.input_t](
@@ -166,7 +165,6 @@ def test_amin_vgf_FP(test_data: Amin.input_t):
166165

167166
@common.parametrize("test_data", Amin.test_data)
168167
@common.SkipIfNoModelConverter
169-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
170168
def test_amin_vgf_INT(test_data: Amin.input_t):
171169
data, dim, keep_dims = test_data()
172170
pipeline = VgfPipeline[Amin.input_t](
@@ -180,7 +178,6 @@ def test_amin_vgf_INT(test_data: Amin.input_t):
180178

181179
@common.parametrize("test_data", Min.test_data)
182180
@common.SkipIfNoModelConverter
183-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
184181
def test_min_dim_vgf_FP_to_amin(test_data: Min.input_t):
185182
data, dim = test_data()
186183
pipeline = VgfPipeline[Min.input_t](
@@ -194,7 +191,6 @@ def test_min_dim_vgf_FP_to_amin(test_data: Min.input_t):
194191

195192
@common.parametrize("test_data", Min.test_data)
196193
@common.SkipIfNoModelConverter
197-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
198194
def test_min_dim_vgf_INT_to_amin(test_data: Min.input_t):
199195
data, dim = test_data()
200196
pipeline = VgfPipeline[Min.input_t](

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

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

77
from typing import List, Tuple
88

9-
import pytest
109
import torch
1110
from executorch.backends.arm.test import common
1211
from executorch.backends.arm.test.tester.test_pipeline import (
@@ -189,7 +188,6 @@ def test_any_u85_INT(test_data: input_t1):
189188

190189
@common.parametrize("test_data", test_data)
191190
@common.SkipIfNoModelConverter
192-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
193191
def test_any_vgf_FP(test_data: input_t1):
194192
op, data_fn = test_data()
195193
pipeline = VgfPipeline[input_t1](
@@ -204,7 +202,6 @@ def test_any_vgf_FP(test_data: input_t1):
204202

205203
@common.parametrize("test_data", test_data)
206204
@common.SkipIfNoModelConverter
207-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
208205
def test_any_vgf_INT(test_data: input_t1):
209206
op, data_fn = test_data()
210207
pipeline = VgfPipeline[input_t1](

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

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
#
55
# This source code is licensed under the BSD-style license found in the
66
# LICENSE file in the root directory of this source tree.
7-
import pytest
87
import torch
98
from executorch.backends.arm.test import common
109
from executorch.backends.arm.test.tester.test_pipeline import (
@@ -84,7 +83,6 @@ def test_adaptive_avg_pool2d_u85_INT(test_data):
8483

8584
@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
8685
@common.SkipIfNoModelConverter
87-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
8886
def test_adaptive_avg_pool2d_vgf_FP(test_data):
8987
pipeline = VgfPipeline[input_t](
9088
AdaptiveAveragePool2d(),
@@ -98,7 +96,6 @@ def test_adaptive_avg_pool2d_vgf_FP(test_data):
9896

9997
@common.parametrize("test_data", AdaptiveAveragePool2d.test_data_suite)
10098
@common.SkipIfNoModelConverter
101-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
10299
def test_adaptive_avg_pool2d_vgf_INT(test_data):
103100
pipeline = VgfPipeline[input_t](
104101
AdaptiveAveragePool2d(),
@@ -331,7 +328,6 @@ def test_mean_dim_u85_INT(test_data):
331328

332329
@common.parametrize("test_data", MeanDim.test_data_suite)
333330
@common.SkipIfNoModelConverter
334-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
335331
def test_mean_dim_vgf_FP(test_data):
336332
test_data_val, dim, keep_dim = test_data()
337333
pipeline = VgfPipeline[input_t](
@@ -346,7 +342,6 @@ def test_mean_dim_vgf_FP(test_data):
346342

347343
@common.parametrize("test_data", MeanDim.test_data_suite)
348344
@common.SkipIfNoModelConverter
349-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
350345
def test_mean_dim_vgf_INT(test_data):
351346
test_data_val, dim, keep_dim = test_data()
352347
pipeline = VgfPipeline[input_t](

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
#
33
# This source code is licensed under the BSD-style license found in the
44
# LICENSE file in the root directory of this source tree.
5-
import pytest
65
import torch
76
from executorch.backends.arm.test import common
87

@@ -102,7 +101,6 @@ def test_scalar_tensor_u85_INT(test_data):
102101

103102
@common.parametrize("test_data", float_test_data_suite)
104103
@common.SkipIfNoModelConverter
105-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
106104
def test_scalar_tensor_vgf_FP(test_data):
107105
scalar, dtype, data = test_data()
108106
pipeline = VgfPipeline(
@@ -116,7 +114,6 @@ def test_scalar_tensor_vgf_FP(test_data):
116114

117115
@common.parametrize("test_data", int_test_data_suite)
118116
@common.SkipIfNoModelConverter
119-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
120117
def test_scalar_tensor_vgf_INT(test_data):
121118
scalar, dtype, data = test_data()
122119
pipeline = VgfPipeline(

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

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

88
from typing import Tuple
99

10-
import pytest
1110
import torch
1211

1312
from executorch.backends.arm.test import common
@@ -170,7 +169,6 @@ def test_select_int_u85_INT(test_data: Tuple):
170169

171170
@common.parametrize("test_data", test_data_suite)
172171
@common.SkipIfNoModelConverter
173-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
174172
def test_select_int_vgf_FP_copy(test_data: Tuple):
175173
pipeline = VgfPipeline[input_t1](
176174
SelectCopy(), test_data(), aten_op_copy, [], tosa_version="TOSA-1.0+FP"
@@ -180,7 +178,6 @@ def test_select_int_vgf_FP_copy(test_data: Tuple):
180178

181179
@common.parametrize("test_data", test_data_suite)
182180
@common.SkipIfNoModelConverter
183-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
184181
def test_select_int_vgf_FP(test_data: Tuple):
185182
pipeline = VgfPipeline[input_t1](
186183
SelectInt(), test_data(), aten_op_int, [], tosa_version="TOSA-1.0+FP"
@@ -190,7 +187,6 @@ def test_select_int_vgf_FP(test_data: Tuple):
190187

191188
@common.parametrize("test_data", test_data_suite)
192189
@common.SkipIfNoModelConverter
193-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
194190
def test_select_int_vgf_INT_copy(test_data: Tuple):
195191
pipeline = VgfPipeline[input_t1](
196192
SelectCopy(),
@@ -204,7 +200,6 @@ def test_select_int_vgf_INT_copy(test_data: Tuple):
204200

205201
@common.parametrize("test_data", test_data_suite)
206202
@common.SkipIfNoModelConverter
207-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
208203
def test_select_int_vgf_INT(test_data: Tuple):
209204
pipeline = VgfPipeline[input_t1](
210205
SelectInt(),

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

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

77
from typing import Tuple
88

9-
import pytest
109
import torch
1110

1211
from executorch.backends.arm.test import common
@@ -214,7 +213,6 @@ def test_var_dim_u85_INT_no_dim(test_data: Tuple):
214213

215214
@common.parametrize("test_data", Var.test_parameters)
216215
@common.SkipIfNoModelConverter
217-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
218216
def test_var_dim_vgf_FP_no_dim(test_data: Tuple):
219217
data, keepdim, correction = test_data()
220218
pipeline = VgfPipeline[input_t1](
@@ -225,7 +223,6 @@ def test_var_dim_vgf_FP_no_dim(test_data: Tuple):
225223

226224
@common.parametrize("test_data", Var.test_parameters)
227225
@common.SkipIfNoModelConverter
228-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
229226
def test_var_dim_vgf_INT_no_dim(test_data: Tuple):
230227
data, keepdim, correction = test_data()
231228
pipeline = VgfPipeline[input_t1](
@@ -296,7 +293,6 @@ def test_var_dim_u85_INT(test_data: Tuple):
296293

297294
@common.parametrize("test_data", VarDim.test_parameters)
298295
@common.SkipIfNoModelConverter
299-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
300296
def test_var_dim_vgf_FP(test_data: Tuple):
301297
data, dim, keepdim, unbiased = test_data()
302298
pipeline = VgfPipeline[input_t1](
@@ -307,7 +303,6 @@ def test_var_dim_vgf_FP(test_data: Tuple):
307303

308304
@common.parametrize("test_data", VarDim.test_parameters)
309305
@common.SkipIfNoModelConverter
310-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
311306
def test_var_dim_vgf_INT(test_data: Tuple):
312307
data, dim, keepdim, unbiased = test_data()
313308
pipeline = VgfPipeline[input_t1](
@@ -377,7 +372,6 @@ def test_var_dim_u85_INT_correction(test_data: Tuple):
377372

378373
@common.parametrize("test_data", VarCorrection.test_parameters)
379374
@common.SkipIfNoModelConverter
380-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
381375
def test_var_dim_vgf_FP_correction(test_data: Tuple):
382376
data, dim, keepdim, corr = test_data()
383377
pipeline = VgfPipeline[input_t1](
@@ -388,7 +382,6 @@ def test_var_dim_vgf_FP_correction(test_data: Tuple):
388382

389383
@common.parametrize("test_data", VarCorrection.test_parameters)
390384
@common.SkipIfNoModelConverter
391-
@pytest.mark.xfail(reason="MLETORCH-1410: Tensor dimension count not supported: 0")
392385
def test_var_dim_vgf_INT_correction(test_data: Tuple):
393386
data, dim, keepdim, corr = test_data()
394387
pipeline = VgfPipeline[input_t1](

0 commit comments

Comments
 (0)