Update PostponeDequantizeOpBelowUseChainPass cost model (#22813)

backends/cadence/aot/reorder_ops.py

@@ -71,6 +71,12 @@
     exir_ops.edge.aten.select_copy,
 }
 
+supported_quantize_ops_overloadpkt = {
+    exir_ops.edge.quantized_decomposed.quantize_per_tensor,
+    exir_ops.edge.quantized_decomposed.quantize_per_channel,
+    exir_ops.edge.cadence.quantize_per_tensor,
+}
+
 
 class AdvanceQuantizeOpAboveDefInBranchPass(ExportPass):
     """
@@ -520,46 +526,43 @@ def postponing_feasible(self, dequant_node: torch.fx.Node):
         ):
             return False
 
-        dequant_shape = get_shape(self.graph_module, dequant_node)
-        slice_shapes = [
-            shape
-            for user in users
-            if (shape := get_shape(self.graph_module, user))
-            and (
-                # skip slices that are the size of the sliced tensor itself.
-                # They should technically get removed in the later passes as nop.
-                shape is None
-                or dequant_shape is None
-                or prod(list(shape)) != prod(list(dequant_shape))
+        quantized_branches = []
+        for user in users:
+            slice_users = list(user.users)
+            quantized_branches.append(
+                bool(slice_users)
+                and all(
+                    slice_user.op == "call_function"
+                    and get_overload_packet(slice_user.target)
+                    in supported_quantize_ops_overloadpkt
+                    for slice_user in slice_users
+                )
             )
-        ]
 
-        if dequant_shape is not None and all(
-            shape is not None for shape in slice_shapes
-        ):
-            dequant_bytes = num_bytes_from_shape_and_dtype(dequant_shape, torch.float32)
-            slice_bytes = sum(
-                [
-                    num_bytes_from_shape_and_dtype(shape, torch.float32)
-                    for shape in slice_shapes
-                ]
+        # Preserve the existing fallback for forks whose branches all requantize.
+        if all(quantized_branches):
+            return True
+
+        dequant_shape = get_shape(self.graph_module, dequant_node)
+        if dequant_shape is None:
+            return False
+
+        dequant_bytes = num_bytes_from_shape_and_dtype(dequant_shape, torch.float32)
+        surviving_dequant_bytes = 0
+        for user, quantized_branch in zip(users, quantized_branches):
+            if quantized_branch:
+                continue
+            slice_shape = get_shape(self.graph_module, user)
+            if slice_shape is None:
+                return False
+            # Nop slices are removed later and do not add another materialized copy.
+            if prod(list(slice_shape)) == prod(list(dequant_shape)):
+                continue
+            surviving_dequant_bytes += num_bytes_from_shape_and_dtype(
+                slice_shape, torch.float32
             )
-            if slice_bytes <= dequant_bytes:
-                return True
-
-        # If the users of each slice op is quantize op, then we can postpone
-        # dequantize, and convert slice -> dequantize -> quantize to
-        # slice -> requantize.
-        users = [x for y in users for x in y.users if x.op != "output"]
-        return all(
-            get_overload_packet(x.target)
-            in {
-                exir_ops.edge.quantized_decomposed.quantize_per_tensor,
-                exir_ops.edge.quantized_decomposed.quantize_per_channel,
-                exir_ops.edge.cadence.quantize_per_tensor,
-            }
-            for x in users
-        )
+
+        return surviving_dequant_bytes <= dequant_bytes
 
     def postpone_dequantize_op(self, graph_module: torch.fx.GraphModule) -> bool:
         # Different supported dequant ops have their own default variants
@@ -598,7 +601,8 @@ def postpone_dequantize_op(self, graph_module: torch.fx.GraphModule) -> bool:
             graph.erase_node(node)
             modified = True
 
-        graph_module.recompile()
+        if modified:
+            graph_module.recompile()
         return modified
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:

backends/cadence/aot/tests/test_reorder_ops_passes.py

@@ -608,6 +608,132 @@ def test_postpone_dequantize_branched(self) -> None:
             ),
         )
 
+    def _build_mixed_slice_quant_graph(
+        self, third_branch_quantized: bool
+    ) -> tuple[torch.fx.GraphModule, tuple[torch.Tensor, ...]]:
+        builder = GraphBuilder()
+        x_data = torch.randint(0, 255, [12, 4], dtype=torch.uint8)
+        weights_data = torch.randn([4, 4], dtype=torch.float32)
+
+        x = builder.placeholder("x", x_data)
+        weights = builder.placeholder("weights", weights_data)
+        dequant = builder.call_operator(
+            op=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x, 0.1, 10, 0, 255, torch.uint8),
+        )
+
+        outputs = []
+        branch_specs = [
+            (0, 8, (0.2, 5)),
+            (2, 10, None),
+            (4, 12, (0.3, 7) if third_branch_quantized else None),
+        ]
+        for start, end, target_qparams in branch_specs:
+            slice_node = builder.call_operator(
+                op=exir_ops.edge.aten.slice_copy.Tensor,
+                args=(dequant, 0, start, end),
+            )
+            if target_qparams is not None:
+                scale, zero_point = target_qparams
+                output = builder.call_operator(
+                    op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+                    args=(slice_node, scale, zero_point, 0, 255, torch.uint8),
+                )
+            else:
+                output = builder.call_operator(
+                    op=exir_ops.edge.aten.mm.default,
+                    args=(slice_node, weights),
+                )
+            outputs.append(output)
+
+        builder.output(outputs)
+        return builder.get_graph_module(), (x_data, weights_data)
+
+    def test_postpone_dequantize_mixed_slice_quant_branches(self) -> None:
+        original_graph, inputs = self._build_mixed_slice_quant_graph(True)
+        result = transform_and_check_numerics(
+            original_graph,
+            inputs,
+            PostponeDequantizeOpBelowUseChainPass(),
+        )
+        self.assertTrue(result.modified)
+
+        converted_graph = result.graph_module
+        graph_nodes = list(converted_graph.graph.nodes)
+        slice_nodes = [
+            node
+            for node in graph_nodes
+            if node.op == "call_function"
+            and node.target == exir_ops.edge.aten.slice_copy.Tensor
+        ]
+        dequant_nodes = [
+            node
+            for node in graph_nodes
+            if node.op == "call_function"
+            and node.target
+            == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
+        ]
+        self.assertEqual(len(slice_nodes), 3)
+        self.assertEqual(len(dequant_nodes), 3)
+        self.assertTrue(
+            all(node.meta["val"].dtype == torch.uint8 for node in slice_nodes)
+        )
+        self.assertTrue(
+            all(tuple(node.meta["val"].shape) == (8, 4) for node in dequant_nodes)
+        )
+        self.assertTrue(
+            all(node.meta["val"].dtype == torch.float32 for node in dequant_nodes)
+        )
+        for slice_node in slice_nodes:
+            source = slice_node.args[0]
+            self.assertIsInstance(source, torch.fx.Node)
+            self.assertEqual(cast(torch.fx.Node, source).meta["val"].dtype, torch.uint8)
+            downstream_dequants = [
+                user for user in slice_node.users if user in dequant_nodes
+            ]
+            self.assertEqual(len(downstream_dequants), 1)
+            self.assertLess(
+                graph_nodes.index(slice_node),
+                graph_nodes.index(downstream_dequants[0]),
+            )
+
+        fused_result = cast(
+            PassResult, FuseQuantDequantToRequantizePass()(converted_graph)
+        )
+        self.assertTrue(fused_result.modified)
+        self.assertEqual(
+            count_node(
+                fused_result.graph_module,
+                exir_ops.edge.cadence.requantize.per_tensor,
+            ),
+            2,
+        )
+        self.assertEqual(
+            count_node(
+                fused_result.graph_module,
+                exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            ),
+            1,
+        )
+
+    def test_postpone_dequantize_rejects_expensive_mixed_float_branches(
+        self,
+    ) -> None:
+        original_graph, inputs = self._build_mixed_slice_quant_graph(False)
+        result = transform_and_check_numerics(
+            original_graph,
+            inputs,
+            PostponeDequantizeOpBelowUseChainPass(),
+        )
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            count_node(
+                result.graph_module,
+                exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            ),
+            1,
+        )
+
     # 4d -> permute -> 4d -> view -> 3d
     def test_permute3_view4_chains(self) -> None:
         builder = GraphBuilder()