asan: Preload asan runtime into python interpreter (#4046)

Same as llvm/llvm-project#123303 upstream

Author: mgehre-amd

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

@@ -37,10 +37,23 @@
 # test_exec_root: The root path where tests should be run.
 config.test_exec_root = os.path.join(config.torch_mlir_obj_root, "test")
 
+if "asan" in config.available_features and "Linux" in config.host_os:
+    _asan_rt = (
+        subprocess.check_output(
+            [
+                config.host_cxx.strip(),
+                f"-print-file-name=libclang_rt.asan-{config.host_arch}.so",
+            ]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+    config.python_executable = f"env LD_PRELOAD={_asan_rt} {config.python_executable}"
+    config.environment["ASAN_OPTIONS"] = "detect_leaks=0"
 # 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/python/test/lit.site.cfg.py.in

@@ -19,6 +19,10 @@ config.llvm_lib_dir = "@LLVM_LIBRARY_DIR@"
 config.llvm_shlib_dir = "@SHLIBDIR@"
 config.llvm_shlib_ext = "@SHLIBEXT@"
 config.llvm_exe_ext = "@EXEEXT@"
+config.host_os = "@HOST_OS@"
+config.host_cxx = "@HOST_CXX@"
+config.host_arch = "@HOST_ARCH@"
+config.llvm_use_sanitizer = "@LLVM_USE_SANITIZER@"
 config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
 config.python_executable = "@Python3_EXECUTABLE@"
 config.gold_executable = "@GOLD_EXECUTABLE@"

projects/pt1/test/lit.cfg.py

@@ -66,10 +66,23 @@
     "PATH", os.path.join(config.llvm_build_dir, "bin"), append_path=True
 )
 
+if "asan" in config.available_features and "Linux" in config.host_os:
+    _asan_rt = (
+        subprocess.check_output(
+            [
+                config.host_cxx.strip(),
+                f"-print-file-name=libclang_rt.asan-{config.host_arch}.so",
+            ]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+    config.python_executable = f"env LD_PRELOAD={_asan_rt} {config.python_executable}"
+    config.environment["ASAN_OPTIONS"] = "detect_leaks=0"
 # 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,9 @@ 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_use_sanitizer = "@LLVM_USE_SANITIZER@"
 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,23 @@
     "PATH", os.path.join(config.llvm_build_dir, "bin"), append_path=True
 )
 
+if "asan" in config.available_features and "Linux" in config.host_os:
+    _asan_rt = (
+        subprocess.check_output(
+            [
+                config.host_cxx.strip(),
+                f"-print-file-name=libclang_rt.asan-{config.host_arch}.so",
+            ]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+    config.python_executable = f"env LD_PRELOAD={_asan_rt} {config.python_executable}"
+    config.environment["ASAN_OPTIONS"] = "detect_leaks=0"
 # 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,9 @@ 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_use_sanitizer = "@LLVM_USE_SANITIZER@"
 config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
 config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"