fix: replace front dynamo export to torch export

Author: bitzyz

python/src/infinitensor/converter/registry.py

@@ -8,57 +8,60 @@
 class ConverterRegistry:
 
     def __init__(self):
-        self._module_converters: Dict[torch.nn.Module, Callable[[fx.Node], None]] = {}
-        self._method_converters: Dict[str, Callable[[fx.Node], None]] = {}
+        # { "aten.matmul": { "default": fn, "out": fn, None: fn } }
+        self._method_converters: Dict[str, Dict[Optional[str], Callable]] = {}
 
-    def register_module(self, module_class):
-        """装饰器：注册模块转换器"""
-
-        def decorator(func):
-            self._module_converters[module_class] = func
-            return func
-
-        return decorator
-
-    def register_method(self, method_name: str):
+    def register(self, op_name: str, overload: Optional[str] = None):
         """装饰器：注册方法和函数转换器"""
 
         def decorator(func):
-            self._method_converters[method_name] = func
+            self._method_converters.setdefault(op_name, {})[overload] = func
             return func
 
         return decorator
 
-    def get_module_converter(self, module_class) -> Optional[Callable]:
-        """获取模块转换器"""
-        # 也检查父类
-        if module_class in self._module_converters:
-            return self._module_converters[module_class]
-        return None
-
-    def get_method_converter(self, method_name: str) -> Optional[Callable]:
+    def get_method_converter(self, op_name: str, overload: Optional[str] = None) -> Optional[Callable]:
         """获取方法和函数转换器"""
-        if method_name in self._method_converters:
-            return self._method_converters[method_name]
-        return None
+        if op_name in self._method_converters:
+            table = self._method_converters[op_name]
+            if overload:
+                if overload in table:
+                    return table[overload]
+                else:
+                    raise ValueError(f"Unsupported op.overload : {op_name}_{overload}")
+            else:
+                if None in table:
+                    return table[None]
+                else:
+                    raise ValueError(f"Unsupported op.overload : {op_name}")
+        else:
+            raise ValueError(f"Unsupported op : {op_name}")
+
 
     def update(self, custom_converters: Dict):
-        """更新转换器"""
+        """更新转换器
+        Args:
+            custom_converters:
+            {
+                (op_name, overload): converter
+            }
+        """
         for key, converter in custom_converters.items():
-            if inspect.isclass(key) and issubclass(key, nn.Module):
-                self._module_converters[key] = converter
-            elif isinstance(key, str):
+            if isinstance(key, tuple) and len(key) == 2:
+                op_name, overload = key
+                self._method_converters.setdefault(op_name, {})[overload] = converter
+            if isinstance(key, str):
                 self._method_converters[key] = converter
+            else:
+                raise TypeError(f"Invalid key type: {type(key)}")
 
     def clear(self):
         """清空所有转换器"""
-        self._module_converters.clear()
         self._method_converters.clear()
 
     def list_all_converters(self):
         """列出所有转换器"""
         return {
-            "modules": list(self._module_converters.keys()),
             "methods": list(self._method_converters.keys()),
         }
 

python/src/infinitensor/converter/unified_converters.py

@@ -1,17 +1,9 @@
 import torch.nn as nn
 from .registry import registry
 
+#https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
 
-@registry.register_module(nn.Linear)
-def convert_linear(translator, node, module):
-    x = translator.tensors[node.args[0]]
-    module = translator.named_modules[node.target]
-    weight = translator.params[module.weight]
-    bias = translator.params.get(module.bias, None)
-    translator.tensors[node] = translator.builder.gemm(x, weight, bias, transB=True)
-
-
-@registry.register_method("matmul")
+@registry.register("matmul","default")
 def convert_matmul(translator, node):
     a = translator.tensors[node.args[0]]
     b = translator.tensors[node.args[1]]