1111 TOSAQuantizer ,
1212)
1313from 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+ )
1518from torchao .quantization .pt2e .quantize_pt2e import prepare_pt2e
1619from 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
153174def 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