Skip to content

Commit 888c94f

Browse files
committed
Arm backend: Remove special annotation for uint8
When using uint8 input, don't do special annotation for operators. The computation for uint8 and int8 inputs is the same, just for uint8, we need to RESCALE the input to the int8 range. Special annotation for uint8 can result in operators left over in FP32(decomposed by ExecuTorch). Signed-off-by: George Gekov <george.gekov@arm.com> Change-Id: I5fdd06b9c3d72c80fe6937d821fb532e20ff7103
1 parent 9520b05 commit 888c94f

2 files changed

Lines changed: 58 additions & 23 deletions

File tree

‎backends/arm/quantizer/arm_quantizer_utils.py‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -511,13 +511,6 @@ class SharedQspecQuantizer(Quantizer, QuantizerReporterUser):
511511
torch.ops.higher_order.while_loop,
512512
torch.ops.higher_order.cond,
513513
]
514-
_UINT8_IO_BRIDGE_OPS: set[Callable[..., object]] = {
515-
torch.ops.aten.cat.default,
516-
torch.ops.aten.concatenate.default,
517-
torch.ops.aten.stack.default,
518-
torch.ops.aten.pixel_shuffle.default,
519-
torch.ops.aten.slice.Tensor,
520-
}
521514

522515
def __init__(self, targets: Optional[list[Callable[..., object]]] = None) -> None:
523516
super().__init__()
@@ -747,15 +740,6 @@ def _annotate_shared_cluster(self, root_node: Node) -> None:
747740
node_order = {node: index for index, node in enumerate(root_node.graph.nodes)}
748741
ordered_nodes = sorted(shared_nodes, key=lambda node: node_order.get(node, 0))
749742

750-
if touches_uint8_quantized_io and any(
751-
node.target in self._UINT8_IO_BRIDGE_OPS for node in shared_nodes
752-
):
753-
self.report_reject(
754-
ordered_nodes,
755-
"Shared-qspec bridge cluster touches uint8 model IO.",
756-
)
757-
return
758-
759743
if self._annotate_while_with_additional_inputs(root_node, adjacent_qspecs):
760744
return
761745

‎backends/arm/test/quantizer/test_uint8_io_quantization.py‎

Lines changed: 58 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@
1111
TOSAQuantizer,
1212
)
1313
from executorch.backends.arm.test import common
14-
from executorch.backends.arm.test.tester.test_pipeline import QuantizationPipeline
14+
from executorch.backends.arm.test.tester.test_pipeline import (
15+
QuantizationPipeline,
16+
TosaPipelineINT,
17+
)
1518
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e
1619
from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
1720

@@ -123,9 +126,13 @@ def test_io_boundary_shared_cluster_is_quantized():
123126
), "clone node has no output_qspec — IO-boundary cluster stayed in float"
124127

125128

126-
def test_cat_does_not_bridge_shared_qspec_clusters():
127-
"""Regression: cat must not merge image IO and high-range activations into
128-
one fallback shared-qspec observer clique.
129+
def test_cat_annotation_uint8_input():
130+
"""Depending on the location of the cat, we use different observers.
131+
132+
If we have uint8 inputs followed immediately by a cat, we observer each
133+
uint8 input independently. If we have a cat further-on in the graph, outside
134+
of the uint8 boundary, we annotate the cat with a SharedObserver as usual.
135+
129136
"""
130137
model = CatWithHighRangeBranch().eval()
131138
test_data = (torch.rand(1, 3, 8, 8), torch.rand(1, 3, 8, 8))
@@ -144,11 +151,25 @@ def test_cat_does_not_bridge_shared_qspec_clusters():
144151
img1_observer = next(iter(graph_nodes["img1"].users))
145152
clone_observer = next(iter(graph_nodes["clone"].users))
146153

147-
assert _get_observer_scale(prepared, img0_observer.target) < 0.01
148-
assert _get_observer_scale(prepared, img1_observer.target) < 0.01
149-
assert Q_ANNOTATION_KEY not in graph_nodes["cat_1"].meta
154+
img0_scale = _get_observer_scale(prepared, img0_observer.target)
155+
img1_scale = _get_observer_scale(prepared, img1_observer.target)
156+
assert img0_scale < 0.01
157+
assert img1_scale < 0.01
158+
# For uint8 input, the two inputs are observerd independently, hence need to
159+
# have different scales on the two input images.
160+
assert img0_scale != img1_scale
150161
assert _get_observer_scale(prepared, clone_observer.target) > 0.05
151162

163+
second_cat = graph_nodes["cat_1"]
164+
img_observer, high_range_observer = second_cat.args[0]
165+
166+
img_scale = _get_observer_scale(prepared, img_observer.target)
167+
high_range_scale = _get_observer_scale(prepared, high_range_observer.target)
168+
# For the second cat, we need to have the same scale on the inputs being concatenated
169+
# because everything outside of the IO boundary is annotated as per the
170+
# backends/arm/quantizer/quantization_annotator.py
171+
assert img_scale == high_range_scale
172+
152173

153174
def test_internal_cat_still_shares_qspec_with_uint8_io():
154175
"""Regression: preserved uint8 model IO must not disable shared-qspec
@@ -175,3 +196,33 @@ def test_internal_cat_still_shares_qspec_with_uint8_io():
175196

176197
assert Q_ANNOTATION_KEY in cat_node.meta
177198
assert cat_node.meta[Q_ANNOTATION_KEY].output_qspec is not None
199+
200+
201+
class SliceModel(torch.nn.Module):
202+
def __init__(self):
203+
super().__init__()
204+
self.conv = torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)
205+
206+
def forward(self, x: torch.Tensor) -> torch.Tensor:
207+
x = x[:, :3, :, :]
208+
return self.conv(x)
209+
210+
211+
def test_uint8_slice_tosa() -> None:
212+
"""Make sure that the SLICE operator is delegated after
213+
to_edge_transform_and_lower.
214+
"""
215+
test_data = (torch.rand(1, 6, 768, 384),)
216+
model = SliceModel().eval()
217+
pipeline = TosaPipelineINT(
218+
model,
219+
test_data,
220+
aten_op="torch.ops.aten.conv2d.default",
221+
run_on_tosa_ref_model=False,
222+
)
223+
pipeline.quantizer.set_io(get_uint8_io_quantization_config())
224+
pipeline.change_args(
225+
"check_not.exir",
226+
["executorch_exir_dialects_edge__ops_aten_"],
227+
)
228+
pipeline.run()

0 commit comments

Comments
 (0)