asan: replace used python in various lit.cfg's with shim script

projects/pt1/python/test/lit.cfg.py

@@ -37,10 +37,26 @@
 # test_exec_root: The root path where tests should be run.
 config.test_exec_root = os.path.join(config.torch_mlir_obj_root, "test")
 
+# Python configuration with sanitizer requires some magic preloading. This will only work on clang/linux.
+# TODO: detect Darwin/Windows situation (or mark these tests as unsupported on these platforms).
+if os.environ.get("IS_ASAN") and "Linux" in config.host_os:
+    # Write shim script that preloads the necessary shared object for ASAN tests. Fallback to such script for two reasons:
+    #
+    # (1) Provide full support for LLVM's test utils like `not`, which are prepended to the original statement containing the `LD_PRELOAD` env definition.
+    #     Having environment definitions in the middle of a command line is syntactically illegal.
+    # (2) Mitigate issues with LIT's internal shell that puts single quotes around the environment definition,
+    #     which leads to malformed command lines:
+    #     `LD_PRELOAD=$(/usr/bin/clang++-17' '-print-file-name=libclang_rt.asan-x86_64.so)' python (...)`
+    with open("python-asan-shim", "w") as file:
+        file.write(
+            f"#!/usr/bin/env bash\nLD_PRELOAD=$({config.host_cxx} -print-file-name=libclang_rt.asan-{config.host_arch}.so) {config.python_executable} $@\n"
+        )
+    os.chmod(os.path.abspath("python-asan-shim"), 0o700)
+    config.python_executable = os.path.abspath("python-asan-shim")
 # On Windows the path to python could contains spaces in which case it needs to
 # be provided in quotes.  This is the equivalent of how %python is setup in
 # llvm/utils/lit/lit/llvm/config.py.
-if "Windows" in config.host_os:
+elif "Windows" in config.host_os:
     config.python_executable = '"%s"' % (config.python_executable)
 
 config.substitutions.append(("%PATH%", config.environment["PATH"]))

projects/pt1/test/lit.cfg.py

@@ -66,10 +66,26 @@
     "PATH", os.path.join(config.llvm_build_dir, "bin"), append_path=True
 )
 
+# Python configuration with sanitizer requires some magic preloading. This will only work on clang/linux.
+# TODO: detect Darwin/Windows situation (or mark these tests as unsupported on these platforms).
+if os.environ.get("IS_ASAN") and "Linux" in config.host_os:
+    # Write shim script that preloads the necessary shared object for ASAN tests. Fallback to such script for two reasons:
+    #
+    # (1) Provide full support for LLVM's test utils like `not`, which are prepended to the original statement containing the `LD_PRELOAD` env definition.
+    #     Having environment definitions in the middle of a command line is syntactically illegal.
+    # (2) Mitigate issues with LIT's internal shell that puts single quotes around the environment definition,
+    #     which leads to malformed command lines:
+    #     `LD_PRELOAD=$(/usr/bin/clang++-17' '-print-file-name=libclang_rt.asan-x86_64.so)' python (...)`
+    with open("python-asan-shim", "w") as file:
+        file.write(
+            f"#!/usr/bin/env bash\nLD_PRELOAD=$({config.host_cxx} -print-file-name=libclang_rt.asan-{config.host_arch}.so) {config.python_executable} $@\n"
+        )
+    os.chmod(os.path.abspath("python-asan-shim"), 0o700)
+    config.python_executable = os.path.abspath("python-asan-shim")
 # On Windows the path to python could contains spaces in which case it needs to
 # be provided in quotes.  This is the equivalent of how %python is setup in
 # llvm/utils/lit/lit/llvm/config.py.
-if "Windows" in config.host_os:
+elif "Windows" in config.host_os:
     config.python_executable = '"%s"' % (config.python_executable)
 
 tool_dirs = [

projects/pt1/test/lit.site.cfg.py.in

@@ -6,6 +6,8 @@ config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@
 config.torch_mlir_obj_root = "@TORCH_MLIR_BINARY_DIR@"
 config.torch_mlir_python_packages_dir = "@TORCH_MLIR_PYTHON_PACKAGES_DIR@"
 config.host_os = "@HOST_OS@"
+config.host_cxx = "@HOST_CXX@"
+config.host_arch = "@HOST_ARCH@"
 config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
 config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"

test/lit.cfg.py

@@ -66,10 +66,26 @@
     "PATH", os.path.join(config.llvm_build_dir, "bin"), append_path=True
 )
 
+# Python configuration with sanitizer requires some magic preloading. This will only work on clang/linux.
+# TODO: detect Darwin/Windows situation (or mark these tests as unsupported on these platforms).
+if os.environ.get("IS_ASAN") and "Linux" in config.host_os:
+    # Write shim script that preloads the necessary shared object for ASAN tests. Fallback to such script for two reasons:
+    #
+    # (1) Provide full support for LLVM's test utils like `not`, which are prepended to the original statement containing the `LD_PRELOAD` env definition.
+    #     Having environment definitions in the middle of a command line is syntactically illegal.
+    # (2) Mitigate issues with LIT's internal shell that puts single quotes around the environment definition,
+    #     which leads to malformed command lines:
+    #     `LD_PRELOAD=$(/usr/bin/clang++-17' '-print-file-name=libclang_rt.asan-x86_64.so)' python (...)`
+    with open("python-asan-shim", "w") as file:
+        file.write(
+            f"#!/usr/bin/env bash\nLD_PRELOAD=$({config.host_cxx} -print-file-name=libclang_rt.asan-{config.host_arch}.so) {config.python_executable} $@\n"
+        )
+    os.chmod(os.path.abspath("python-asan-shim"), 0o700)
+    config.python_executable = os.path.abspath("python-asan-shim")
 # On Windows the path to python could contains spaces in which case it needs to
 # be provided in quotes.  This is the equivalent of how %python is setup in
 # llvm/utils/lit/lit/llvm/config.py.
-if "Windows" in config.host_os:
+elif "Windows" in config.host_os:
     config.python_executable = '"%s"' % (config.python_executable)
 
 tool_dirs = [

test/lit.site.cfg.py.in

@@ -7,6 +7,8 @@ config.torch_mlir_obj_root = "@TORCH_MLIR_BINARY_DIR@"
 config.torch_mlir_python_packages_dir = "@TORCH_MLIR_PYTHON_PACKAGES_DIR@"
 config.torch_mlir_enable_refbackend = @TORCH_MLIR_ENABLE_REFBACKEND@
 config.host_os = "@HOST_OS@"
+config.host_cxx = "@HOST_CXX@"
+config.host_arch = "@HOST_ARCH@"
 config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
 config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"