[Backend]add_triton_version_event (flagos-ai#3336)

* add_triton_version_event

* fix codestyle

* fix bug

* fix codestyle

* fix code mode

Author: Galaxy1458

src/flag_gems/__init__.py

@@ -14,15 +14,20 @@
 from flag_gems.ops import *  # noqa: F403
 from flag_gems.patches import *  # noqa: F403
 from flag_gems.runtime import flagtune
-from flag_gems.runtime.register import Register
+from flag_gems.runtime.backend import SpecOpRegistrar
+from flag_gems.runtime.op_registrar import GeneralOpRegistrar
 
 __version__ = "5.0.2"
 device = runtime.device.name
 vendor_name = runtime.device.vendor_name
+backend_info = runtime.device
 aten_lib = torch.library.Library("aten", "IMPL")
-registrar = Register
+
+# Register all ops in the current backend with SpecOpRegistrar to support architecture-specialized implementations
+SpecOpRegistrar(globals()).apply()
+
+registrar = GeneralOpRegistrar
 current_work_registrar = None
-runtime.replace_customized_ops(globals())
 AUTOGRAD_DISPATCH_KEY = torch._C.DispatchKey.Autograd.name
 
 
@@ -683,7 +688,7 @@ def __init__(self, exclude=None, include=None, record=False, once=False, path=No
         self.lib = torch.library.Library("aten", "IMPL")
         self.exclude = exclude if isinstance(exclude, (list, tuple, set, str)) else []
         self.include = include if isinstance(include, (list, tuple, set, str)) else []
-        self.registrar = Register
+        self.registrar = GeneralOpRegistrar
         self.record = record
         self.once = once
         self.path = path

src/flag_gems/runtime/__init__.py

@@ -1,9 +1,9 @@
 from . import backend, common, error
 from .backend.device import DeviceDetector
-from .configloader import ConfigLoader
+from .configs_loader import TunedConfigLoader
 from .flagtune import flagtune, flagtune_enabled
 
-config_loader = ConfigLoader()
+config_loader = TunedConfigLoader()
 device = DeviceDetector()
 
 """
@@ -26,24 +26,6 @@ def get_heuristic_config(op_name):
     return config_loader.get_heuristics_config(op_name)
 
 
-def replace_customized_ops(_globals):
-    event = backend.BackendArchEvent()
-    arch_specific_ops = event.get_arch_ops() if event.has_arch else None
-    extended_ops = backend.get_customized_ops(device.vendor_name)
-    if device.vendor != common.vendors.NVIDIA:
-        try:
-            for fn_name, fn in extended_ops:
-                _globals[fn_name] = fn
-        except RuntimeError as e:
-            error.customized_op_replace_error(e)
-    if arch_specific_ops:
-        try:
-            for fn_name, fn in arch_specific_ops:
-                _globals[fn_name] = fn
-        except RuntimeError as e:
-            error.customized_op_replace_error(e)
-
-
 def get_expand_config(op_name, yaml_path=None):
     return config_loader.get_expand_config(op_name=op_name, yaml_path=yaml_path)
 
@@ -57,7 +39,7 @@ def ops_get_configs(op_name, pre_hook=None, yaml_path=None):
 
 
 __all__ = [
-    "ConfigLoader",
+    "TunedConfigLoader",
     "DeviceDetector",
     "backend",
     "common",

src/flag_gems/runtime/backend/__init__.py

@@ -8,6 +8,7 @@
 
 from ..common import vendors
 from . import backend_utils
+from .backend_utils import BackendEventBase
 
 
 class BackendState:
@@ -42,7 +43,63 @@ def __init__(self):
 _state = BackendState()
 
 
-class BackendArchEvent:
+class TritonVersionEvent(BackendEventBase):
+    _instance = None
+    has_version_spec = False
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, version=None):
+        self.has_version_spec = False
+        self.version = version if version is not None else self.get_version()
+        self.dir = self.get_version_spec_dir()
+        if self.dir and Path(self.dir).exists():
+            self.module = self.get_version_spec_module()
+            self.has_version_spec = True
+
+    def is_available(self):
+        return self.has_version_spec
+
+    def get_version_spec_dir(self, path=None):
+        dir_name = f"triton_{self.version}"
+        backend_path = Path(path or _state.vendor_module.__path__[0])
+        backend_path = backend_path.parent if backend_path.is_file() else backend_path
+        excluded = ("ops", "fused")
+        return {
+            p.name: str(p)
+            for p in backend_path.iterdir()
+            if p.is_dir() and p.name not in excluded and not p.name.startswith("_")
+        }.get(dir_name, None)
+
+    def get_functions_from_module(self, module):
+        return inspect.getmembers(module, inspect.isfunction) if module else []
+
+    def get_version_spec_module(self):
+        module_name = f"triton_{self.version}"
+        path_dir = os.path.dirname(self.dir)
+        sys.path.insert(0, str(path_dir))
+        version_module = importlib.import_module(module_name)
+        sys.path.remove(str(path_dir))
+        return version_module
+
+    def get_ops(self):
+        return self.get_version_ops()
+
+    def get_version_ops(self):
+        pass
+
+    def get_version(self):
+        try:
+            import triton
+        except ImportError:
+            return None
+        return triton.__version__
+
+
+class BackendArchEvent(BackendEventBase):
     has_arch: bool = False
     _instance = None
     _initialized: bool = False
@@ -67,6 +124,9 @@ def __init__(self, backend=None):
             self.autotune_configs = self.get_autotune_configs()
             self.heuristics_configs = self.get_heuristics_configs()
 
+    def is_available(self):
+        return self.has_arch
+
     def get_functions_from_module(self, module):
         return inspect.getmembers(module, inspect.isfunction) if module else []
 
@@ -126,6 +186,10 @@ def get_arch_module(self):
         sys.path.remove(str(path_dir))
         return current_arch_module
 
+    def get_ops(self):
+        """Provide a unified interface for the upper layer"""
+        return self.get_arch_ops()
+
     def get_arch_ops(self):
         arch_specialized_ops = []
         sys.path.append(self.current_arch_path)
@@ -147,6 +211,23 @@ def get_arch_ops(self):
         return arch_specialized_ops
 
 
+class SpecOpRegistrar:
+    def __init__(self, _globals):
+        self._globals = _globals
+
+    def apply(self):
+        spec_events = self._get_specific_events()
+        for event in spec_events:
+            if not event.is_available():
+                continue
+            operators = event.get_ops()
+            for fn_name, fn in operators:
+                self._globals[fn_name] = fn
+
+    def _get_specific_events(self):
+        return (BackendArchEvent(), TritonVersionEvent())
+
+
 def _import_module_safe(module_name, vendor_name, module_type):
     """Helper to import a module with proper error handling."""
     try:
@@ -295,6 +376,11 @@ def get_customized_ops(vendor_name=None):
     return _state.customized_ops
 
 
+def get_ops(vendor_name=None):
+    """Provide a unified interface for the upper layer"""
+    return get_customized_ops(vendor_name)
+
+
 def get_unused_ops(vendor_name=None):
     global vendor_module  # noqa: F824
     get_vendor_module(vendor_name)

src/flag_gems/runtime/backend/backend_utils.py

@@ -39,6 +39,17 @@ def get_tune_config(vendor_name=None, file_mode="r", file_path=None):
     return config
 
 
+class BackendEventBase:
+    def __init__(self):
+        ...
+
+    def get_ops(self):
+        ...
+
+    def is_available(self):
+        ...
+
+
 @functools.lru_cache(maxsize=None)
 def _load_expand_config(file_path, file_mode="r"):
     with open(file_path, file_mode) as file:

src/flag_gems/runtime/configloader.py src/flag_gems/runtime/configs_loader.pysrc/flag_gems/runtime/configloader.py renamed to src/flag_gems/runtime/configs_loader.py

@@ -9,12 +9,12 @@
 from .backend.device import DeviceDetector
 
 
-class ConfigLoader(object):
+class TunedConfigLoader(object):
     _instance = None
 
     def __new__(cls, *args, **kargs):
         if cls._instance is None:
-            cls._instance = super(ConfigLoader, cls).__new__(cls)
+            cls._instance = super(TunedConfigLoader, cls).__new__(cls)
         return cls._instance
 
     def __init__(self):

src/flag_gems/runtime/register.py src/flag_gems/runtime/op_registrar.pysrc/flag_gems/runtime/register.py renamed to src/flag_gems/runtime/op_registrar.py

@@ -4,7 +4,7 @@
 from .backend.device import DeviceDetector
 
 
-class Register:
+class GeneralOpRegistrar:
     def __init__(
         self,
         config,

src/flag_gems/runtime/precision_register.py

@@ -16,7 +16,7 @@
     precision_config,
     write_precision_result,
 )
-from .register import Register
+from .op_registrar import GeneralOpRegistrar
 
 # Maximum tensor element count allowed for precision check
 # (skip if exceeded to avoid large tensor copy overhead)
@@ -177,7 +177,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-class PrecisionCheckRegister(Register):
+class PrecisionCheckRegister(GeneralOpRegistrar):
     """Register subclass that wraps every operator with precision checking.
 
     This class is only instantiated when the user has explicitly called