Reapply "[IMP][Launch Latency] misc improvements to compiler/kernel" and fix too aggressive optimization (#7851)

Author: stadlmax

python/test/unit/runtime/test_driver.py

@@ -10,11 +10,11 @@ def test_is_lazy():
     from importlib import reload
     reload(sys.modules["triton.runtime.driver"])
     reload(sys.modules["triton.runtime"])
-    mod = sys.modules[triton.runtime.driver.__module__]
-    assert isinstance(triton.runtime.driver.active, getattr(mod, "LazyProxy"))
-    assert triton.runtime.driver.active._obj is None
+    assert triton.runtime.driver._active is None
+    assert triton.runtime.driver._default is None
+    assert isinstance(triton.runtime.driver.active, getattr(triton.backends.driver, "DriverBase"))
+    assert isinstance(triton.runtime.driver.default, getattr(triton.backends.driver, "DriverBase"))
     utils = triton.runtime.driver.active.utils  # noqa: F841
-    assert issubclass(triton.runtime.driver.active._obj.__class__, getattr(triton.backends.driver, "DriverBase"))
 
 
 def test_kernel_in_thread(device):

python/triton/compiler/compiler.py

@@ -434,13 +434,14 @@ def __init__(self, src, metadata_group, hash):
         # (e.g., checking amount of shared memory on current device)
         self.module = None
         self.function = None
+        self._run = None
 
     def _init_handles(self):
         if self.module is not None:
             return
         device = driver.active.get_current_device()
         # create launcher
-        self.run = driver.active.launcher_cls(self.src, self.metadata)
+        self._run = driver.active.launcher_cls(self.src, self.metadata)
         # not enough shared memory to run the kernel
         max_shared = max_shared_mem(device)
         if self.metadata.shared > max_shared:
@@ -461,10 +462,10 @@ def _init_handles(self):
         if knobs.runtime.kernel_load_end_hook is not None:
             knobs.runtime.kernel_load_end_hook(self.module, self.function, self.name, self.metadata_group, self.hash)
 
-    def __getattribute__(self, name):
-        if name == 'run':
-            self._init_handles()
-        return super().__getattribute__(name)
+    @property
+    def run(self):
+        self._init_handles()
+        return self._run
 
     def launch_metadata(self, grid, stream, *args):
         if knobs.runtime.launch_enter_hook is None:

python/triton/runtime/driver.py

@@ -2,8 +2,6 @@
 
 from ..backends import backends, DriverBase
 
-from typing import Any, Callable, Generic, TypeVar, Union
-
 
 def _create_driver() -> DriverBase:
     active_drivers = [x.driver for x in backends.values() if x.driver.is_active()]
@@ -12,52 +10,29 @@ def _create_driver() -> DriverBase:
     return active_drivers[0]()
 
 
-T = TypeVar("T")
-
-
-class LazyProxy(Generic[T]):
-
-    def __init__(self, init_fn: Callable[[], T]) -> None:
-        self._init_fn = init_fn
-        self._obj: Union[T, None] = None
-
-    def _initialize_obj(self) -> T:
-        if self._obj is None:
-            self._obj = self._init_fn()
-        return self._obj
-
-    def __getattr__(self, name) -> Any:
-        return getattr(self._initialize_obj(), name)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        if name in ["_init_fn", "_obj"]:
-            super().__setattr__(name, value)
-        else:
-            setattr(self._initialize_obj(), name, value)
-
-    def __delattr__(self, name: str) -> None:
-        delattr(self._initialize_obj(), name)
-
-    def __repr__(self) -> str:
-        if self._obj is None:
-            return f"<{self.__class__.__name__} for {self._init_fn} not yet initialized>"
-        return repr(self._obj)
-
-    def __str__(self) -> str:
-        return str(self._initialize_obj())
-
-
 class DriverConfig:
 
     def __init__(self) -> None:
-        self.default: LazyProxy[DriverBase] = LazyProxy(_create_driver)
-        self.active: Union[LazyProxy[DriverBase], DriverBase] = self.default
+        self._default: DriverBase | None = None
+        self._active: DriverBase | None = None
+
+    @property
+    def default(self) -> DriverBase:
+        if self._default is None:
+            self._default = _create_driver()
+        return self._default
+
+    @property
+    def active(self) -> DriverBase:
+        if self._active is None:
+            self._active = self.default
+        return self._active
 
     def set_active(self, driver: DriverBase) -> None:
-        self.active = driver
+        self._active = driver
 
     def reset_active(self) -> None:
-        self.active = self.default
+        self._active = self.default
 
 
 driver = DriverConfig()