decorator

Author: yujiaoliang

front/py/deepx/nn/__init__.py

@@ -1,6 +1,6 @@
 from .deepxir import *
 from .modules import __all__ as _modules_all
 __all__ = [
-    "DeepxIR","DeepxIRResp",
+    "DeepxIR","DeepxIRResp","deepx_op"
     *_modules_all
     ]

front/py/deepx/nn/deepxir.py

@@ -250,4 +250,42 @@ def FuseFunc(f):
 
     f()
 
-    return
+    return
+
+from functools import wraps
+import inspect
+
+def deepx_op(opname):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            sig = inspect.signature(func)
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+
+            params = [Param.tensor(v) for k, v in bound.arguments.items() if k != 'out']
+            returns = [Param.tensor(bound.arguments['out'])]
+
+            ir = DeepxIR(opname, params, returns, kwargs.get('author', None))
+            send(ir)
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator
+
+def deepx_subgraph(opname):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            sig = inspect.signature(func)
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+
+            params = [Param.tensor(v) for k, v in bound.arguments.items() if k != 'out']
+            returns = [Param.tensor(bound.arguments['out'])]
+
+            ir = DeepxIR(opname, params, returns, kwargs.get('author', None))
+            # 修改这里的逻辑
+            send(ir)
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator

front/py/deepx/nn/functional/rtf_matmul.py

@@ -2,11 +2,16 @@
 from deepx.nn import DeepxIR,Param
 from deepx.scheduler import send
 
+
 def rtf_matmul(a:Tensor,b:Tensor,out: Tensor ,author='cublas',bench:int=None):
     args=[Param.tensor(a),Param.tensor(b)]
     returns=[Param.tensor(out)]
     ir=DeepxIR("matmul", args, returns, author)
     if bench is not None:
         ir._metadata.openbench(bench)
     send(ir)
+    return out
+
+@deepx_op("matmul")
+def rtf_matmul(a:Tensor,b:Tensor,out: Tensor ,author='cublas',bench:int=None):
     return out