diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py index c28623123d999..cbdac91d5ba7d 100644 --- a/mlir/test/lit.cfg.py +++ b/mlir/test/lit.cfg.py @@ -82,17 +82,32 @@ def add_runtime(name): # available. This is darwin specific since it's currently only needed on darwin. # Stolen from llvm/test/lit.cfg.py with a few modifications def get_asan_rtlib(): - if not "asan" in config.available_features or not "Darwin" in config.host_os: + if not "asan" in config.available_features: return "" - # Find the asan rt lib - resource_dir = ( - subprocess.check_output([config.host_cc.strip(), "-print-resource-dir"]) - .decode("utf-8") - .strip() - ) - return os.path.join( - resource_dir, "lib", "darwin", "libclang_rt.asan_osx_dynamic.dylib" - ) + + if "Darwin" in config.host_os: + # Find the asan rt lib + resource_dir = ( + subprocess.check_output([config.host_cc.strip(), "-print-resource-dir"]) + .decode("utf-8") + .strip() + ) + return os.path.join( + resource_dir, "lib", "darwin", "libclang_rt.asan_osx_dynamic.dylib" + ) + if "Linux" in config.host_os: + return ( + subprocess.check_output( + [ + config.host_cxx.strip(), + f"-print-file-name=libclang_rt.asan-{config.host_arch}.so", + ] + ) + .decode("utf-8") + .strip() + ) + + return "" # On macOS, we can't do the DYLD_INSERT_LIBRARIES trick with a shim python @@ -247,7 +262,9 @@ def find_real_python_interpreter(): # TODO: detect Windows situation (or mark these tests as unsupported on these platforms). if "asan" in config.available_features: if "Linux" in config.host_os: - python_executable = f"LD_PRELOAD=$({config.host_cxx} -print-file-name=libclang_rt.asan-{config.host_arch}.so) {config.python_executable}" + python_executable = ( + f"env LD_PRELOAD={get_asan_rtlib()} {config.python_executable}" + ) if "Darwin" in config.host_os: # Ensure we use a non-shim Python executable, for the `DYLD_INSERT_LIBRARIES` # env variable to take effect