diff --git a/backends/apple/coreai/passes/__init__.py b/backends/apple/coreai/passes/__init__.py new file mode 100644 index 00000000000..e1355736369 --- /dev/null +++ b/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"] diff --git a/backends/apple/coreai/passes/narrow_dtypes.py b/backends/apple/coreai/passes/narrow_dtypes.py new file mode 100644 index 00000000000..6bdd88cc885 --- /dev/null +++ b/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 diff --git a/backends/apple/coreai/passes/test/__init__.py b/backends/apple/coreai/passes/test/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/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. diff --git a/backends/apple/coreai/passes/test/test_narrow_dtypes.py b/backends/apple/coreai/passes/test/test_narrow_dtypes.py new file mode 100644 index 00000000000..d9d295ec0b6 --- /dev/null +++ b/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()