Replace custom op pad with aten op, post-export

Differential Revision: D60941693

Pull Request resolved: #4603

Author: lucylq

examples/models/flamingo/passes/__init__.py

@@ -0,0 +1,31 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import torch
+from executorch.exir.pass_base import ExportPass
+from executorch.extension.llm.custom_ops import preprocess_custom_ops  # noqa
+
+
+class ReplaceCustomOpsWithAtenOpsPass(ExportPass):
+    """
+    Goes through all ops and replaces custom ops with aten ops. In some cases
+    aten ops cannot be exported due to dynamism, eg. pad in flamingo preprocess.
+    Use a custom op to pass export, and replace it with the aten op post-export,
+    which avoids re-writing the op in C++.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op._name == "preprocess::pad":
+            return super().call_operator(
+                torch.ops.aten.constant_pad_nd.default, args, kwargs, meta
+            )
+
+        return super().call_operator(op, args, kwargs, meta)

examples/models/flamingo/passes/replace_custom_ops_with_aten_ops_pass.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import unittest
+
+from typing import List
+
+import torch
+from executorch.exir import EdgeCompileConfig, to_edge
+
+from .replace_custom_ops_with_aten_ops_pass import ReplaceCustomOpsWithAtenOpsPass
+
+
+class TestPasses(unittest.TestCase):
+    def test_replace_custom_ops_with_aten_ops_pass(self) -> None:
+        from executorch.extension.llm.custom_ops import preprocess_custom_ops  # noqa
+
+        class Pad(torch.nn.Module):
+            def forward(self, x: torch.Tensor, padding: List[int]) -> torch.Tensor:
+                return torch.ops.preprocess.pad.default(x, padding)
+
+        pad = Pad()
+
+        image_tensor = torch.ones([3, 4, 5])
+        padding = [0, 2, 0, 1]
+
+        edge_prog = to_edge(
+            torch.export.export(pad, (image_tensor, padding), strict=False),
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
+
+        # Check that the custom op exists in the graph, and aten op does not.
+        edge_nodes = [node.name for node in edge_prog.exported_program().graph.nodes]
+        assert "constant_pad_nd" not in edge_nodes
+        assert "preprocess_pad_default" in edge_nodes
+
+        edge_prog = edge_prog.transform([ReplaceCustomOpsWithAtenOpsPass()])
+
+        # After running replace_custom_ops_with_aten_ops pass, the custom op
+        # should be replaced with aten op.
+        post_transform_nodes = [
+            node.name for node in edge_prog.exported_program().graph.nodes
+        ]
+        assert "constant_pad_nd" in post_transform_nodes
+        assert "preprocess_pad_default" not in post_transform_nodes

examples/models/flamingo/passes/test_passes.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import torch
 from executorch.exir.dialects._ops import ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload