Skip to content

Commit a3bf556

Browse files
Treat aten.slice.Tensor as a shared-observer op in XNNPACKQuantizer (pytorch#21568)
### Summary `_is_share_obs_or_fq_op` lists `aten.slice_copy.Tensor` but not `aten.slice.Tensor`. The quantizer runs before `to_edge()`, so the op it actually sees is `aten.slice.Tensor`, the `_copy` entry never matches, and `propagate_annotation` does not give the slice its producer's `SharedQuantizationSpec`. The other data-movement ops in that list are registered under both names (`permute.default`/`permute_copy.default`, `squeeze.dim`/`squeeze_copy.dim`, `view.default`/`view_copy.default`); slice only had the `_copy` one. `slice` has no `.default` overload, which is why the pre-functionalization name is `slice.Tensor` and not `slice_copy.default`. Without the shared spec, the slice input and output end up with different quant params. XNNPACK's `static_slice` requires them to be equal, so the model partitions and lowers fine and then fails at runtime init: ``` [XNNCompiler.cpp:1421] Failed to create static slice node 377 with code: xnn_status_invalid_parameter [XNNPACKBackend.cpp:119] XNNCompiler::compileModel failed: 0x1 [method.cpp:110] Init failed for backend XnnpackBackend: 0x1 ``` I hit this quantizing a Gemma-family embedding model with `XNNPACKQuantizer` (per-channel W8A8 static), where all 96 slices from the RoPE `rotate_half` pattern were left without quant params. `_test_qs8_slice_copy` was already disabled for this ("T156004676 - slice isn't propagated"), and its comment names the same cause, so I re-enabled it here. ### Test plan ``` pytest backends/xnnpack/test/ops/test_slice_copy.py -k qs8 ``` The test runs end to end now: quantize, lower, serialize, run and compare outputs. Reverting the one-line change makes it fail with 2 `quantize_per_tensor` nodes instead of 5. Two assertions in it needed updating first. `check_node_count` takes op overloads rather than the `"aten::slice.Tensor"` strings the test was written with, so it reported `Expected 3 aten::slice.Tensor nodes but found 0` on a graph that did contain them. And the expected `quantize_per_tensor` count is 5, not 3, since the three slices are now quantized along with the input and the add output. `lintrunner -m upstream/main` is clean. cc @GregoryComer @digantdesai @cbilgin @JakeStevens
1 parent 52b176f commit a3bf556

2 files changed

Lines changed: 7 additions & 6 deletions

File tree

‎backends/xnnpack/quantizer/xnnpack_quantizer_utils.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1102,6 +1102,7 @@ def _is_share_obs_or_fq_op(op: Callable) -> bool:
11021102
torch.ops.aten.view_copy.default,
11031103
torch.ops.aten.view.default,
11041104
torch.ops.aten.reshape.default,
1105+
torch.ops.aten.slice.Tensor,
11051106
torch.ops.aten.slice_copy.Tensor,
11061107
torch.ops.aten.flatten.using_ints,
11071108
]

‎backends/xnnpack/test/ops/test_slice_copy.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -140,10 +140,9 @@ def forward(self, x):
140140
.check_not(["torch.ops.higher_order.executorch_call_delegate"])
141141
)
142142

143-
# Note: Slice ends up as slice_copy later in the process, but during quantization,
144-
# it's still slice, which isn't supported by the XNNPACK quantizer.
145-
@unittest.skip("T156004676 - slice isn't propagated")
146-
def _test_qs8_slice_copy(self):
143+
# Note: Slice ends up as slice_copy later in the process, but during quantization
144+
# it's still slice, so the quantizer shares observers on aten.slice.Tensor.
145+
def test_qs8_slice_copy(self):
147146
class SliceCopy(torch.nn.Module):
148147
def forward(self, x):
149148
y = x + x
@@ -157,8 +156,9 @@ def forward(self, x):
157156
.export()
158157
.check_node_count(
159158
{
160-
"aten::slice.Tensor": 3,
161-
"quantized_decomposed::quantize_per_tensor": 3,
159+
torch.ops.aten.slice.Tensor: 3,
160+
# 3 slices, plus the input and the add output
161+
torch.ops.quantized_decomposed.quantize_per_tensor.default: 5,
162162
}
163163
)
164164
.to_edge_transform_and_lower()

0 commit comments

Comments
 (0)