Reland upstream commit f7d2674: "[Knobs] Reduce overhead of reading knobs (#7841)" (#4966)

whitneywhtsang · web-flow · commit 558841e6a1a5 · 2025-08-27T13:07:09.000-04:00
This reverts commit 298480d.
diff --git a/python/src/ir.cc b/python/src/ir.cc
@@ -1904,6 +1904,85 @@ void init_triton_ir(py::module &&m) {
           py::call_guard<py::gil_scoped_release>());
 }
 
+bool str_eq_ignore_case(const char *s1, const char *s2, int n) {
+  for (int i = 0; i < n; ++i) {
+    if (tolower(s1[i]) != s2[i])
+      return false;
+  }
+  return true;
+}
+
+int strlen_max(const char *str, int max) {
+  for (int i = 0; i <= max; ++i) {
+    if (str[i] == '\0') {
+      return i;
+    }
+  }
+  return 0;
+}
+
+bool is_truthy(char *str) {
+  int len = strlen_max(str, 4);
+  switch (len) {
+  case 1:
+    return str[0] == '1' || tolower(str[0]) == 'y';
+  case 2:
+    return str_eq_ignore_case(str, "on", len);
+  case 3:
+    return str_eq_ignore_case(str, "yes", len);
+  case 4:
+    return str_eq_ignore_case(str, "true", len);
+  default:
+    return false;
+  }
+}
+
+PyObject *py_getenv(PyObject *self, PyObject *const *args, Py_ssize_t nargs) {
+  if (!(nargs == 1 || nargs == 2)) {
+    PyErr_SetString(PyExc_TypeError, "getenv expected 1 or 2 arguments");
+    return NULL;
+  }
+  PyObject *name = args[0];
+  PyObject *default_val = nargs == 2 ? args[1] : Py_None;
+  if (!PyUnicode_CheckExact(name)) {
+    PyErr_SetString(PyExc_TypeError, "name must be a string");
+    return NULL;
+  }
+  char *env_val = getenv(PyUnicode_AsUTF8(name));
+  if (!env_val) {
+    Py_INCREF(default_val);
+    return default_val;
+  }
+  return PyUnicode_FromString(env_val);
+}
+
+PyObject *py_getenv_bool(PyObject *self, PyObject *const *args,
+                         Py_ssize_t nargs) {
+  if (nargs != 2) {
+    PyErr_SetString(PyExc_TypeError, "getenv_bool expected 2 arguments");
+    return NULL;
+  }
+  PyObject *name = args[0];
+  PyObject *default_val = args[1];
+  if (!PyUnicode_CheckExact(name)) {
+    PyErr_SetString(PyExc_TypeError, "name must be a string");
+    return NULL;
+  }
+  char *env_val = getenv(PyUnicode_AsUTF8(name));
+  PyObject *res = default_val;
+  if (env_val) {
+    res = is_truthy(env_val) ? Py_True : Py_False;
+  }
+  Py_INCREF(res);
+  return res;
+}
+
+static PyMethodDef ModuleMethods[] = {
+    {"getenv", (PyCFunction)py_getenv, METH_FASTCALL, NULL},
+    {"getenv_bool", (PyCFunction)py_getenv_bool, METH_FASTCALL, NULL},
+    {NULL, NULL, 0, NULL} // sentinel
+};
+
 void init_triton_env_vars(py::module &m) {
   m.def("get_cache_invalidating_env_vars",
         []() -> std::map<std::string, std::string> {
@@ -1920,4 +1999,5 @@ void init_triton_env_vars(py::module &m) {
           }
           return ret;
         });
+  PyModule_AddFunctions(m.ptr(), ModuleMethods);
 }
diff --git a/python/triton/knobs.py b/python/triton/knobs.py
@@ -12,6 +12,8 @@
 from contextlib import contextmanager
 from typing import cast, Any, Callable, Generator, Generic, Optional, Protocol, Type, TypeVar, TypedDict, TYPE_CHECKING, Union
 
+from triton._C.libtriton import getenv, getenv_bool
+
 if TYPE_CHECKING:
     from .runtime.cache import CacheManager, RemoteCacheBackend
     from .runtime.jit import JitFunctionInfo, KernelParam
@@ -27,11 +29,6 @@ class Env:
 propagate_env: bool = True
 
 
-def getenv(key: str) -> Optional[str]:
-    res = os.getenv(key)
-    return res.strip() if res is not None else res
-
-
 def setenv(key: str, value: Optional[str]) -> None:
     if not propagate_env:
         return
@@ -64,32 +61,25 @@ def toenv(val: Any) -> Union[None, tuple[Optional[str]]]:
 SetType = TypeVar("SetType")
 GetType = TypeVar("GetType")
 
+_NOTHING = object()
+
 
 class env_base(Generic[SetType, GetType]):
 
-    def __init__(self, key: str, default: Union[SetType, Callable[[], SetType]]) -> None:
+    def __init__(self, key: str) -> None:
         self.key = key
-        self.default: Callable[[], SetType] = default if callable(default) else lambda: default
 
     def __set_name__(self, objclass: Type[object], name: str) -> None:
         self.name = name
 
     def __get__(self, obj: Optional[object], objclass: Optional[Type[object]]) -> GetType:
-        if obj is None:
-            raise AttributeError(f"Cannot access {type(self)} on non-instance")
-
-        if self.name in obj.__dict__:
-            return self.transform(obj.__dict__[self.name])
-        else:
+        py_val = obj.__dict__.get(self.name, _NOTHING)
+        if py_val is _NOTHING:
             return self.get()
-
-    @property
-    def env_val(self) -> str | None:
-        return getenv(self.key)
+        return self.transform(py_val)
 
     def get(self) -> GetType:
-        env = self.env_val
-        return self.transform(self.default() if env is None else self.from_env(env))
+        raise NotImplementedError()
 
     def __set__(self, obj: object, value: Union[SetType, Env]) -> None:
         if isinstance(value, Env):
@@ -107,54 +97,70 @@ def transform(self, val: SetType) -> GetType:
         # if GetType != SetType.
         return cast(GetType, val)
 
-    def from_env(self, val: str) -> SetType:
-        raise NotImplementedError()
-
 
 class env_str(env_base[str, str]):
 
-    def from_env(self, val: str) -> str:
-        return val
+    def __init__(self, key: str, default: str):
+        super().__init__(key)
+        self.default = default
+
+    def get(self) -> str:
+        return getenv(self.key, self.default)
+
+
+class env_str_callable_default(env_base[str, str]):
+
+    def __init__(self, key: str, default_factory: Callable[[], str]):
+        super().__init__(key)
+        self.default_factory = default_factory
+
+    def get(self) -> str:
+        env_val = getenv(self.key)
+        if env_val is None:
+            return self.default_factory()
+        return env_val
 
 
 class env_bool(env_base[bool, bool]):
 
-    def __init__(self, key: str, default: Union[bool, Callable[[], bool]] = False) -> None:
-        super().__init__(key, default)
+    def __init__(self, key: str, default: bool = False) -> None:
+        super().__init__(key)
+        self.default = default
 
-    def from_env(self, val: str) -> bool:
-        return val.lower() in ("1", "true", "yes", "on", "y")
+    def get(self) -> bool:
+        return getenv_bool(self.key, self.default)
 
 
 class env_int(env_base[int, int]):
 
-    def __init__(self, key: str, default: Union[int, Callable[[], int]] = 0) -> None:
-        super().__init__(key, default)
+    def __init__(self, key: str, default: int = 0) -> None:
+        super().__init__(key)
+        self.default = default
 
-    def from_env(self, val: str) -> int:
+    def get(self) -> int:
+        val = getenv(self.key)
+        if val is None:
+            return self.default
         try:
             return int(val)
         except ValueError as exc:
             raise RuntimeError(f"Unable to use {self.key}={val}: expected int") from exc
 
 
-class env_opt_base(Generic[GetType, SetType], env_base[Optional[GetType], Optional[SetType]]):
-
-    def __init__(self, key: str) -> None:
-        super().__init__(key, None)
-
-
 ClassType = TypeVar("ClassType")
 
 
-class env_class(Generic[ClassType], env_opt_base[Type[ClassType], Type[ClassType]]):
+class env_class(Generic[ClassType], env_base[Optional[Type[ClassType]], Optional[Type[ClassType]]]):
 
     def __init__(self, key: str, type: str) -> None:
         super().__init__(key)
         # We can't pass the type directly to avoid import cycles
         self.type = type
 
-    def from_env(self, val: str) -> Type[ClassType]:
+    def get(self) -> Optional[Type[ClassType]]:
+        val = getenv(self.key)
+        if val is None:
+            return None
         comps = val.split(":", 1)
         if len(comps) != 2:
             raise RuntimeError(f"Unable to read {self.key}: '{val}' isn't of the form MODULE:CLASS")
@@ -201,7 +207,7 @@ def from_path(path: str) -> Optional[IntelTool]:
             if version is None:
                 return None
             return IntelTool(path, version.group(1))
-        except subprocess.CalledProcessError:
+        except (subprocess.CalledProcessError, FileNotFoundError):
             return None
 
 
@@ -210,73 +216,65 @@ class env_nvidia_tool(env_base[str, NvidiaTool]):
     def __init__(self, binary: str) -> None:
         binary += sysconfig.get_config_var("EXE")
         self.binary = binary
-        super().__init__(f"TRITON_{binary.upper()}_PATH", lambda: os.path.join(
-            os.path.dirname(__file__),
-            "backends",
-            "nvidia",
-            "bin",
-            self.binary,
-        ))
+        self.default_path = os.path.join(os.path.dirname(__file__), "backends", "nvidia", "bin", binary)
+        super().__init__(f"TRITON_{binary.upper()}_PATH")
+
+    def get(self) -> NvidiaTool:
+        return self.transform(getenv(self.key))
 
     def transform(self, path: str) -> NvidiaTool:
-        paths = [
-            path,
-            # We still add default as fallback in case the pointed binary isn't
-            # accessible.
-            self.default(),
-        ]
+        # We still add default as fallback in case the pointed binary isn't
+        # accessible.
+        if path is not None:
+            paths = [path, self.default_path]
+        else:
+            paths = [self.default_path]
+
         for path in paths:
-            if not path or not os.access(path, os.X_OK):
-                continue
             if tool := NvidiaTool.from_path(path):
                 return tool
 
         raise RuntimeError(f"Cannot find {self.binary}")
 
-    def from_env(self, val: str) -> str:
-        return val
-
 
 class env_intel_tool(env_base[str, IntelTool]):
 
     def __init__(self, binary: str) -> None:
         binary += sysconfig.get_config_var("EXE")
         self.binary = binary
-        super().__init__(f"TRITON_{binary.upper().replace('-', '_')}_PATH", lambda: os.path.join(
-            os.path.dirname(__file__),
-            "backends",
-            "intel",
-            "bin",
-            self.binary,
-        ))
+        self.default_path = os.path.join(os.path.dirname(__file__), "backends", "intel", "bin", binary)
+        super().__init__(f"TRITON_{binary.upper()}_PATH")
+
+    def get(self) -> IntelTool:
+        return self.transform(getenv(self.key))
 
     def transform(self, path: str) -> IntelTool:
-        paths = [
-            path,
-            # We still add default as fallback in case the pointed binary isn't
-            # accessible.
-            self.default(),
-            shutil.which(self.binary) or "",
-        ]
+        # We still add default as fallback in case the pointed binary isn't
+        # accessible.
+        if path is not None:
+            paths = [path, self.default_path]
+        else:
+            paths = [self.default_path]
+        if shutil_path := shutil.which(self.binary):
+            paths += [shutil_path]
         for path in paths:
-            if not path or not os.access(path, os.X_OK):
-                continue
             if tool := IntelTool.from_path(path):
                 return tool
 
         raise RuntimeError(f"Cannot find {self.binary}")
 
-    def from_env(self, val: str) -> str:
-        return val
-
 
 # Separate classes so that types are correct
-class env_opt_str(env_opt_base[str, str], env_str):
-    pass
+class env_opt_str(env_base[Optional[str], Optional[str]]):
 
+    def get(self) -> Optional[str]:
+        return getenv(self.key)
 
-class env_opt_bool(env_opt_base[bool, bool], env_bool):
-    pass
+
+class env_opt_bool(env_base):
+
+    def get(self) -> Optional[str]:
+        return getenv_bool(self.key, None)
 
 
 @dataclass(frozen=True)
@@ -344,7 +342,7 @@ def reset(self: knobs_type) -> knobs_type:
     @contextmanager
     def scope(self) -> Generator[None, None, None]:
         try:
-            initial_env = {knob.key: knob.env_val for knob in self.knob_descriptors.values()}
+            initial_env = {knob.key: getenv(knob.key) for knob in self.knob_descriptors.values()}
             orig = dict(self.__dict__)
             yield
         finally:
@@ -389,11 +387,11 @@ class redis_knobs(base_knobs):
 
 
 class cache_knobs(base_knobs):
-    home_dir: env_str = env_str("TRITON_HOME", lambda: os.path.expanduser("~/"))
+    home_dir: env_str = env_str("TRITON_HOME", os.path.expanduser("~/"))
 
-    dump_dir: env_str = env_str("TRITON_DUMP_DIR", lambda: cache.get_triton_dir("dump"))
-    override_dir: env_str = env_str("TRITON_OVERRIDE_DIR", lambda: cache.get_triton_dir("override"))
-    dir: env_str = env_str("TRITON_CACHE_DIR", lambda: cache.get_triton_dir("cache"))
+    dump_dir = env_str_callable_default("TRITON_DUMP_DIR", lambda: cache.get_triton_dir("dump"))
+    override_dir = env_str_callable_default("TRITON_OVERRIDE_DIR", lambda: cache.get_triton_dir("override"))
+    dir = env_str_callable_default("TRITON_CACHE_DIR", lambda: cache.get_triton_dir("cache"))
 
     manager_class: env_class[CacheManager] = env_class("TRITON_CACHE_MANAGER", "CacheManager")
     remote_manager_class: env_class[RemoteCacheBackend] = env_class("TRITON_REMOTE_CACHE_BACKEND", "RemoteCacheBackend")
