Reapply "fix vector.splat -> vector.broadcast" (#170)

This reverts commit b6729e9.

Author: makslevental

mlir/extras/runtime/refbackend.py

@@ -41,6 +41,8 @@
 CONSUME_RETURN_CALLBACK_ATTR = "refbackend_consume_return_callback"
 refback_cb_attr = CONSUME_RETURN_CALLBACK_ATTR
 
+_exec_engine_shared_libs = []
+
 if ASYNC_RUNTIME_LIB_PATH := os.getenv("ASYNC_RUNTIME_LIB_PATH"):
     ASYNC_RUNTIME_LIB_PATH = Path(ASYNC_RUNTIME_LIB_PATH)
 else:
@@ -49,6 +51,9 @@
         / f"{shlib_prefix()}mlir_async_runtime.{shlib_ext()}"
     )
 
+if ASYNC_RUNTIME_LIB_PATH.exists():
+    _exec_engine_shared_libs.append(ASYNC_RUNTIME_LIB_PATH)
+
 if C_RUNNER_UTILS_LIB_PATH := os.getenv("C_RUNNER_UTILS_LIB_PATH"):
     C_RUNNER_UTILS_LIB_PATH = Path(C_RUNNER_UTILS_LIB_PATH)
 else:
@@ -57,6 +62,9 @@
         / f"{shlib_prefix()}mlir_c_runner_utils.{shlib_ext()}"
     )
 
+if C_RUNNER_UTILS_LIB_PATH.exists():
+    _exec_engine_shared_libs.append(C_RUNNER_UTILS_LIB_PATH)
+
 if RUNNER_UTILS_LIB_PATH := os.getenv("RUNNER_UTILS_LIB_PATH"):
     RUNNER_UTILS_LIB_PATH = Path(RUNNER_UTILS_LIB_PATH)
 else:
@@ -65,6 +73,20 @@
         / f"{shlib_prefix()}mlir_runner_utils.{shlib_ext()}"
     )
 
+if RUNNER_UTILS_LIB_PATH.exists():
+    _exec_engine_shared_libs.append(RUNNER_UTILS_LIB_PATH)
+
+if CUDA_RUNTIME_LIB_PATH := os.getenv("CUDA_RUNTIME_LIB_PATH"):
+    CUDA_RUNTIME_LIB_PATH = Path(CUDA_RUNTIME_LIB_PATH)
+else:
+    CUDA_RUNTIME_LIB_PATH = (
+        Path(_mlir_libs.__file__).parent
+        / f"{shlib_prefix()}mlir_cuda_runtime.{shlib_ext()}"
+    )
+
+if CUDA_RUNTIME_LIB_PATH.exists():
+    _exec_engine_shared_libs.append(CUDA_RUNTIME_LIB_PATH)
+
 
 def get_ctype_func(mlir_ret_types):
     ctypes_arg = [None]
@@ -202,14 +224,10 @@ def __init__(
         shared_lib_paths=None,
     ):
         if shared_lib_paths is None:
-            shared_lib_paths = []
+            shared_lib_paths = set()
         if platform.system() != "Windows":
-            shared_lib_paths += [
-                ASYNC_RUNTIME_LIB_PATH,
-                C_RUNNER_UTILS_LIB_PATH,
-                RUNNER_UTILS_LIB_PATH,
-            ]
-        self.shared_lib_paths = shared_lib_paths
+            shared_lib_paths |= set(_exec_engine_shared_libs)
+        self.shared_lib_paths = list(shared_lib_paths)
         self.return_func_types = None
         self.return_func_name = None
 

tests/test_nvgpu_nvvm.py

@@ -27,7 +27,7 @@
 from mlir.extras.dialects.ext.gpu import smem_space
 from mlir.extras.dialects.ext.llvm import llvm_ptr_t
 from mlir.extras.runtime.passes import Pipeline, run_pipeline
-from mlir.extras.runtime.refbackend import LLVMJITBackend
+from mlir.extras.runtime.refbackend import LLVMJITBackend, CUDA_RUNTIME_LIB_PATH
 
 # noinspection PyUnresolvedReferences
 from mlir.extras.testing import (
@@ -202,7 +202,8 @@ def payload():
         compute_linspace_val.emit()
 
         @func
-        def printMemrefF32(x: T.memref(T.f32())): ...
+        def printMemrefF32(x: T.memref(T.f32())):
+            ...
 
         printMemrefF32_.append(printMemrefF32)
 
@@ -413,10 +414,11 @@ def main(module: any_op_t()):
     # CHECK:    }
     # CHECK:  }
 
-    filecheck_with_comments(mod)
+    mod.operation.verify()
 
+    if CUDA_RUNTIME_LIB_PATH.exists():
+        filecheck_with_comments(mod)
 
-CUDA_RUNTIME_LIB_PATH = Path(_mlir_libs.__file__).parent / f"libmlir_cuda_runtime.so"
 
 NVIDIA_GPU = False
 try:
@@ -553,7 +555,8 @@ def payload():
         compute_linspace_val.emit()
 
         @func
-        def printMemrefF32(x: T.memref(T.f32())): ...
+        def printMemrefF32(x: T.memref(T.f32())):
+            ...
 
         printMemrefF32_.append(printMemrefF32)