Add a pass to replace nodes with empty tensors with full.

Summary: Remove subgraphs of ops that produce empty tensors at the end. `ReplaceEmptyTensorsWithFullPass` both does the replacement and dead code elimination.

Reviewed By: zonglinpeng

Differential Revision: D68907459

Author: hsharma35

backends/cadence/aot/graph_builder.py

@@ -6,10 +6,10 @@
 from typing import Optional, Sequence, Union
 
 import torch
-from executorch.exir.pass_base import ExportPass, NodeMetadata, PassResult, ProxyValue
+from executorch.exir.pass_base import ExportPass, NodeMetadata, PassResult, ProxyValue, Argument
 from torch._dispatch.python import enable_python_dispatcher
 from torch._subclasses import FakeTensor, FakeTensorMode
-from torch.fx.node import Argument, Target
+from torch.fx.node import Target
 from torch.utils import _pytree as pytree
 
 

backends/cadence/aot/replace_ops.py

@@ -2070,11 +2070,32 @@ def call_operator(
             meta,
         )
 
+@register_cadence_pass(CadencePassAttribute(opt_level=1))
+class ReplaceEmptyTensorsWithFullPass(ExportPass):
+    """Replaces nodes that produce empty tensors with full nodes."""
+
+    def call_operator(self, op, args, kwargs, meta):
+        val = meta.data.get("val", None)
+        if isinstance(val, torch.Tensor) and val.numel() == 0:
+            return super().call_operator(
+                exir_ops.edge.aten.full.default,
+                args=(val.shape, 0),
+                kwargs={"dtype": val.dtype},
+                meta=meta,
+            )
+        return super().call_operator(op, args, kwargs, meta)
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        ret = super().call(graph_module)
+        modified = ret.graph_module.graph.eliminate_dead_code() or ret.modified
+        return PassResult(ret.graph_module, modified)
+
 
 # This class encapsulates all the functions that replace/switch one op in the
 # graph with another.
 class CadenceReplaceOpsInGraph:
     passes = [
+        ReplaceEmptyTensorsWithFullPass,
         ReplaceFunctionallyEquivalentOpTargets,
         ReplaceTCopyWithTransposePass,
         ReplacePermuteWithTransposePass,

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -7,7 +7,10 @@
 import torch.nn.functional as F
 from executorch.backends.cadence.aot import compiler
 from executorch.backends.cadence.aot.compiler import export_to_edge, quantize_pt2
-from executorch.backends.cadence.aot.graph_builder import single_op_builder
+from executorch.backends.cadence.aot.graph_builder import (
+    GraphBuilder,
+    single_op_builder,
+)
 from executorch.backends.cadence.aot.pass_utils import count_node
 from executorch.backends.cadence.aot.replace_ops import (
     ForceChannelLastForConvPass,
@@ -18,6 +21,7 @@
     ReplaceConstantPadNdWithSlicePass,
     ReplaceConvolutionOptionalArgsWithConcreteArgsPass,
     ReplaceConvWithIm2RowAndLinear,
+    ReplaceEmptyTensorsWithFullPass,
     ReplaceFunctionallyEquivalentOpTargets,
     ReplaceIm2RowWithViewPass,
     ReplaceLinearWithFullyConnectedOpPass,
@@ -1681,3 +1685,57 @@ def test_cat_insert_transpose(self):
             count_node(gm_after_pass, exir_ops.edge.aten.transpose_copy.int),
             3,
         )
+
+
+class TestReplaceEmptyTensorsWithFullPass(unittest.TestCase):
+    def _get_slice_empty_gm(self) -> torch.fx.GraphModule:
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(4))
+        # This is empty (numel == 0).
+        slice0 = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor, (x, 0, 0, 0)
+        )
+        # Copy of x.
+        slice1 = builder.call_operator(exir_ops.edge.aten.slice_copy.Tensor, (x,))
+        cat = builder.call_operator(
+            exir_ops.edge.aten.cat.default,
+            ((slice0, slice1),),
+        )
+        builder.output([cat])
+        return builder.get_graph_module()
+
+    def test_slice_no_transpose_if_already_outermost(self):
+        gm = self._get_slice_empty_gm()
+        self.assertEqual(
+            len(
+                gm.graph.find_nodes(
+                    op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+                )
+            ),
+            2,
+        )
+        self.assertEqual(
+            len(
+                gm.graph.find_nodes(
+                    op="call_function", target=exir_ops.edge.aten.full.default
+                )
+            ),
+            0,
+        )
+        updated_gm = ReplaceEmptyTensorsWithFullPass()(gm).graph_module
+        self.assertEqual(
+            len(
+                updated_gm.graph.find_nodes(
+                    op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+                )
+            ),
+            1,
+        )
+        self.assertEqual(
+            len(
+                updated_gm.graph.find_nodes(
+                    op="call_function", target=exir_ops.edge.aten.full.default
+                )
+            ),
+            1,
+        )