Merge remote-tracking branch 'origin/feature/backport_ea1_ops' into b…

…ump_to_e9ed4af9

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

@@ -18,6 +18,24 @@
 
 # Configuration file for the 'lit' test runner.
 
+
+# Find path to the ASan runtime required for the Python interpreter.
+def find_asan_runtime():
+    if not "asan" in config.available_features or not "Linux" in config.host_os:
+        return ""
+    # Find the asan rt lib
+    return (
+        subprocess.check_output(
+            [
+                config.host_cxx.strip(),
+                f"-print-file-name=libclang_rt.asan-{config.host_arch}.so",
+            ]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+
+
 # name: The name of this test suite.
 config.name = "TORCH_MLIR_PYTHON"
 
@@ -37,10 +55,15 @@
 # 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 "asan" in config.available_features and "Linux" in config.host_os:
+    _asan_rt = find_asan_runtime()
+    config.python_executable = f"env LD_PRELOAD={_asan_rt} {config.python_executable}"
 # 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

@@ -18,6 +18,24 @@
 
 # Configuration file for the 'lit' test runner.
 
+
+# Find path to the ASan runtime required for the Python interpreter.
+def find_asan_runtime():
+    if not "asan" in config.available_features or not "Linux" in config.host_os:
+        return ""
+    # Find the asan rt lib
+    return (
+        subprocess.check_output(
+            [
+                config.host_cxx.strip(),
+                f"-print-file-name=libclang_rt.asan-{config.host_arch}.so",
+            ]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+
+
 # name: The name of this test suite.
 config.name = "TORCH_MLIR_PT1"
 
@@ -66,10 +84,15 @@
     "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 "asan" in config.available_features and "Linux" in config.host_os:
+    _asan_rt = find_asan_runtime()
+    config.python_executable = f"env LD_PRELOAD={_asan_rt} {config.python_executable}"
 # 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

@@ -18,6 +18,24 @@
 
 # Configuration file for the 'lit' test runner.
 
+
+# Find path to the ASan runtime required for the Python interpreter.
+def find_asan_runtime():
+    if not "asan" in config.available_features or not "Linux" in config.host_os:
+        return ""
+    # Find the asan rt lib
+    return (
+        subprocess.check_output(
+            [
+                config.host_cxx.strip(),
+                f"-print-file-name=libclang_rt.asan-{config.host_arch}.so",
+            ]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+
+
 # name: The name of this test suite.
 config.name = "TORCH_MLIR"
 
@@ -66,10 +84,15 @@
     "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 "asan" in config.available_features and "Linux" in config.host_os:
+    _asan_rt = find_asan_runtime()
+    config.python_executable = f"env LD_PRELOAD={_asan_rt} {config.python_executable}"
 # 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@"