[frontend]: update api for frontend

Author: wlt-cortex

python/mrt/frontend/api.py

@@ -1,45 +1,107 @@
 import os
 import importlib
 import sys
+from functools import wraps
 
-from mrt.mir.optype import _DEFAULT_TYPE_INFER
+from mrt.mir.symbol import *
+from mrt.common.types import *
 
-FRONTEND = os.environ.get("FRONTEND", "pytorch")
+class Singleton(object):
+    def __new__(cls, *args, **kw):
+        if not hasattr(cls, '_instance'):
+            orig = super(Singleton, cls)
+        cls._instance = orig.__new__(cls, *args, **kw)
+        return cls._instance
+
+class DynamicModule(Singleton):
+    def __init__(self):
+        self._funcs = {}
+
+    def load_mod(self, frontend):
+        try:
+            frontend_module = importlib.import_module(f".{FRONTEND}", package="mrt.frontend")
+        except ImportError as e:
+            print(f"Error: Frontend '{FRONTEND}' cannot be imported: {e}")
+            return
+
+        for f in self._funcs:
+            if hasattr(frontend_module, f):
+                self._funcs[f] = getattr(frontend_module, f)
+            else:
+                print(f"Error: function '{f}' not found in frontend '{FRONTEND}'")
+
+        return self
+
+    def typedef_mod_function(self, func):
+        fname = func.__name__
+        self._funcs.setdefault(fname, None)
+
+        @wraps(func)
+        def _func_impl(*args, **kwargs):
+            assert self._funcs[fname] is not None, f"func:{fname} not registered in mod: {self._funcs.keys()}"
+            func(*args, **kwargs)
+            return self._funcs[fname](*args, **kwargs)
+        return _func_impl
+
+
+mod = DynamicModule()
+
+@mod.typedef_mod_function
+def create_executor(
+        symbol: MultiHeadSymbol, params: ParametersT,
+        device: str = "cpu",
+        target: str = "", # no use in pytorch frontend
+        ):
+    """ Create Runtime Executor for Model Inference. """
+    pass
 
-# Dynamically load the frontend module
-try:
-    frontend_module = importlib.import_module(f".{FRONTEND}", package="mrt.frontend")
-except ImportError as e:
-    print(f"Error: Frontend '{FRONTEND}' is not supported or cannot be imported: {e}")
-    sys.exit(1)
-
-# Register default type infer functions
-if hasattr(frontend_module, "type_infer"):
-    _DEFAULT_TYPE_INFER = frontend_module.type_infer
-else:
-    print(f"Error: Required function 'type_infer' not found in frontend '{FRONTEND}'")
-    sys.exit(1)
-
-# Try to get frontend_to_mrt function
-frontend_to_mrt = None
-if hasattr(frontend_module, "from_frontend"):
-    frontend_to_mrt = frontend_module.from_frontend
-elif hasattr(frontend_module, "pytorch_to_mrt"):
-    frontend_to_mrt = frontend_module.pytorch_to_mrt
-elif hasattr(frontend_module, "expr2symbol"):
-    frontend_to_mrt = frontend_module.expr2symbol
-else:
-    print(f"Error: Required function 'frontend_to_mrt' not found in frontend '{FRONTEND}'")
-    sys.exit(1)
-
-# Try to get mrt_to_frontend function
-mrt_to_frontend = None
-if hasattr(frontend_module, "to_frontend"):
-    mrt_to_frontend = frontend_module.to_frontend
-elif hasattr(frontend_module, "mrt_to_pytorch"):
-    mrt_to_frontend = frontend_module.mrt_to_pytorch
-elif hasattr(frontend_module, "symbol2expr"):
-    mrt_to_frontend = frontend_module.symbol2expr
-else:
-    print(f"Error: Required function 'mrt_to_frontend' not found in frontend '{FRONTEND}'")
-    sys.exit(1)
+@mod.typedef_mod_function
+def run_executor(
+        executor,
+        data: typing.Optional[np.ndarray] = None,
+        data_dict: ParametersT = {}
+        ) -> OpNumpyT:
+    """ Apply data to executor. """
+    pass
+
+@mod.typedef_mod_function
+def infer(
+        graph: MultiHeadSymbol,
+        params: ParametersT,
+        data: typing.Optional[np.ndarray] = None,
+        data_dict: ParametersT = {},
+        device: str = "cpu",
+        **kwargs):
+    """ Convinent Method to infer model. """
+    pass
+
+@mod.typedef_mod_function
+def data_from_frontend(data: typing.Any) -> OpNumpyT:
+    """ Convert Frontend Tensor to MRT DType. """
+    pass
+
+@mod.typedef_mod_function
+def data_to_frontend(data: OpNumpyT):
+    """ Convert MRT DType to Frontend Tensor. """
+    pass
+
+@mod.typedef_mod_function
+def model_from_frontend(
+        fe_model,
+        func_names: typing.List[str] = [ "main", ]
+        ) -> typing.Tuple[MultiHeadSymbol, ParametersT]:
+    """ Convert Frontend Graph to MRT Symbol/Params. """
+    pass
+
+@mod.typedef_mod_function
+def model_to_frontend(graph: MultiHeadSymbol, params: ParametersT,):
+    """ Convert MRT Symbol/Params to Frontend Graph. """
+    pass
+
+@mod.typedef_mod_function
+def type_infer(symbol: Symbol) -> Symbol:
+    """ Shape/DType Inference use Frontend API. """
+
+
+FRONTEND = os.environ.get("FRONTEND", "pytorch")
+mod.load_mod(FRONTEND)

python/mrt/frontend/pytorch/__init__.py

@@ -7,9 +7,17 @@
 """
 
 from .converter import pytorch_to_mrt, mrt_to_pytorch, type_infer
+from .types import data_to_mrt, data_to_torch
+from .vm import create_executor, run_executor, infer
 
 # Expose the required functions for the frontend API
 from_frontend = pytorch_to_mrt
 to_frontend = mrt_to_pytorch
 
-__all__ = ["pytorch_to_mrt", "mrt_to_pytorch", "from_frontend", "to_frontend", "type_infer"]
+model_from_frontend = pytorch_to_mrt
+model_to_frontend = mrt_to_pytorch
+
+data_from_frontend = data_to_torch
+data_to_frontend = data_to_mrt
+
+# __all__ = ["pytorch_to_mrt", "mrt_to_pytorch", "from_frontend", "to_frontend", "type_infer"]

python/mrt/frontend/pytorch/converter.py

@@ -139,8 +139,9 @@ def pytorch_to_mrt(
         ep: torch.export.ExportedProgram,
         func_names: typing.List[str] = [ "main", ]
         ) -> typing.Tuple[MultiHeadSymbol, ParametersT]:
-    env: typing.Dict[fx.Node, Symbol] = {}
+    env: typing.Dict[torch.Node, Symbol] = {}
 
+    assert isinstance(ep, torch.export.ExportedProgram), f"input not torch ExportedProgram, but {type(ep)}"
     param_vars, params = create_parameters(ep)
 
     def _retrieve_args(node):
@@ -163,8 +164,6 @@ def _retrieve_args(node):
         if "tensor_meta" in node.meta:
             meta_data = node.meta["tensor_meta"]
             shape = data_to_mrt(meta_data.shape)
-            #  shape = [str(s) if isinstance(s, torch.SymInt) else int(s) \
-            #          for s in meta_data.shape]
             dtype = data_to_mrt(meta_data.dtype)
         #  else:
         #      print(node.name, "has no tensor meta")
@@ -284,8 +283,6 @@ def _infer_single_op(sym: Symbol, env: typing.Dict[str, F.Tensor]) -> F.Tensor:
 
 def type_infer(symbol: Symbol) -> Symbol:
     """Infer shape and dtype for all symbols in the graph.
-
-    This function works by using torch.Tensor inference.
     """
     env: Dict[str, F.Tensor] = {}
 

python/mrt/frontend/pytorch/vm.py

@@ -30,7 +30,8 @@ def run_executor(
         data_dict[k] = torch.from_numpy(v).to(device)
     if data is not None:
         data = torch.from_numpy(data).to(device)
-    out = vm(data, **data_dict)
+    with torch.no_grad():
+        out = vm(data, **data_dict)
     return data_to_mrt(out.detach().cpu())
 
 def infer(graph: MultiHeadSymbol, params: ParametersT,

python/mrt/mir/optype.py

@@ -2,6 +2,7 @@
 
 from mrt.common.types import *
 from mrt.common import config
+from mrt.frontend import api
 #  from mrt.mir import op, opns
 #  from mrt.symbol import Symbol, transform
 
@@ -10,7 +11,6 @@
 from .symbol import Symbol, transform
 
 InferTypeT = typing.Callable[[Symbol], Symbol]
-_DEFAULT_TYPE_INFER = None
 _INFER_TYPE_REG: typing.Dict[str, InferTypeT] = {}
 
 def register_type_infer(
@@ -29,7 +29,7 @@ def infer_single(symbol: Symbol) -> Symbol:
     C = config.LogConfig.G()
 
     out = op.retrieve_operator(symbol)
-    _infer = _INFER_TYPE_REG.get(out.op_name, _DEFAULT_TYPE_INFER)
+    _infer = _INFER_TYPE_REG.get(out.op_name, api.type_infer)
     assert _infer is not None
 
     if symbol.is_near(*C.log_type_infer):

tests/frontend/pytorch/test_resnet50.py

@@ -162,4 +162,4 @@ def test_resnet50_type_infer():
     print("\n" + "="*60 + "\n")
 
     test_resnet50_infer()
-    #  test_resnet50_type_infer()
+    test_resnet50_type_infer()

tests/frontend/test_frontend_loading.py

@@ -9,20 +9,41 @@
 # Add the project root to the path
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
+from mrt.mir.symbol import *
+from mrt.mir import op
+
 def test_frontend_loading():
     print("Testing frontend loading...")
-    
+
+    x = op.variable("x", (3, 224), "float")
+    w = op.variable("w", (10, 224), "float")
+    y = op.nn_dense(x, w)
+
     # Test PyTorch frontend
     print("\n1. Testing PyTorch frontend:")
     os.environ["FRONTEND"] = "pytorch"
     try:
-        from mrt.frontend.api import FRONTEND, _DEFAULT_TYPE_INFER, frontend_to_mrt, mrt_to_frontend
+        from mrt.frontend.api import FRONTEND, model_from_frontend, model_to_frontend
         print(f"   Loaded frontend: {FRONTEND}")
-        print(f"   _DEFAULT_TYPE_INFER: {_DEFAULT_TYPE_INFER}")
-        print(f"   frontend_to_mrt: {frontend_to_mrt}")
-        print(f"   mrt_to_frontend: {mrt_to_frontend}")
+        print(f"   model_from_frontend: {model_from_frontend}")
+        print(f"   model_to_frontend: {model_to_frontend}")
+
+        import torch
+        fe_model = model_to_frontend(MultiHeadSymbol(main=y), {})
+        out = fe_model(torch.randn(3, 224), w = torch.rand(10, 224))
+        #  out = fe_model(torch.randn(3, 224))
+        #  print(out)
+        fe_model = torch.export.export(
+                fe_model,
+                args=(torch.randn(3, 224),),
+                kwargs={
+                    "w": torch.randn(10, 224)})
+        #  print(fe_model)
+        y1 = model_from_frontend(fe_model)
+
         print("   PyTorch frontend loaded successfully!")
     except Exception as e:
+        #  raise e
         print(f"   Error loading PyTorch frontend: {e}")
         return False