Materialize zero bias for conv replacement (pytorch#23354)

Summary:

`aten::convolution` declares `Tensor? bias`, so a conv built with `bias=False` lowers to a node whose bias argument is `None`. `cadence::conv1d/conv2d/conv3d` declare a non-optional `Tensor bias`, with no way to express absence. Replacing a biasless convolution therefore yields a node that strict Edge validation rejects for schema mismatch.

Materialize an explicit zero bias via `edge.aten.full.default`, sized to the convolution's output channels and typed from the weight, so the replacement is schema-valid. Adding zeros is numerically a no-op.

If the Cadence kernels gain optional-bias support, as discussed on the diff, this method can be deleted outright.

Reviewed By: DrJessop

Differential Revision: D110672307

Pull Request resolved: pytorch#23354

Author: pcwu0329

backends/cadence/aot/replace_ops.py

@@ -722,6 +722,35 @@ class ReplaceAtenConvolutionWithCadenceConvolutionPass(RemoveOrReplacePassInterf
     def targets(self) -> list[EdgeOpOverload]:
         return [exir_ops.edge.aten.convolution.default]
 
+    def materialize_zero_bias(
+        self,
+        node: torch.fx.Node,
+        weight: torch.fx.Node,
+    ) -> torch.fx.Node:
+        out_val = node.meta["val"]
+        weight_val = weight.meta["val"]
+        out_shape = out_val.shape
+        if len(out_shape) < 2:
+            raise ValueError(
+                f"Expected convolution output rank >= 2, got shape {out_shape}"
+            )
+        # edge.aten.full accepts SymInt[] sizes, so retain symbolic channels.
+        bias_size = out_shape[1]
+        bias_dtype = weight_val.dtype
+        with node.graph.inserting_before(node):
+            zero_bias = node.graph.call_function(
+                exir_ops.edge.aten.full.default,
+                args=([bias_size], 0.0),
+                kwargs={"dtype": bias_dtype},
+            )
+        # Populate per-field rather than aliasing node.meta: "val" has to
+        # describe this bias, and "stack_trace" is the convolution's provenance,
+        # so copying it would misattribute the synthesized constant.
+        zero_bias.meta["val"] = weight_val.new_zeros((bias_size,), dtype=bias_dtype)
+        if "debug_handle" in node.meta:
+            zero_bias.meta["debug_handle"] = node.meta["debug_handle"]
+        return zero_bias
+
     def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         # There must be 9 total args.
         if len(node.args) != 9:
@@ -770,6 +799,16 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         else:  # len(stride) == 3
             target = exir_ops.edge.cadence.conv3d.default
 
+        # Must reject before the bias is materialized below; bailing out after
+        # would strand the synthesized full node in the graph. Transposed
+        # convolutions accept non-zero output_padding, hence the guard.
+        if not transposed and not all(x == 0 for x in output_padding):
+            return False
+
+        if bias is None:
+            assert isinstance(weight, torch.fx.Node)
+            bias = self.materialize_zero_bias(node, weight)
+
         with node.graph.inserting_before(node):
             if transposed:
                 # Flip the height and width dimensions of weight, since we apply a
@@ -804,10 +843,6 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
                     False,
                 )
             else:
-                # Verify that output_padding is 0.
-                if not all(x == 0 for x in output_padding):
-                    return False
-
                 # Keep the original stride to maintain correct output dimensions
                 new_stride = stride
 

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -520,6 +520,8 @@ def test_replace_aten_conv_with_cadence_conv(
                 groups,
             ),
         )
+        if not bias_enabled:
+            convolution.node.meta["debug_handle"] = 1234
         builder.output([convolution])
         original_gm = builder.get_graph_module()
 
@@ -550,13 +552,171 @@ def test_replace_aten_conv_with_cadence_conv(
             count_node(graph_after_passes, exir_ops.edge.cadence.conv1d.default),
             1,
         )
+        for node in graph_after_passes.graph.nodes:
+            if node.target == exir_ops.edge.cadence.conv1d.default:
+                self.assertIsNotNone(node.args[2])
+                if not bias_enabled:
+                    bias_node = node.args[2]
+                    self.assertIsInstance(bias_node, torch.fx.Node)
+                    self.assertEqual(
+                        exir_ops.edge.aten.full.default,
+                        cast(torch.fx.Node, bias_node).target,
+                    )
+                    self.assertEqual(
+                        1234, cast(torch.fx.Node, bias_node).meta["debug_handle"]
+                    )
+                    self.assertEqual(
+                        (out_channels,),
+                        tuple(cast(torch.fx.Node, bias_node).meta["val"].shape),
+                    )
         self.assertEqual(
             count_node(
                 graph_after_passes, exir_ops.edge.cadence.transposed_convolution.default
             ),
             0,
         )
 
+    @expand([[2], [3]])
+    @torch.no_grad()
+    def test_replace_aten_multidimensional_conv_materializes_zero_bias(
+        self,
+        spatial_dims: int,
+    ) -> None:
+        in_channels = 4
+        out_channels = 6
+        builder = GraphBuilder()
+        x_tensor = torch.randn(
+            1,
+            in_channels,
+            *([8] * spatial_dims),
+            dtype=torch.float32,
+        )
+        weights_tensor = torch.randn(
+            out_channels,
+            in_channels,
+            *([3] * spatial_dims),
+            dtype=torch.float32,
+        )
+        x = builder.placeholder("x", x_tensor)
+        weights = builder.placeholder("weights", weights_tensor)
+        convolution = builder.call_operator(
+            op=exir_ops.edge.aten.convolution.default,
+            args=(
+                x,
+                weights,
+                None,
+                [1] * spatial_dims,
+                [0] * spatial_dims,
+                [1] * spatial_dims,
+                False,
+                [0] * spatial_dims,
+                1,
+            ),
+        )
+        convolution.node.meta["debug_handle"] = 1234
+        convolution.node.meta["stack_trace"] = "convolution-only provenance"
+        builder.output([convolution])
+        original_gm = builder.get_graph_module()
+        gm_before = copy.deepcopy(original_gm)
+
+        result = cast(
+            PassResult,
+            ReplaceAtenConvolutionWithCadenceConvolutionPass()(original_gm),
+        )
+        self.assertTrue(result.modified)
+        graph_after_passes = result.graph_module
+        validate(
+            gm_before,
+            graph_after_passes,
+            (x_tensor, weights_tensor),
+            "ReplaceAtenConvolutionWithCadenceConvolutionPass",
+        )
+
+        target = (
+            exir_ops.edge.cadence.conv2d.default
+            if spatial_dims == 2
+            else exir_ops.edge.cadence.conv3d.default
+        )
+        replacement = next(
+            node for node in graph_after_passes.graph.nodes if node.target == target
+        )
+        bias_node = replacement.args[2]
+        self.assertIsInstance(bias_node, torch.fx.Node)
+        assert isinstance(bias_node, torch.fx.Node)
+        self.assertEqual(exir_ops.edge.aten.full.default, bias_node.target)
+        self.assertEqual((out_channels,), tuple(bias_node.meta["val"].shape))
+        self.assertEqual(1234, bias_node.meta["debug_handle"])
+        self.assertNotIn("stack_trace", bias_node.meta)
+
+    @torch.no_grad()
+    def test_materialized_zero_bias_metadata_matches_runtime_value(self) -> None:
+        for weight_dtype, bias_dtype in (
+            (torch.float32, torch.float32),
+            (torch.float64, torch.float64),
+        ):
+            with self.subTest(weight_dtype=weight_dtype):
+                graph = torch.fx.Graph()
+                input_node = graph.placeholder("x")
+                weight = graph.placeholder("weight")
+                weight.meta["val"] = torch.ones(6, 2, 3, dtype=weight_dtype)
+                source = graph.call_function(
+                    exir_ops.edge.aten.relu.default,
+                    args=(input_node,),
+                )
+                source.meta["val"] = torch.ones(2, 6, 4, dtype=torch.float32)
+                source.meta["debug_handle"] = 1234
+
+                zero_bias = ReplaceAtenConvolutionWithCadenceConvolutionPass().materialize_zero_bias(
+                    source,
+                    weight,
+                )
+
+                self.assertEqual(exir_ops.edge.aten.full.default, zero_bias.target)
+                self.assertEqual(([6], 0.0), zero_bias.args)
+                self.assertEqual({"dtype": bias_dtype}, zero_bias.kwargs)
+                torch.testing.assert_close(
+                    zero_bias.meta["val"],
+                    weight.meta["val"].new_zeros(6, dtype=bias_dtype),
+                )
+                self.assertEqual(bias_dtype, zero_bias.meta["val"].dtype)
+                self.assertEqual(
+                    weight.meta["val"].device,
+                    zero_bias.meta["val"].device,
+                )
+                self.assertEqual(1234, zero_bias.meta["debug_handle"])
+                self.assertLess(
+                    list(graph.nodes).index(zero_bias),
+                    list(graph.nodes).index(source),
+                )
+
+    @torch.no_grad()
+    def test_rejected_convolution_does_not_materialize_zero_bias(self) -> None:
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(1, 4, 8))
+        weights = builder.placeholder("weights", torch.randn(6, 4, 3))
+        convolution = builder.call_operator(
+            op=exir_ops.edge.aten.convolution.default,
+            args=(x, weights, None, [1], [0], [1], False, [1], 1),
+        )
+        builder.output([convolution])
+
+        result = cast(
+            PassResult,
+            ReplaceAtenConvolutionWithCadenceConvolutionPass()(
+                builder.get_graph_module()
+            ),
+        )
+
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            1,
+            count_node(result.graph_module, exir_ops.edge.aten.convolution.default),
+        )
+        self.assertEqual(
+            0,
+            count_node(result.graph_module, exir_ops.edge.aten.full.default),
+        )
+
     @expand(
         [
             [(1, 8, 16), 8, 16, 3],
@@ -649,6 +809,16 @@ def test_replace_aten_transposed_conv_with_cadence_transposed_conv(
             ),
             1,
         )
+        for node in graph_after_passes.graph.nodes:
+            if node.target == exir_ops.edge.cadence.transposed_convolution.default:
+                self.assertIsNotNone(node.args[2])
+                if not bias_enabled:
+                    bias_node = node.args[2]
+                    self.assertIsInstance(bias_node, torch.fx.Node)
+                    self.assertEqual(
+                        exir_ops.edge.aten.full.default,
+                        cast(torch.fx.Node, bias_node).target,
+                    )
 
     @expand(
         [
@@ -947,7 +1117,7 @@ def test_replace_repeat_with_cat(self) -> None:
     def test_replace_masked_scalar_tensor_with_full(
         self,
         shape: Tuple[int],
-        mask_shape: Union[Tuple[int, ...], None] = None,
+        mask_shape: Optional[Tuple[int, ...]] = None,
     ) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn(*shape, dtype=torch.float32))