[ExecuTorch][Core AI][1/x] Graph pass: int64/float64 dtype narrowing

backends/apple/coreai/passes/__init__.py

@@ -0,0 +1,11 @@
+# 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.
+
+from executorch.backends.apple.coreai.passes.narrow_dtypes import (
+    NarrowToCoreAIDtypesPass,
+)
+
+__all__ = ["NarrowToCoreAIDtypesPass"]

backends/apple/coreai/passes/narrow_dtypes.py

@@ -0,0 +1,104 @@
+# 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.
+
+"""Boundary dtype-narrowing pass for the Core AI backend."""
+
+import torch
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
+from executorch.exir.pass_base import PassResult
+from torch.fx import GraphModule
+from torch.fx.passes.fake_tensor_prop import FakeTensorProp
+
+# Core AI represents only up-to-32-bit dtypes (coreai-torch narrows int64->si32
+# and float64->f32 internally), so it can't carry a 64-bit tensor across the
+# delegate boundary.
+_NARROW = {torch.int64: torch.int32, torch.float64: torch.float32}
+_SIXTY_FOUR_BIT = (torch.int64, torch.float64)
+
+
+class NarrowToCoreAIDtypesPass:
+    """Narrow int64/float64 to 32-bit inside the graph while preserving I/O.
+
+    Casts int64/float64 graph inputs to 32-bit right after the placeholder,
+    re-propagates dtypes, then casts 32-bit values feeding int64/float64 outputs
+    back to 64-bit. The model's external input/output dtypes stay unchanged,
+    while the interior (what Core AI sees) becomes 32-bit.
+
+    The inserted ``_to_copy`` casts keep a 64-bit operand/result, so
+    :class:`CoreAIPartitioner` (which rejects nodes with int64/float64 tensors)
+    leaves them outside the delegate; the delegate boundary then sees only 32-bit
+    tensors that match what coreai declares.
+
+    Runs pre-partition (e.g. via ``get_default_passes``) at either the ATen or
+    edge level (it picks the matching ``_to_copy`` overload).
+    """
+
+    def __call__(self, graph_module: GraphModule) -> PassResult:
+        graph = graph_module.graph
+        placeholders = [n for n in graph.nodes if n.op == "placeholder"]
+        if not any(_is_64bit(p.meta.get("val")) for p in placeholders):
+            return PassResult(graph_module, False)
+
+        uses_edge = any(
+            isinstance(node.target, EdgeOpOverload)
+            for node in graph.nodes
+            if node.op == "call_function"
+        )
+        if uses_edge:
+            from executorch.exir.dialects._ops import ops as exir_ops
+
+            to_copy = exir_ops.edge.aten._to_copy.default
+        else:
+            to_copy = torch.ops.aten._to_copy.default
+
+        output_node = next(n for n in graph.nodes if n.op == "output")
+        # Remember which outputs were 64-bit so we can restore them afterwards.
+        orig_output_dtypes = [
+            val.dtype if _is_64bit(val := _node_val(a)) else None
+            for a in output_node.args[0]
+        ]
+
+        # 1. Narrow 64-bit inputs right after the placeholder.
+        for placeholder in placeholders:
+            val = placeholder.meta.get("val")
+            if not _is_64bit(val):
+                continue
+            narrow_dtype = _NARROW[val.dtype]
+            with graph.inserting_after(placeholder):
+                cast = graph.call_function(
+                    to_copy, (placeholder,), {"dtype": narrow_dtype}
+                )
+            placeholder.replace_all_uses_with(cast)
+            cast.args = (placeholder,)  # restore the cast's own (detached) input
+            cast.meta["val"] = val.to(narrow_dtype)
+
+        # 2. Re-propagate dtypes through the (now-32-bit) interior.
+        FakeTensorProp(graph_module).propagate(*[p.meta["val"] for p in placeholders])
+
+        # 3. Widen 32-bit values feeding originally-64-bit outputs back.
+        out_args = list(output_node.args[0])
+        for i, arg in enumerate(out_args):
+            orig_dtype = orig_output_dtypes[i]
+            if orig_dtype is None or not isinstance(arg, torch.fx.Node):
+                continue
+            cur = _node_val(arg)
+            if isinstance(cur, torch.Tensor) and cur.dtype != orig_dtype:
+                with graph.inserting_before(output_node):
+                    widen = graph.call_function(to_copy, (arg,), {"dtype": orig_dtype})
+                widen.meta["val"] = cur.to(orig_dtype)
+                out_args[i] = widen
+        output_node.args = (tuple(out_args),)
+
+        graph_module.recompile()
+        return PassResult(graph_module, True)
+
+
+def _node_val(arg):
+    return arg.meta.get("val") if isinstance(arg, torch.fx.Node) else None
+
+
+def _is_64bit(val) -> bool:
+    return isinstance(val, torch.Tensor) and val.dtype in _SIXTY_FOUR_BIT

backends/apple/coreai/passes/test/__init__.py

@@ -0,0 +1,5 @@
+# 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.

backends/apple/coreai/passes/test/test_narrow_dtypes.py

@@ -0,0 +1,95 @@
+# 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.
+
+"""Unit tests for :class:`NarrowToCoreAIDtypesPass`."""
+
+import unittest
+
+import torch
+import torch.nn as nn
+
+from executorch.backends.apple.coreai.passes.narrow_dtypes import (
+    NarrowToCoreAIDtypesPass,
+)
+from executorch.exir import to_edge
+
+
+def _edge_gm(module, example_inputs):
+    ep = torch.export.export(module.eval(), example_inputs)
+    return to_edge(ep).exported_program().graph_module
+
+
+def _placeholders(gm):
+    return [n for n in gm.graph.nodes if n.op == "placeholder"]
+
+
+def _copy_nodes(gm):
+    return [
+        n
+        for n in gm.graph.nodes
+        if n.op == "call_function" and "_to_copy" in str(n.target)
+    ]
+
+
+def _dtype(node):
+    val = node.meta.get("val")
+    return getattr(val, "dtype", None)
+
+
+class _Add(nn.Module):
+    def forward(self, x):
+        return x + x
+
+
+class NarrowToCoreAIDtypesPassTest(unittest.TestCase):
+    def test_int64_boundary_narrowed_and_widened(self):
+        # x + x with an int64 input: output dtype follows the input, so the pass
+        # must narrow the input to i32 and widen the output back to i64.
+        gm = _edge_gm(_Add(), (torch.randint(-8, 8, (2, 8), dtype=torch.int64),))
+        before = len(_copy_nodes(gm))
+        result = NarrowToCoreAIDtypesPass()(gm)
+        self.assertTrue(result.modified)
+
+        # External input placeholder dtype is unchanged.
+        self.assertEqual(_dtype(_placeholders(gm)[0]), torch.int64)
+        # Boundary casts were inserted (input narrow + output widen).
+        cast_dtypes = {_dtype(c) for c in _copy_nodes(gm)}
+        self.assertGreater(len(_copy_nodes(gm)), before)
+        self.assertIn(torch.int32, cast_dtypes)  # input narrowed to i32
+        self.assertIn(torch.int64, cast_dtypes)  # output widened back to i64
+        # External output dtype preserved as int64.
+        out = gm.graph.output_node().args[0][0]
+        self.assertEqual(_dtype(out), torch.int64)
+
+    def test_interior_is_narrowed_to_32bit(self):
+        gm = _edge_gm(_Add(), (torch.randint(-8, 8, (2, 8), dtype=torch.int64),))
+        NarrowToCoreAIDtypesPass()(gm)
+        add = next(
+            n
+            for n in gm.graph.nodes
+            if n.op == "call_function" and "add" in str(n.target)
+        )
+        self.assertEqual(_dtype(add), torch.int32)
+
+    def test_float64_narrowed_io_preserved(self):
+        gm = _edge_gm(_Add(), (torch.randn(2, 8, dtype=torch.float64),))
+        result = NarrowToCoreAIDtypesPass()(gm)
+        self.assertTrue(result.modified)
+        self.assertEqual(_dtype(_placeholders(gm)[0]), torch.float64)  # I/O kept
+        cast_dtypes = {_dtype(c) for c in _copy_nodes(gm)}
+        self.assertIn(torch.float32, cast_dtypes)
+        self.assertIn(torch.float64, cast_dtypes)
+
+    def test_noop_when_no_64bit(self):
+        gm = _edge_gm(_Add(), (torch.randn(2, 8),))  # float32 only
+        before = len(_copy_nodes(gm))
+        result = NarrowToCoreAIDtypesPass()(gm)
+        self.assertFalse(result.modified)
+        self.assertEqual(len(_copy_nodes(gm)), before)  # nothing inserted
+
+
+if __name__ == "__main__":
+    unittest.main()