Duplicate cat if sandwiched between deq/quant (#19925)

Differential Revision: D107174424

Pull Request resolved: #19925

Author: DrJessop

backends/cadence/aot/reorder_ops.py

@@ -895,6 +895,77 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         return should_swap(parent, node) and do_swap(parent, node)
 
 
+_QUANT_OVERLOAD_PACKETS = {
+    exir_ops.edge.quantized_decomposed.quantize_per_tensor,
+    exir_ops.edge.cadence.quantize_per_tensor,
+}
+
+_DEQUANT_OVERLOAD_PACKETS = {
+    exir_ops.edge.quantized_decomposed.dequantize_per_tensor,
+    exir_ops.edge.cadence.dequantize_per_tensor,
+}
+
+
+@register_cadence_pass(CadencePassAttribute(opt_level=1))
+class SplitDequantizedCatPass(RemoveOrReplacePassInterface):
+    """Split a cat node so that quantize consumers get their own copy.
+
+    Fires when a cat has all floating-point inputs, at least one dequantize
+    input, and at least one quantize consumer.  Quant consumers are grouped
+    by matching qparams; each group receives a dedicated duplicate of the
+    cat node.  Non-quant consumers stay on the original cat, whose
+    semantics are unchanged.
+
+    A later pass (e.g. AdvanceQuantizeOpAboveDefChainPass extended for cat)
+    can then hoist each quant above its single-consumer cat copy without
+    affecting the non-quant paths.
+    """
+
+    @property
+    def targets(self) -> list[EdgeOpOverload]:
+        return [exir_ops.edge.aten.cat.default]
+
+    def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
+        cat_inputs = node.args[0]
+        if not isinstance(cat_inputs, (list, tuple)):
+            return False
+
+        has_dequant_input = False
+        for inp in cat_inputs:
+            assert isinstance(inp, torch.fx.Node)
+            val = inp.meta["val"]
+            if val is None or not val.is_floating_point():
+                return False
+            if get_overload_packet(inp.target) in _DEQUANT_OVERLOAD_PACKETS:
+                has_dequant_input = True
+
+        if not has_dequant_input:
+            return False
+
+        quant_groups: DefaultDict[Tuple, List[torch.fx.Node]] = defaultdict(list)
+        for user in list(node.users.keys()):
+            if get_overload_packet(user.target) in _QUANT_OVERLOAD_PACKETS:
+                quant_groups[user.args[1:]].append(user)
+
+        if not quant_groups:
+            return False
+
+        graph = node.graph
+        dim = get_arg(node, "dim", int)
+        for quant_consumers in quant_groups.values():
+            with graph.inserting_after(node):
+                dup_cat = graph.call_function(
+                    exir_ops.edge.aten.cat.default,
+                    args=(list(cat_inputs), dim),
+                )
+                dup_cat.meta = node.meta.copy()
+
+            for q_node in quant_consumers:
+                q_node.replace_input_with(node, dup_cat)
+
+        return True
+
+
 # The following class consolidates functions to reoder ops (i.e., either hoist
 # or sink some ops in the graph).
 class CadenceReorderOpsInGraph:

backends/cadence/aot/tests/test_reorder_ops_passes.py

@@ -28,6 +28,7 @@
     PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
     PropagateSlice,
     SinkOpsCloserToUsePass,
+    SplitDequantizedCatPass,
 )
 from executorch.backends.test.graph_builder import GraphBuilder
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -1024,3 +1025,246 @@ def test_no_swap_binary_same_shape(self) -> None:
         result = PropagateSlice().call(gm)
 
         self.assertFalse(result.modified)
+
+
+class TestSplitDequantizedCat(unittest.TestCase):
+    def test_no_dequant_input_noop(self) -> None:
+        """Cat with only float (non-dequant) inputs should not be split."""
+        builder = GraphBuilder()
+        a = builder.placeholder("a", torch.randn(2, 4))
+        b = builder.placeholder("b", torch.randn(2, 4))
+        cat = builder.call_operator(exir_ops.edge.aten.cat.default, args=([a, b], 0))
+        q = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.01, 0, -128, 127, torch.int8),
+        )
+        dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q, 0.01, 0, -128, 127, torch.int8),
+        )
+        builder.output([dq])
+        gm = builder.get_graph_module()
+
+        result = SplitDequantizedCatPass().call(gm)
+
+        self.assertFalse(result.modified)
+        self.assertEqual(count_node(gm, exir_ops.edge.aten.cat.default), 1)
+
+    def test_no_quant_output_noop(self) -> None:
+        """Cat with a dequant input but no quant consumer should not be split."""
+        builder = GraphBuilder()
+        x_int8 = builder.placeholder(
+            "x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
+        )
+        dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x_int8, 0.01, 0, -128, 127, torch.int8),
+        )
+        b = builder.placeholder("b", torch.randn(2, 4))
+        cat = builder.call_operator(exir_ops.edge.aten.cat.default, args=([dq, b], 0))
+        builder.output([cat])
+        gm = builder.get_graph_module()
+
+        result = SplitDequantizedCatPass().call(gm)
+
+        self.assertFalse(result.modified)
+        self.assertEqual(count_node(gm, exir_ops.edge.aten.cat.default), 1)
+
+    def test_one_dequant_input_one_quant_output(self) -> None:
+        """Cat with one dequant input and one quant consumer should be split."""
+        builder = GraphBuilder()
+        x_int8 = builder.placeholder(
+            "x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
+        )
+        dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x_int8, 0.01, 0, -128, 127, torch.int8),
+        )
+        b = builder.placeholder("b", torch.randn(2, 4))
+        cat = builder.call_operator(exir_ops.edge.aten.cat.default, args=([dq, b], 0))
+        sliced = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
+        )
+        q = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.02, -5, -128, 127, torch.int8),
+        )
+        q_dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q, 0.02, -5, -128, 127, torch.int8),
+        )
+        builder.output([sliced, q_dq])
+        gm = builder.get_graph_module()
+
+        result = SplitDequantizedCatPass().call(gm)
+
+        self.assertTrue(result.modified)
+        converted = result.graph_module
+        self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 2)
+
+        # The slice should still be on the original cat, which has no quant consumers.
+        slice_nodes = converted.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+        )
+        for node in slice_nodes:
+            cat_input = node.args[0]
+            self.assertEqual(cat_input.target, exir_ops.edge.aten.cat.default)
+            quant_users = [
+                u
+                for u in cat_input.users
+                if u.target
+                == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
+            ]
+            self.assertEqual(len(quant_users), 0)
+
+    def test_non_quant_consumers_stay_on_original_cat(self) -> None:
+        """All non-quant consumers should remain on the original cat."""
+        builder = GraphBuilder()
+        x_int8 = builder.placeholder(
+            "x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
+        )
+        dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x_int8, 0.01, 0, -128, 127, torch.int8),
+        )
+        b = builder.placeholder("b", torch.randn(2, 4))
+        cat = builder.call_operator(exir_ops.edge.aten.cat.default, args=([dq, b], 0))
+        sliced = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
+        )
+        abs_val = builder.call_operator(exir_ops.edge.aten.abs.default, args=(cat,))
+        q = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.02, -5, -128, 127, torch.int8),
+        )
+        q_dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q, 0.02, -5, -128, 127, torch.int8),
+        )
+        builder.output([sliced, abs_val, q_dq])
+        gm = builder.get_graph_module()
+
+        result = SplitDequantizedCatPass().call(gm)
+
+        self.assertTrue(result.modified)
+        converted = result.graph_module
+        self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 2)
+
+        # Both non-quant consumers (slice and abs) should use the same cat.
+        slice_nodes = converted.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+        )
+        abs_nodes = converted.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.abs.default
+        )
+        self.assertEqual(len(slice_nodes), 1)
+        self.assertEqual(len(abs_nodes), 1)
+        self.assertIs(slice_nodes[0].args[0], abs_nodes[0].args[0])
+
+        # That shared cat should have no quant consumers.
+        original_cat = slice_nodes[0].args[0]
+        quant_users = [
+            u
+            for u in original_cat.users
+            if u.target
+            == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
+        ]
+        self.assertEqual(len(quant_users), 0)
+
+    def test_two_quant_outputs_same_params_shared_cat(self) -> None:
+        """Two quant consumers with identical params should share one duplicate cat."""
+        builder = GraphBuilder()
+        x_int8 = builder.placeholder(
+            "x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
+        )
+        dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x_int8, 0.01, 0, -128, 127, torch.int8),
+        )
+        b = builder.placeholder("b", torch.randn(2, 4))
+        cat = builder.call_operator(exir_ops.edge.aten.cat.default, args=([dq, b], 0))
+        sliced = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
+        )
+        q1 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.02, -5, -128, 127, torch.int8),
+        )
+        q2 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.02, -5, -128, 127, torch.int8),
+        )
+        dq1 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q1, 0.02, -5, -128, 127, torch.int8),
+        )
+        dq2 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q2, 0.02, -5, -128, 127, torch.int8),
+        )
+        builder.output([sliced, dq1, dq2])
+        gm = builder.get_graph_module()
+
+        result = SplitDequantizedCatPass().call(gm)
+
+        self.assertTrue(result.modified)
+        converted = result.graph_module
+        # Original cat + one shared duplicate = 2 cats total
+        self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 2)
+
+        # Both quant nodes should share the same cat input (the duplicate).
+        quant_nodes = converted.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        )
+        quant_cat_inputs = {node.args[0] for node in quant_nodes}
+        self.assertEqual(len(quant_cat_inputs), 1)
+
+    def test_two_quant_outputs_different_params_separate_cats(self) -> None:
+        """Two quant consumers with different params should get separate duplicate cats."""
+        builder = GraphBuilder()
+        x_int8 = builder.placeholder(
+            "x_int8", torch.randint(-128, 127, (2, 4), dtype=torch.int8)
+        )
+        dq = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(x_int8, 0.01, 0, -128, 127, torch.int8),
+        )
+        b = builder.placeholder("b", torch.randn(2, 4))
+        cat = builder.call_operator(exir_ops.edge.aten.cat.default, args=([dq, b], 0))
+        sliced = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor, args=(cat, 0, 0, 2)
+        )
+        q1 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.02, -5, -128, 127, torch.int8),
+        )
+        q2 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            args=(cat, 0.03, 10, -128, 127, torch.int8),
+        )
+        dq1 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q1, 0.02, -5, -128, 127, torch.int8),
+        )
+        dq2 = builder.call_operator(
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            args=(q2, 0.03, 10, -128, 127, torch.int8),
+        )
+        builder.output([sliced, dq1, dq2])
+        gm = builder.get_graph_module()
+
+        result = SplitDequantizedCatPass().call(gm)
+
+        self.assertTrue(result.modified)
+        converted = result.graph_module
+        # Original cat + two separate duplicates = 3 cats total
+        self.assertEqual(count_node(converted, exir_ops.edge.aten.cat.default), 3)
+
+        # Each quant node should have a different cat input.
+        quant_nodes = converted.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        )
+        quant_cat_inputs = {node.args[0] for node in quant_nodes}
+        self.assertEqual(len(quant_cat_inputs), 2)