[python] Allow backend-provided runtime host compiler flags (#7659)

Modifies the `compile_module_from_src` function (and `_build`) to allow
third party plugins to provide arbitrary cc compiler flags when
compiling modules from source at runtime. This is useful for backends
like Intel (`-fsycl`) or for the CPU project (`-fopenmp`) and also for
supporting dynamic lookup for the Python library when running on macOS.

Author: alexbaden

python/triton/runtime/build.py

@@ -16,8 +16,8 @@
 from .. import knobs
 
 
-def _build(name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str],
-           libraries: list[str]) -> str:
+def _build(name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str], libraries: list[str],
+           ccflags: 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')
@@ -48,6 +48,7 @@ def _build(name: str, src: str, srcdir: str, library_dirs: list[str], include_di
     cc_cmd += [f'-l{lib}' for lib in libraries]
     cc_cmd += [f"-L{dir}" for dir in library_dirs]
     cc_cmd += [f"-I{dir}" for dir in include_dirs if dir is not None]
+    cc_cmd.extend(ccflags)
     subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)
     return so
 
@@ -68,7 +69,8 @@ def _load_module_from_path(name: str, path: str) -> ModuleType:
 
 
 def compile_module_from_src(src: str, name: str, library_dirs: list[str] | None = None,
-                            include_dirs: list[str] | None = None, libraries: list[str] | None = None) -> ModuleType:
+                            include_dirs: list[str] | None = None, libraries: list[str] | None = None,
+                            ccflags: list[str] | None = None) -> ModuleType:
     key = hashlib.sha256((src + platform_key()).encode("utf-8")).hexdigest()
     cache = get_cache_manager(key)
     suffix = sysconfig.get_config_var("EXT_SUFFIX")
@@ -85,7 +87,7 @@ def compile_module_from_src(src: str, name: str, library_dirs: list[str] | None
         src_path = os.path.join(tmpdir, name + ".c")
         with open(src_path, "w") as f:
             f.write(src)
-        so = _build(name, src_path, tmpdir, library_dirs or [], include_dirs or [], libraries or [])
+        so = _build(name, src_path, tmpdir, library_dirs or [], include_dirs or [], libraries or [], ccflags or [])
         with open(so, "rb") as f:
             cache_path = cache.put(f.read(), f"{name}{suffix}", binary=True)