[Eager] Add eager to graph fallback API

Author: YWHyuk

PyTorchSimDevice/torch_openreg/openreg/__init__.py

@@ -243,6 +243,43 @@ def launch_model(model, *args, stream_index=0, timestamp=0, **kwargs):
 from .random import *  # noqa: F403
 from .amp import *
 
+def eager_to_compile(op_name):
+    """
+    Register an eager mode operation as a graph-based implementation using torch.compile().
+
+    Args:
+        op_name: Operator name (e.g., "aten::mul.Tensor")
+
+    Example:
+        torch.npu.eager_to_compile("aten::mul.Tensor")
+    """
+    def wrapper(*args, **kwargs):
+        @torch.compile(dynamic=False)
+        def dummy_graph(*args, **kwargs):
+            # Convert "aten::mul.Tensor" -> torch.ops.aten.mul.Tensor
+            namespace, op_path = op_name.split("::", 1)
+            op_path_parts = op_path.split(".")
+            op = torch.ops
+            for part in [namespace] + op_path_parts:
+                op = getattr(op, part)
+            return op(*args, **kwargs)
+        return dummy_graph(*args, **kwargs)
+
+    torch.library.impl(op_name, "npu", wrapper)
+
+def register_eager_to_compile(ops):
+    """
+    Register multiple operators at once using eager_to_compile.
+
+    Args:
+        ops: List of operator names (e.g., ["aten::mul.Tensor", "aten::add.Tensor"])
+
+    Example:
+        torch.npu.register_eager_to_compile(["aten::mul.Tensor", "aten::add.Tensor"])
+    """
+    for op_name in ops:
+        eager_to_compile(op_name)
+
 __all__ = [
     "device",
     "device_count",
@@ -269,4 +306,6 @@ def launch_model(model, *args, stream_index=0, timestamp=0, **kwargs):
     "synchronize",
     "get_tog_simulator",
     "set_tog_simulator",
+    "eager_to_compile",
+    "register_eager_to_compile",
 ]

tests/test_eager.py

@@ -1,13 +1,12 @@
 import torch
 
-@torch.library.impl("aten::mul.Tensor", "npu")
-def my_fallback(x, y):
-    raise NotImplementedError("Fallback called")
+torch.npu.register_eager_to_compile(["aten::mul.Tensor", "aten::add.Tensor"])
 
 if __name__ == "__main__":
     #torch.npu.register_fallback_op("aten::add.out", my_fallback)
     device = torch.device("npu:0")
     x = torch.ones(10, 10).to(device)
     y = torch.ones(10, 10).to(device)
     z = x * y
+    z = x + z
     print(z.cpu())