Add LLVM_TOOLS_INSTALL_DIR to configure

Author: xhochy

Tools/jit/_llvm.py

@@ -71,10 +71,16 @@ async def _get_brew_llvm_prefix(llvm_version: str, *, echo: bool = False) -> str
 
 
 @_async_cache
-async def _find_tool(tool: str, llvm_version: str, *, echo: bool = False) -> str | None:
+async def _find_tool(
+    tool: str,
+    llvm_version: str,
+    llvm_tools_install_dir: str | None,
+    *,
+    echo: bool = False,
+) -> str | None:
     # Explicitly defined LLVM installation location
-    if llvm_tools_dir := os.getenv("LLVM_TOOLS_INSTALL_DIR"):
-        path = os.path.join(llvm_tools_dir, tool)
+    if llvm_tools_install_dir:
+        path = os.path.join(llvm_tools_install_dir, tool)
         if await _check_tool_version(path, llvm_version, echo=echo):
             return path
     # Unversioned executables:
@@ -110,10 +116,11 @@ async def maybe_run(
     args: typing.Iterable[str],
     echo: bool = False,
     llvm_version: str = _LLVM_VERSION,
+    llvm_tools_install_dir: str | None = None,
 ) -> str | None:
     """Run an LLVM tool if it can be found. Otherwise, return None."""
 
-    path = await _find_tool(tool, llvm_version, echo=echo)
+    path = await _find_tool(tool, llvm_version, llvm_tools_install_dir, echo=echo)
     return path and await _run(path, args, echo=echo)
 
 
@@ -122,10 +129,17 @@ async def run(
     args: typing.Iterable[str],
     echo: bool = False,
     llvm_version: str = _LLVM_VERSION,
+    llvm_tools_install_dir: str | None = None,
 ) -> str:
     """Run an LLVM tool if it can be found. Otherwise, raise RuntimeError."""
 
-    output = await maybe_run(tool, args, echo=echo, llvm_version=llvm_version)
+    output = await maybe_run(
+        tool,
+        args,
+        echo=echo,
+        llvm_version=llvm_version,
+        llvm_tools_install_dir=llvm_tools_install_dir,
+    )
     if output is None:
         raise RuntimeError(f"Can't find {tool}-{llvm_version}!")
     return output

Tools/jit/_targets.py

@@ -51,6 +51,7 @@ class _Target(typing.Generic[_S, _R]):
     verbose: bool = False
     cflags: str = ""
     llvm_version: str = _llvm._LLVM_VERSION
+    llvm_tools_install_dir: str | None = None
     known_symbols: dict[str, int] = dataclasses.field(default_factory=dict)
     pyconfig_dir: pathlib.Path = pathlib.Path.cwd().resolve()
 
@@ -83,7 +84,11 @@ async def _parse(self, path: pathlib.Path) -> _stencils.StencilGroup:
         group = _stencils.StencilGroup()
         args = ["--disassemble", "--reloc", f"{path}"]
         output = await _llvm.maybe_run(
-            "llvm-objdump", args, echo=self.verbose, llvm_version=self.llvm_version
+            "llvm-objdump",
+            args,
+            echo=self.verbose,
+            llvm_version=self.llvm_version,
+            llvm_tools_install_dir=self.llvm_tools_install_dir,
         )
         if output is not None:
             # Make sure that full paths don't leak out (for reproducibility):
@@ -103,7 +108,11 @@ async def _parse(self, path: pathlib.Path) -> _stencils.StencilGroup:
             f"{path}",
         ]
         output = await _llvm.run(
-            "llvm-readobj", args, echo=self.verbose, llvm_version=self.llvm_version
+            "llvm-readobj",
+            args,
+            echo=self.verbose,
+            llvm_version=self.llvm_version,
+            llvm_tools_install_dir=self.llvm_tools_install_dir,
         )
         # --elf-output-style=JSON is only *slightly* broken on Mach-O...
         output = output.replace("PrivateExtern\n", "\n")
@@ -177,14 +186,22 @@ async def _compile(
             *shlex.split(self.cflags),
         ]
         await _llvm.run(
-            "clang", args_s, echo=self.verbose, llvm_version=self.llvm_version
+            "clang",
+            args_s,
+            echo=self.verbose,
+            llvm_version=self.llvm_version,
+            llvm_tools_install_dir=self.llvm_tools_install_dir,
         )
         self.optimizer(
             s, label_prefix=self.label_prefix, symbol_prefix=self.symbol_prefix
         ).run()
         args_o = [f"--target={self.triple}", "-c", "-o", f"{o}", f"{s}"]
         await _llvm.run(
-            "clang", args_o, echo=self.verbose, llvm_version=self.llvm_version
+            "clang",
+            args_o,
+            echo=self.verbose,
+            llvm_version=self.llvm_version,
+            llvm_tools_install_dir=self.llvm_tools_install_dir,
         )
         return await self._parse(o)
 

Tools/jit/build.py

@@ -43,6 +43,9 @@
         "--cflags", help="additional flags to pass to the compiler", default=""
     )
     parser.add_argument("--llvm-version", help="LLVM version to use")
+    parser.add_argument(
+        "--llvm-tools-install-dir", help="Installation location of LLVM tools"
+    )
     args = parser.parse_args()
     for target in args.target:
         target.debug = args.debug
@@ -52,6 +55,8 @@
         target.pyconfig_dir = args.pyconfig_dir
         if args.llvm_version:
             target.llvm_version = args.llvm_version
+        if args.llvm_tools_install_dir:
+            target.llvm_tools_install_dir = args.llvm_tools_install_dir
         target.build(
             comment=comment,
             force=args.force,

configure

@@ -2804,7 +2804,7 @@ AS_VAR_IF([jit_flags],
           [],
           [AS_VAR_APPEND([CFLAGS_NODIST], [" $jit_flags"])
            AS_VAR_SET([REGEN_JIT_COMMAND],
-                      ["\$(PYTHON_FOR_REGEN) \$(srcdir)/Tools/jit/build.py ${ARCH_TRIPLES:-$host} --output-dir . --pyconfig-dir . --cflags=\"$CFLAGS_JIT\" --llvm-version=\"$LLVM_VERSION\""])
+                      ["\$(PYTHON_FOR_REGEN) \$(srcdir)/Tools/jit/build.py ${ARCH_TRIPLES:-$host} --output-dir . --pyconfig-dir . --cflags=\"$CFLAGS_JIT\" --llvm-version=\"$LLVM_VERSION\" --llvm-tools-install-dir=\"$LLVM_TOOLS_INSTALL_DIR\""])
            AS_VAR_IF([Py_DEBUG],
                      [true],
                      [AS_VAR_APPEND([REGEN_JIT_COMMAND], [" --debug"])],