[python][knobs] Implement _build hook (triton-lang#6710)

Right now there's a TODO in `_build` to expand support for compilers.
With `knobs` a natural way to expand support is with a custom build
implementation, hence this PR.

Author: danzimm

python/triton/knobs.py

@@ -272,13 +272,22 @@ def scope(self) -> Generator[None, None, None]:
             self.__dict__.update(orig)
 
 
+class BuildImpl(Protocol):
+
+    def __call__(self, name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str],
+                 libraries: list[str], /) -> str:
+        ...
+
+
 class build_knobs(base_knobs):
     """Configuration controlling how the native compiler is invoked"""
     cc: env_opt_str = env_opt_str("CC")
 
     cudacrt_path: env_opt_str = env_opt_str("TRITON_CUDACRT_PATH")
     cudart_path: env_opt_str = env_opt_str("TRITON_CUDART_PATH")
 
+    impl: Optional[BuildImpl] = None
+
     @property
     def backend_dirs(self) -> set[str]:
         return {path for path in (self.cudacrt_path, self.cudart_path) if path is not None}

python/triton/runtime/build.py

@@ -6,23 +6,26 @@
 from .. import knobs
 
 
-def _build(name, src, srcdir, library_dirs, include_dirs, libraries):
+def _build(name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str],
+           libraries: list[str]) -> str:
+    if impl := knobs.build.impl:
+        return impl(name, src, srcdir, library_dirs, include_dirs, libraries)
     suffix = sysconfig.get_config_var('EXT_SUFFIX')
     so = os.path.join(srcdir, '{name}{suffix}'.format(name=name, suffix=suffix))
     # try to avoid setuptools if possible
     cc = os.environ.get("CC")
     if cc is None:
-        # TODO: support more things here.
         clang = shutil.which("clang")
         gcc = shutil.which("gcc")
         cc = gcc if gcc is not None else clang
         if cc is None:
-            raise RuntimeError("Failed to find C compiler. Please specify via CC environment variable.")
+            raise RuntimeError(
+                "Failed to find C compiler. Please specify via CC environment variable or set triton.knobs.build.impl.")
     # This function was renamed and made public in Python 3.10
     if hasattr(sysconfig, 'get_default_scheme'):
         scheme = sysconfig.get_default_scheme()
     else:
-        scheme = sysconfig._get_default_scheme()
+        scheme = sysconfig._get_default_scheme()  # type: ignore
     # 'posix_local' is a custom scheme on Debian. However, starting Python 3.10, the default install
     # path changes to include 'local'. This change is required to use triton with system-wide python.
     if scheme == 'posix_local':