Fix constant_pad_nd->cat lowering dtype for quantized graphs

Differential Revision: D107545428

Pull Request resolved: pytorch#20039

Author: mcremon-meta

backends/cadence/aot/replace_ops.py

@@ -632,6 +632,7 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         value = 0 if len(node.args) == 2 else node.args[2]
 
         arg_shape = input_node.meta["val"].shape
+        dtype = input_node.meta["val"].dtype
 
         # Convert orig_padding to a list for manipulation
         # pyre-ignore[6]: Argument type
@@ -663,7 +664,7 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
                         left_padding_shape,
                         value,
                     ),
-                    kwargs={"dtype": torch.float32},
+                    kwargs={"dtype": dtype},
                 )
                 left_padding_node.meta = node.meta
             cat_tensors.append(left_padding_node)
@@ -683,7 +684,7 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
                         right_padding_shape,
                         value,
                     ),
-                    kwargs={"dtype": torch.float32},
+                    kwargs={"dtype": dtype},
                 )
                 right_padding_node.meta = node.meta
             cat_tensors.append(right_padding_node)

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -839,6 +839,30 @@ def test_replace_pad_with_cat(self, shape: Tuple[int], padding: Tuple[int]) -> N
             0,
         )
 
+    @torch.no_grad()
+    def test_replace_pad_with_cat_preserves_dtype(self) -> None:
+        # The padding constant tensors must match the input dtype, otherwise the
+        # resulting cat mixes dtypes and fails edge dialect dtype verification
+        # (e.g. for quantized int8 graphs).
+        x = torch.randint(-128, 127, (1, 2, 3), dtype=torch.int8)
+        original_gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.constant_pad_nd.default,
+            args=(x, [1, 1]),
+        )
+
+        p = ReplacePadWithCatPass()
+        result = cast(PassResult, p(original_gm))
+        self.assertTrue(result.modified)
+        graph_after_passes = result.graph_module
+
+        full_nodes = graph_after_passes.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.full.default
+        )
+        self.assertEqual(len(full_nodes), 2)
+        for full_node in full_nodes:
+            self.assertEqual(full_node.kwargs["dtype"], torch.int8)
+
     @torch.no_grad()
     def test_replace_repeat_with_cat(self) -> None:
         x = torch.randn([3, 5])