[MLIR][mlir-spirv-cpu-runner] A SPIR-V cpu runner prototype

[georgemitenkov]

Summary

This patch introduces a SPIR-V runner. The aim is to run a gpu
kernel on a CPU via GPU -> SPIRV -> LLVM conversions. This is a first
prototype, so more features will be added in due time.

• Overview

The runner follows similar flow as the other runners in-tree. However,
having converted the kernel to SPIR-V, we encode the bind attributes of
global variables that represent kernel arguments. Then SPIR-V module is
converted to LLVM. On the host side, we emulate passing the data to device
by creating in main module globals with the same symbolic name as in kernel
module. These global variables are later linked with ones from the nested
module. We copy data from kernel arguments to globals, call the kernel
function from nested module and then copy the data back.

• Current state

At the moment, the runner is capable of running 2 modules, nested one in
another. The kernel module must contain exactly one kernel function. Also,
the runner supports rank 1 integer memref types as arguments (to be scaled).

• Enhancement of JitRunner and ExecutionEngine

To translate nested modules to LLVM IR, JitRunner and ExecutionEngine were
altered to take an optional (default to nullptr) function reference that
is a custom LLVM IR module builder. This allows to customize LLVM IR module
creation from MLIR modules.

[rriddle]

Can you split this up? Seems like the passes could be split and reviewed separately from the runner code.

[georgemitenkov]

In D86108#2222554, @rriddle wrote:

Can you split this up? Seems like the passes could be split and reviewed separately from the runner code.

Sure, thanks!

[georgemitenkov]

Splitted the patch.

[georgemitenkov]

Fixed a typo.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

69

Can you document the else case please?

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

231

I'm not convinced this injection is really properly layered: it has some loose contract right now.

[mravishankar]

Thanks George!

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

231

How about the translateModuleToLLVMIR be modified to handle multiple modules. So instead of just taking a single ModuleOp takes an ArrayRef<ModuleOp> as arguments and then links them. This way you wouldnt have to thread through the llvmModuleBuilder callback through the entire stack. The spirv-runner can split the nested module out and pass the outer module and the nested modules as two separate modules that are to be compiled and link. Based on a quick survey of the code here, I think that should work.

@mehdi_amini @ftynse, any comments?

mlir/tools/mlir-spirv-runner/runtime-wrappers.cpp

27 ↗

(On Diff #286159)

THis seems like a test utility function. If so it might be better to move this into a test/ folder. Something like test/mlir-spirv-runner. See similar examples here : https://github.com/llvm/llvm-project/blob/master/mlir/test/mlir-cpu-runner/mlir_test_cblas.cpp

[georgemitenkov]

Moved the library to tests. Documented the else case for llvmModuleBuilder in ExecutionEngine.h.

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

231

I am a bit confused by what you're suggesting here, where would the ArrayRef<ModuleOp> come from? ExecutionEngine operates on a single module right now.

[mravishankar]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

231

I am saying that ExecutionEngine::create can take an argument ArrayRef<ModuleOp> instead of just Module. The translateModuleToLLVMIR can be modified to handle multiple modules that get linked together before a single LLVM module is returned

[georgemitenkov]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

231

I am not sure why passing ArrayRef<ModuleOp> will be better? We then dedicate linking to translateModuleToLLVMIR(). If we use LLVM's Linker (https://github.com/llvm/llvm-project/blob/master/llvm/include/llvm/Linker/Linker.h), we may also want to specify some of the optional parameters?

[georgemitenkov]

FYI, I made a post on the forum so it is easier to discuss this.

[georgemitenkov]

Fixed an error on building SPIR-V runner.

[georgemitenkov]

Renamed the runner to mlir-spirv-cpu-runner to resolve the ambiguity of where the runner is actually executed. Also, rebased on updated pass infrastructure.

[georgemitenkov]

Updated the runner based on the previous patches and added a test for 3D memrefs with floats.

[mravishankar]

Looking at the stack here, I think only the change below this one in the stack (which is good to go) and this would probably need me to submit this. It would be good to commit the other two dependent changes and then I could try these two out locally.

[georgemitenkov]

In D86108#2240916, @mravishankar wrote:

Looking at the stack here, I think only the change below this one in the stack (which is good to go) and this would probably need me to submit this. It would be good to commit the other two dependent changes and then I could try these two out locally.

Yes, only the pass converting gpu.launch_func and host code is needed. I am committing all accepted patches now :)

[georgemitenkov]

Rebase against master.

[mravishankar]

Trying this change out, I needed the following changes to get this to build locally. Can you verify this is needed?

index 560dcae8200..41f5c8a9ce5 100644
--- a/mlir/lib/Conversion/GPUCommon/CMakeLists.txt
+++ b/mlir/lib/Conversion/GPUCommon/CMakeLists.txt
@@ -34,6 +34,8 @@ add_mlir_conversion_library(MLIRGPUToGPURuntimeTransforms
   MLIRIR
   MLIRLLVMIR
   MLIRPass
+  MLIRSPIRV
   MLIRSupport
   MLIRStandardToLLVM
+  MLIRSPIRVToLLVM
 )
diff --git a/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt b/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
index 69080ae66dc..8e0c2a6ad67 100644
--- a/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
+++ b/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
@@ -15,6 +15,7 @@ if (MLIR_SPIRV_CPU_RUNNER_ENABLED)
   get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
   target_link_libraries(mlir-spirv-cpu-runner PRIVATE
+    LLVMLinker
     ${conversion_libs}
     ${dialect_libs}
     MLIRAnalysis

Also there seems to be some failure that is unrelated to this changes, and I am not sure how to run the tests individually. I ran this command locally and I see this error

$ $HOME/llvm/llvm-project/build/bin/mlir-spirv-cpu-runner -e main --entry-point-result=void --shared-libs=$HOME/llvm/llvm-project/build/lib/libmlir_runner_utils.so,$HOME/llvm/llvm-project/build/lib/libmlir_test_spirv_cpu_runner_c_wrappers.so simple_add.mlir 
loc("simple_add.mlir":50:5): error: 'gpu.launch_func' op SPIR-V kernel module '__spv__kernels' is not found
loc("simple_add.mlir":50:5): error: 'gpu.launch_func' op kernel module 'kernels' is undefined

Any ideas about what I might be missing here?

[georgemitenkov]

In D86108#2241210, @mravishankar wrote:

Any ideas about what I might be missing here?

I think this is to do with strings and names in LowerHostToLLVMPass, cause there the build also fails.

[georgemitenkov]

In D86108#2241210, @mravishankar wrote:

Trying this change out, I needed the following changes to get this to build locally. Can you verify this is needed?

I will have a look at this one as well!

[georgemitenkov]

Rebased on the correct version of LowerHostToLLVMPass.

[georgemitenkov]

In D86108#2241210, @mravishankar wrote:

Trying this change out, I needed the following changes to get this to build locally. Can you verify this is needed?

index 560dcae8200..41f5c8a9ce5 100644
--- a/mlir/lib/Conversion/GPUCommon/CMakeLists.txt
+++ b/mlir/lib/Conversion/GPUCommon/CMakeLists.txt
@@ -34,6 +34,8 @@ add_mlir_conversion_library(MLIRGPUToGPURuntimeTransforms
   MLIRIR
   MLIRLLVMIR
   MLIRPass
+  MLIRSPIRV
   MLIRSupport
   MLIRStandardToLLVM
+  MLIRSPIRVToLLVM
 )
diff --git a/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt b/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
index 69080ae66dc..8e0c2a6ad67 100644
--- a/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
+++ b/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
@@ -15,6 +15,7 @@ if (MLIR_SPIRV_CPU_RUNNER_ENABLED)
   get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
   target_link_libraries(mlir-spirv-cpu-runner PRIVATE
+    LLVMLinker
     ${conversion_libs}
     ${dialect_libs}
     MLIRAnalysis

I am building locally without these changes and it is working for me?

Also there seems to be some failure that is unrelated to this changes, and I am not sure how to run the tests individually. I ran this command locally and I see this error

$ $HOME/llvm/llvm-project/build/bin/mlir-spirv-cpu-runner -e main --entry-point-result=void --shared-libs=$HOME/llvm/llvm-project/build/lib/libmlir_runner_utils.so,$HOME/llvm/llvm-project/build/lib/libmlir_test_spirv_cpu_runner_c_wrappers.so simple_add.mlir 
loc("simple_add.mlir":50:5): error: 'gpu.launch_func' op SPIR-V kernel module '__spv__kernels' is not found
loc("simple_add.mlir":50:5): error: 'gpu.launch_func' op kernel module 'kernels' is undefined

Any ideas about what I might be missing here?

This was actually a LowerHostToLLVMPass problem - StringRef objects were destroyed and symbol names were messed up because of that. I fixed that so now everything should be running fine?

[mravishankar]

In D86108#2242654, @georgemitenkov wrote:

In D86108#2241210, @mravishankar wrote:

Trying this change out, I needed the following changes to get this to build locally. Can you verify this is needed?

index 560dcae8200..41f5c8a9ce5 100644
--- a/mlir/lib/Conversion/GPUCommon/CMakeLists.txt
+++ b/mlir/lib/Conversion/GPUCommon/CMakeLists.txt
@@ -34,6 +34,8 @@ add_mlir_conversion_library(MLIRGPUToGPURuntimeTransforms
   MLIRIR
   MLIRLLVMIR
   MLIRPass
+  MLIRSPIRV
   MLIRSupport
   MLIRStandardToLLVM
+  MLIRSPIRVToLLVM
 )
diff --git a/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt b/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
index 69080ae66dc..8e0c2a6ad67 100644
--- a/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
+++ b/mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt
@@ -15,6 +15,7 @@ if (MLIR_SPIRV_CPU_RUNNER_ENABLED)
   get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
   target_link_libraries(mlir-spirv-cpu-runner PRIVATE
+    LLVMLinker
     ${conversion_libs}
     ${dialect_libs}
     MLIRAnalysis

I am building locally without these changes and it is working for me?

I was getting missing symbols error. But I dont see the errors in buildbot either. Thats strange.

Also there seems to be some failure that is unrelated to this changes, and I am not sure how to run the tests individually. I ran this command locally and I see this error

$ $HOME/llvm/llvm-project/build/bin/mlir-spirv-cpu-runner -e main --entry-point-result=void --shared-libs=$HOME/llvm/llvm-project/build/lib/libmlir_runner_utils.so,$HOME/llvm/llvm-project/build/lib/libmlir_test_spirv_cpu_runner_c_wrappers.so simple_add.mlir 
loc("simple_add.mlir":50:5): error: 'gpu.launch_func' op SPIR-V kernel module '__spv__kernels' is not found
loc("simple_add.mlir":50:5): error: 'gpu.launch_func' op kernel module 'kernels' is undefined

Any ideas about what I might be missing here?

This was actually a LowerHostToLLVMPass problem - StringRef objects were destroyed and symbol names were messed up because of that. I fixed that so now everything should be running fine?

Ill try it out again. THanks!

[mravishankar]

Latest patch works, but I still need the fixes to the builds (actually the fixes to GPUTransforms/CMakeLists.txt above is for the patch earlier in the stack).

Thanks George for making this available!

[georgemitenkov]

In D86108#2242754, @mravishankar wrote:

Latest patch works, but I still need the fixes to the builds (actually the fixes to GPUTransforms/CMakeLists.txt above is for the patch earlier in the stack).

Thanks George for making this available!

Great! Let me know if you have any comments.

[ftynse]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

231

I'm fine either way. The injection may be also useful to, e.g., separate the LLVM-only translation from LLVM+OpenMP translation that are now intermixed. It may look strange that a single runner does the module linking, but it removes the need for all other runners to (transitively) depend on the LLVM linker.

On the other hand, the ExecutionEngine already has a linking layer for resolving calls to dynamic libraries, so may be reasonable to have it also link multiple MLIR/LLVM modules together.

The growing complexity of ExecutionEngine API suggests that its users actually want lower-level control over the process, so it may be the right time to reconsider this abstraction. In the longer term, it is preferable to move MLIR-independent parts (better lookup APIs, void ** argument packing, dynamic library loading) into LLVM itself. Ideally, we should be able to produce an LLVM module from MLIR that "just runs" with (augmented) lli instead of having multiple custom runners. For GPU, we may still need some simple shims that split nested modules and link them.

mlir/lib/ExecutionEngine/JitRunner.cpp

304–305

Can we typedef this as TranslatioCallback or something? It's verbose and used too many times.

mlir/test/mlir-spirv-cpu-runner/mlir_test_spirv_cpu_runner_c_wrappers.cpp

15–24 ↗

(On Diff #288389)

I suspect we already have this declared elsewhere, e.g. https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h#L120, could we rather reuse existing things than copy?

mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp

57 ↗

(On Diff #288389)

I'd consider having a check that there is only one nested module.

[georgemitenkov]

Addressed comments:

• Added a typedef for LLVM module builder function in JitRunner.
• Added a check for a single nested module and refactored convertMLIRModule function (now getOps<>() is used over the walk function).
• Removed a duplicate MemRefDescriptor declaration.

[georgemitenkov]

Changed includes to match new file locations in lower-host-to-llvm pass.

[mravishankar]

Thanks George, let me try it. I am trying to add BUILD files for this to make it runnable internally. This refactoring should help. (Will try to land this soon).

mlir/CMakeLists.txt

if ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD)
   set(MLIR_ROCM_CONVERSIONS_ENABLED 1)
 else()
   set(MLIR_ROCM_CONVERSIONS_ENABLED 0)
 endif()
 add_definitions(-DMLIR_ROCM_CONVERSIONS_ENABLED=${MLIR_ROCM_CONVERSIONS_ENABLED})
 
 set(MLIR_CUDA_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir CUDA runner")
 set(MLIR_ROCM_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir ROCm runner")
+set(MLIR_SPIRV_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir SPIR-V runner")
 set(MLIR_VULKAN_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir Vulkan runner")
 
 option(MLIR_INCLUDE_TESTS
        "Generate build targets for the MLIR unit tests."
        ${LLVM_INCLUDE_TESTS})
 
 option(MLIR_INCLUDE_INTEGRATION_TESTS
        "Generate build targets for the MLIR integration tests.")

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 /// be used to invoke the JIT-compiled function.
 class ExecutionEngine {
 public:
   ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener,
                   bool enablePerfNotificationListener);
 
   /// Creates an execution engine for the given module.
   ///
+  /// If `llvmModuleBuilder` is provided, it will be used to create LLVM module
+  /// from the given MLIR module. Otherwise, a default `translateModuleToLLVMIR`
+  /// function will be used to translate MLIR module to LLVM IR.
+  ///
   /// If `transformer` is provided, it will be called on the LLVM module during
   /// JIT-compilation and can be used, e.g., for reporting or optimization.
   ///
   /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for
   /// target code generation.
   ///
   /// If `sharedLibPaths` are provided, the underlying JIT-compilation will
   /// open and link the shared libraries for symbol resolution.
   ///
   /// If `enableObjectCache` is set, the JIT compiler will create one to store
   /// the object generated for the given module.
   ///
   /// If enable `enableGDBNotificationListener` is set, the JIT compiler will
   /// notify the llvm's global GDB notification listener.
   ///
   /// If `enablePerfNotificationListener` is set, the JIT compiler will notify
   /// the llvm's global Perf notification listener.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
   create(ModuleOp m,
+         llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
+                                                          llvm::LLVMContext &)>
+             llvmModuleBuilder = nullptr,
          llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {},
          Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
          ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = true,
          bool enableGDBNotificationListener = true,
          bool enablePerfNotificationListener = true);
 
   /// Looks up a packed-argument function with the given name and returns a
   /// pointer to it.  Propagates errors in case of failure.

mlir/include/mlir/ExecutionEngine/JitRunner.h

 // transformation.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_JITRUNNER_H_
 #define MLIR_SUPPORT_JITRUNNER_H_
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/Module.h"
 
 namespace mlir {
 
 class ModuleOp;
 struct LogicalResult;
 
 // Entry point for all CPU runners. Expects the common argc/argv arguments for
-// standard C++ main functions and an mlirTransformer.
-// The latter is applied after parsing the input into MLIR IR and before passing
-// the MLIR module to the ExecutionEngine.
+// standard C++ main functions, `mlirTransformer` and `llvmModuleBuilder`.
+/// `mlirTransformer` is applied after parsing the input into MLIR IR and before
+/// passing the MLIR module to the ExecutionEngine.
+/// `llvmModuleBuilder` is a custom function that is passed to ExecutionEngine.
+/// It processes MLIR module and creates LLVM IR module.
 int JitRunnerMain(
     int argc, char **argv,
-    llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer);
+    llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer,
+    llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
+                                                     llvm::LLVMContext &)>
+        llvmModuleBuilder = nullptr);
 
 } // namespace mlir
 
 #endif // MLIR_SUPPORT_JITRUNNER_H_

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

: cache(enableObjectCache ? new SimpleObjectCache() : nullptr),
       gdbListener(enableGDBNotificationListener
                       ? llvm::JITEventListener::createGDBRegistrationListener()
                       : nullptr),
       perfListener(enablePerfNotificationListener
                        ? llvm::JITEventListener::createPerfJITEventListener()
                        : nullptr) {}
 
 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
-    ModuleOp m, llvm::function_ref<Error(llvm::Module *)> transformer,
+    ModuleOp m,
+    llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
+                                                     llvm::LLVMContext &)>
+        llvmModuleBuilder,
+    llvm::function_ref<Error(llvm::Module *)> transformer,
     Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel,
     ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache,
     bool enableGDBNotificationListener, bool enablePerfNotificationListener) {
   auto engine = std::make_unique<ExecutionEngine>(
       enableObjectCache, enableGDBNotificationListener,
       enablePerfNotificationListener);
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
-  auto llvmModule = translateModuleToLLVMIR(m, *ctx);
+  auto llvmModule = llvmModuleBuilder ? llvmModuleBuilder(m, *ctx)
+                                      : translateModuleToLLVMIR(m, *ctx);
   if (!llvmModule)
     return make_string_error("could not convert to LLVM IR");
   // FIXME: the triple should be passed to the translation or dialect conversion
   // instead of this.  Currently, the LLVM module created above has no triple
   // associated with it.
   setupTargetTriple(llvmModule.get());
   packFunctionArguments(llvmModule.get());
 

mlir/lib/ExecutionEngine/JitRunner.cpp

 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
-#include "llvm/IR/Module.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include <cstdint>
 #include <numeric>
 
if (flag) {
       optLevel = j;
       break;
     }
   }
   return optLevel;
 }
 
 // JIT-compile the given module and run "entryPoint" with "args" as arguments.
-static Error
-compileAndExecute(Options &options, ModuleOp module, StringRef entryPoint,
-                  std::function<llvm::Error(llvm::Module *)> transformer,
-                  void **args) {
+static Error compileAndExecute(
+    Options &options, ModuleOp module,
+    function_ref<std::unique_ptr<llvm::Module>(ModuleOp, llvm::LLVMContext &)>
+        llvmModuleBuilder,
+    StringRef entryPoint,
+    std::function<llvm::Error(llvm::Module *)> transformer, void **args) {
   Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
   if (auto clOptLevel = getCommandLineOptLevel(options))
     jitCodeGenOptLevel =
         static_cast<llvm::CodeGenOpt::Level>(clOptLevel.getValue());
   SmallVector<StringRef, 4> libs(options.clSharedLibs.begin(),
                                  options.clSharedLibs.end());
-  auto expectedEngine = mlir::ExecutionEngine::create(module, transformer,
-                                                      jitCodeGenOptLevel, libs);
+  auto expectedEngine = mlir::ExecutionEngine::create(
+      module, llvmModuleBuilder, transformer, jitCodeGenOptLevel, libs);
   if (!expectedEngine)
     return expectedEngine.takeError();
 
   auto engine = std::move(*expectedEngine);
   auto expectedFPtr = engine->lookup(entryPoint);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
 
   if (options.dumpObjectFile)
     engine->dumpToObjectFile(options.objectFilename.empty()
                                  ? options.inputFilename + ".o"
                                  : options.objectFilename);
 
   void (*fptr)(void **) = *expectedFPtr;
   (*fptr)(args);
 
   return Error::success();
 }
 
 static Error compileAndExecuteVoidFunction(
-    Options &options, ModuleOp module, StringRef entryPoint,
+    Options &options, ModuleOp module,
+    function_ref<std::unique_ptr<llvm::Module>(ModuleOp, llvm::LLVMContext &)>
+        llvmModuleBuilder,
+    StringRef entryPoint,
     std::function<llvm::Error(llvm::Module *)> transformer) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
   if (!mainFunction || mainFunction.empty())
     return make_string_error("entry point not found");
   void *empty = nullptr;
-  return compileAndExecute(options, module, entryPoint, transformer, &empty);
+  return compileAndExecute(options, module, llvmModuleBuilder, entryPoint,
+                           transformer, &empty);
 }
 
 template <typename Type>
 Error checkCompatibleReturnType(LLVM::LLVMFuncOp mainFunction);
 template <>
 Error checkCompatibleReturnType<int32_t>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getType().getFunctionResultType().isIntegerTy(32))
     return make_string_error("only single llvm.i32 function result supported");
   return Error::success();
 }
 template <>
 Error checkCompatibleReturnType<int64_t>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getType().getFunctionResultType().isIntegerTy(64))
     return make_string_error("only single llvm.i64 function result supported");
   return Error::success();
 }
 template <>
 Error checkCompatibleReturnType<float>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getType().getFunctionResultType().isFloatTy())
     return make_string_error("only single llvm.f32 function result supported");
   return Error::success();
 }
 template <typename Type>
 Error compileAndExecuteSingleReturnFunction(
-    Options &options, ModuleOp module, StringRef entryPoint,
+    Options &options, ModuleOp module,
+    function_ref<std::unique_ptr<llvm::Module>(ModuleOp, llvm::LLVMContext &)>
+        llvmModuleBuilder,
+    StringRef entryPoint,
     std::function<llvm::Error(llvm::Module *)> transformer) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
   if (!mainFunction || mainFunction.isExternal())
     return make_string_error("entry point not found");
 
   if (mainFunction.getType().getFunctionNumParams() != 0)
     return make_string_error("function inputs not supported");
 
   if (Error error = checkCompatibleReturnType<Type>(mainFunction))
     return error;
 
   Type res;
   struct {
     void *data;
   } data;
   data.data = &res;
-  if (auto error = compileAndExecute(options, module, entryPoint, transformer,
-                                     (void **)&data))
+  if (auto error = compileAndExecute(options, module, llvmModuleBuilder,
+                                     entryPoint, transformer, (void **)&data))
     return error;
 
   // Intentional printing of the output so we can test.
   llvm::outs() << res << '\n';
 
   return Error::success();
 }
 
-/// Entry point for all CPU runners. Expects the common argc/argv
-/// arguments for standard C++ main functions and an mlirTransformer.
-/// The latter is applied after parsing the input into MLIR IR and
-/// before passing the MLIR module to the ExecutionEngine.
+/// Entry point for all CPU runners. Expects the common argc/argv arguments for
+/// standard C++ main functions, `mlirTransformer` and `llvmModuleBuilder`.
+/// `mlirTransformer` is applied after parsing the input into MLIR IR and before
+/// passing the MLIR module to the ExecutionEngine.
+/// `llvmModuleBuilder` is a custom function that is passed to ExecutionEngine.
+/// It processes MLIR module and creates LLVM IR module.
 int mlir::JitRunnerMain(
     int argc, char **argv,
-    function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer) {
+    function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer,
+    function_ref<std::unique_ptr<llvm::Module>(ModuleOp, llvm::LLVMContext &)>
+        llvmModuleBuilder) {
   // Create the options struct containing the command line options for the
   // runner. This must come before the command line options are parsed.
   Options options;
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");
 
   Optional<unsigned> optLevel = getCommandLineOptLevel(options);
   SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
       options.optO0, options.optO1, options.optO2, options.optO3};
if (!tmOrError) {
     return EXIT_FAILURE;
   }
 
   auto transformer = mlir::makeLLVMPassesTransformer(
       passes, optLevel, /*targetMachine=*/tmOrError->get(), optPosition);
 
   // Get the function used to compile and execute the module.
   using CompileAndExecuteFnT =
-      Error (*)(Options &, ModuleOp, StringRef,
-                std::function<llvm::Error(llvm::Module *)>);
+      Error (*)(Options &, ModuleOp,
+                function_ref<std::unique_ptr<llvm::Module>(
+                    ModuleOp, llvm::LLVMContext &)>,
+                StringRef, std::function<llvm::Error(llvm::Module *)>);
   auto compileAndExecuteFn =
       llvm::StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
           .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
           .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
           .Case("f32", compileAndExecuteSingleReturnFunction<float>)
           .Case("void", compileAndExecuteVoidFunction)
           .Default(nullptr);
 
   Error error =
       compileAndExecuteFn
-          ? compileAndExecuteFn(options, m.get(),
+          ? compileAndExecuteFn(options, m.get(), llvmModuleBuilder,
                                 options.mainFuncName.getValue(), transformer)
           : make_string_error("unsupported function type");
 
   int exitCode = EXIT_SUCCESS;
   llvm::handleAllErrors(std::move(error),
                         [&exitCode](const llvm::ErrorInfoBase &info) {
                           llvm::errs() << "Error: ";
                           info.log(llvm::errs());
                           llvm::errs() << '\n';
                           exitCode = EXIT_FAILURE;
                         });
 
   return exitCode;
 }

mlir/test/CMakeLists.txt

llvm_canonicalize_cmake_booleans(
   )
 
 # Passed to lit.site.cfg.py.in to set up the path where to find the libraries
 # for linalg integration tests.
 set(MLIR_DIALECT_LINALG_INTEGRATION_TEST_LIB_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 set(MLIR_RUNNER_UTILS_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 
 # Passed to lit.site.cfg.py.in to set up the path where to find the libraries
-# for the mlir cuda / rocm / vulkan runner tests.
+# for the mlir cuda / rocm / spirv / vulkan runner tests.
 set(MLIR_CUDA_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 set(MLIR_ROCM_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
+set(MLIR_SPIRV_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 set(MLIR_VULKAN_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
   MAIN_CONFIG
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
   )

 endif()
 
 if(MLIR_ROCM_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-rocm-runner
   )
 endif()
 
+if(MLIR_SPIRV_RUNNER_ENABLED)
+  add_subdirectory(mlir-spirv-runner)
+  list(APPEND MLIR_TEST_DEPENDS
+    mlir-spirv-runner
+    mlir_test_spirv_runner_c_wrappers
+  )
+endif()
+
 if(MLIR_VULKAN_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 endif()
 
 if(MLIR_BINDINGS_PYTHON_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS

mlir/test/lit.cfg.py

tools.extend([
     ToolSubst('toy-ch2', unresolved='ignore'),
     ToolSubst('toy-ch3', unresolved='ignore'),
     ToolSubst('toy-ch4', unresolved='ignore'),
     ToolSubst('toy-ch5', unresolved='ignore'),
     ToolSubst('%cuda_wrapper_library_dir', config.cuda_wrapper_library_dir, unresolved='ignore'),
     ToolSubst('%linalg_test_lib_dir', config.linalg_test_lib_dir, unresolved='ignore'),
     ToolSubst('%mlir_runner_utils_dir', config.mlir_runner_utils_dir, unresolved='ignore'),
     ToolSubst('%rocm_wrapper_library_dir', config.rocm_wrapper_library_dir, unresolved='ignore'),
+    ToolSubst('%spirv_wrapper_library_dir', config.spirv_wrapper_library_dir, unresolved='ignore'),
     ToolSubst('%vulkan_wrapper_library_dir', config.vulkan_wrapper_library_dir, unresolved='ignore'),
 ])
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)
 
 
 # FileCheck -enable-var-scope is enabled by default in MLIR test
 # This option avoids to accidentally reuse variable across -LABEL match,

mlir/test/lit.site.cfg.py.in

 config.linalg_test_lib_dir = "@MLIR_DIALECT_LINALG_INTEGRATION_TEST_LIB_DIR@"
 config.build_examples = @LLVM_BUILD_EXAMPLES@
 config.run_cuda_tests = @MLIR_CUDA_CONVERSIONS_ENABLED@
 config.cuda_wrapper_library_dir = "@MLIR_CUDA_WRAPPER_LIBRARY_DIR@"
 config.enable_cuda_runner = @MLIR_CUDA_RUNNER_ENABLED@
 config.run_rocm_tests = @MLIR_ROCM_CONVERSIONS_ENABLED@
 config.rocm_wrapper_library_dir = "@MLIR_ROCM_WRAPPER_LIBRARY_DIR@"
 config.enable_rocm_runner = @MLIR_ROCM_RUNNER_ENABLED@
+config.spirv_wrapper_library_dir = "@MLIR_SPIRV_WRAPPER_LIBRARY_DIR@"
+config.enable_spirv_runner = @MLIR_SPIRV_RUNNER_ENABLED@
 config.vulkan_wrapper_library_dir = "@MLIR_VULKAN_WRAPPER_LIBRARY_DIR@"
 config.enable_vulkan_runner = @MLIR_VULKAN_RUNNER_ENABLED@
 config.enable_bindings_python = @MLIR_BINDINGS_PYTHON_ENABLED@
 
 # Support substitution of the tools_dir with user parameters. This is
 # used when we can't determine the tool dir at configuration time.
 try:
     config.llvm_tools_dir = config.llvm_tools_dir % lit_config.params

mlir/test/mlir-spirv-runner/CMakeLists.txt

+set(LLVM_OPTIONAL_SOURCES
+  mlir_test_spirv_runner_c_wrappers.cpp
+  )
+
+add_llvm_library(mlir_test_spirv_runner_c_wrappers SHARED mlir_test_spirv_runner_c_wrappers.cpp)
+target_compile_definitions(mlir_test_spirv_runner_c_wrappers PRIVATE mlir_test_spirv_runner_c_wrappers_EXPORTS)

mlir/test/mlir-spirv-runner/lit.local.cfg

+if not config.enable_spirv_runner:
+  config.unsupported = True

mlir/test/mlir-spirv-runner/mlir_test_spirv_runner_c_wrappers.cpp

+//===- mlir_test_spirv_runner_c_wrappers.cpp - MLIR SPIR-V runner library -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// A small library for SPIR-V runner testing.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/ExecutionEngine/RunnerUtils.h"
+
+// A struct corresponding to a MemRef argument after C-compatible wrapper
+// emission.
+template <typename T, size_t N>
+struct MemRefDescriptor {
+  T *allocated;
+  T *aligned;
+  intptr_t offset;
+  intptr_t sizes[N];
+  intptr_t strides[N];
+};
+
+extern "C" void
+_mlir_ciface_fillI32Buffer(MemRefDescriptor<int32_t, 1> *mem_ref,
+                           int32_t value) {
+  std::fill_n(mem_ref->allocated, mem_ref->sizes[0], value);
+}

mlir/test/mlir-spirv-runner/simple.mlir

+// RUN: mlir-spirv-runner %s -e main --entry-point-result=void --shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%spirv_wrapper_library_dir/libmlir_test_spirv_runner_c_wrappers%shlibext
+
+// CHECK: [8,  8,  8,  8,  8,  8]
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_variable_pointers]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+  gpu.module @kernels {
+    gpu.func @double(%arg0 : memref<6xi32>, %arg1 : memref<6xi32>)
+      kernel attributes { spv.entry_point_abi = {local_size = dense<[1, 1, 1]>: vector<3xi32>}} {
+      %factor = constant 2 : i32
+
+      %i0 = constant 0 : index
+      %i1 = constant 1 : index
+      %i2 = constant 2 : index
+      %i3 = constant 3 : index
+      %i4 = constant 4 : index
+      %i5 = constant 5 : index
+
+      %x0 = load %arg0[%i0] : memref<6xi32>
+      %x1 = load %arg0[%i1] : memref<6xi32>
+      %x2 = load %arg0[%i2] : memref<6xi32>
+      %x3 = load %arg0[%i3] : memref<6xi32>
+      %x4 = load %arg0[%i4] : memref<6xi32>
+      %x5 = load %arg0[%i5] : memref<6xi32>
+
+      %y0 = muli %x0, %factor : i32
+      %y1 = muli %x1, %factor : i32
+      %y2 = muli %x2, %factor : i32
+      %y3 = muli %x3, %factor : i32
+      %y4 = muli %x4, %factor : i32
+      %y5 = muli %x5, %factor : i32
+
+      store %y0, %arg1[%i0] : memref<6xi32>
+      store %y1, %arg1[%i1] : memref<6xi32>
+      store %y2, %arg1[%i2] : memref<6xi32>
+      store %y3, %arg1[%i3] : memref<6xi32>
+      store %y4, %arg1[%i4] : memref<6xi32>
+      store %y5, %arg1[%i5] : memref<6xi32>
+      gpu.return
+    }
+  }
+  func @main() {
+    %input = alloc() : memref<6xi32>
+    %output = alloc() : memref<6xi32>
+    %four = constant 4 : i32
+    %zero = constant 0 : i32
+    %input_casted = memref_cast %input : memref<6xi32> to memref<?xi32>
+    %output_casted = memref_cast %output : memref<6xi32> to memref<?xi32>
+    call @fillI32Buffer(%input_casted, %four) : (memref<?xi32>, i32) -> ()
+    call @fillI32Buffer(%output_casted, %zero) : (memref<?xi32>, i32) -> ()
+
+    %one = constant 1 : index
+    "gpu.launch_func"(%one, %one, %one,
+                      %one, %one, %one,
+                      %input, %output) { kernel = @kernels::@double }
+        : (index, index, index, index, index, index, memref<6xi32>, memref<6xi32>) -> ()
+    %result = memref_cast %output : memref<6xi32> to memref<*xi32>
+    call @print_memref_i32(%result) : (memref<*xi32>) -> ()
+    return
+  }
+
+  func @fillI32Buffer(%arg0 : memref<?xi32>, %arg1 : i32)
+  func @print_memref_i32(%ptr : memref<*xi32>)
+}

mlir/tools/CMakeLists.txt

 add_subdirectory(mlir-cuda-runner)
 add_subdirectory(mlir-cpu-runner)
 add_subdirectory(mlir-linalg-ods-gen)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-rocm-runner)
 add_subdirectory(mlir-shlib)
+add_subdirectory(mlir-spirv-runner)
 add_subdirectory(mlir-translate)
 add_subdirectory(mlir-vulkan-runner)
 No newline at end of file

mlir/tools/mlir-spirv-runner/CMakeLists.txt

+set(LLVM_OPTIONAL_SOURCES
+  mlir-spirv-runner.cpp
+  )
+
+if (MLIR_SPIRV_RUNNER_ENABLED)
+  message(STATUS "Building SPIR-V runner")
+
+  add_llvm_tool(mlir-spirv-runner
+    mlir-spirv-runner.cpp
+  )
+
+  llvm_update_compile_flags(mlir-spirv-runner)
+
+  get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+  get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+
+  target_link_libraries(mlir-spirv-runner PRIVATE
+    ${conversion_libs}
+    ${dialect_libs}
+    MLIRAnalysis
+    MLIREDSC
+    MLIRExecutionEngine
+    MLIRIR
+    MLIRJitRunner
+    MLIRLLVMIR
+    MLIRParser
+    MLIRTargetLLVMIR
+    MLIRTransforms
+    MLIRTranslation
+    MLIRSupport
+  )
+endif()

mlir/tools/mlir-spirv-runner/mlir-spirv-runner.cpp

+//===- mlir-spirv-runner.cpp - MLIR SPIR-V Execution on CPU ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Main entry point to a command line utility that executes an MLIR file on the
+// CPU by translating MLIR GPU module and host part to LLVM IR before
+// JIT-compiling and executing.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
+#include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
+#include "mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/GPU/Passes.h"
+#include "mlir/Dialect/SPIRV/Passes.h"
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/ExecutionEngine/JitRunner.h"
+#include "mlir/ExecutionEngine/OptUtils.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Target/LLVMIR.h"
+
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
+
+using namespace mlir;
+
+/// A utility function that builds llvm::Module from two nested MLIR modules.
+///
+/// module @main {
+///   module @kernel {
+///     // Some ops
+///   }
+///   // Some other ops
+/// }
+///
+/// Each of these two modules is translated to LLVM IR module, then they are
+/// linked together and returned.
+static std::unique_ptr<llvm::Module>
+convertMLIRModule(ModuleOp module, llvm::LLVMContext &context) {
+  std::unique_ptr<llvm::Module> kernelModule;
+  module.walk([&](ModuleOp nested) -> WalkResult {
+    if (nested.getParentOp()) {
+      // Found kernel module. Translate it to LLVM IR, remove from the main
+      // module and terminate.
+      kernelModule = translateModuleToLLVMIR(nested, context);
+      nested.erase();
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  std::unique_ptr<llvm::Module> mainModule =
+      translateModuleToLLVMIR(module, context);
+  llvm::Linker::linkModules(*mainModule, std::move(kernelModule));
+  return mainModule;
+}
+
+static LogicalResult runMLIRPasses(ModuleOp module) {
+  PassManager passManager(module.getContext());
+  applyPassManagerCLOptions(passManager);
+  passManager.addPass(createGpuKernelOutliningPass());
+  passManager.addPass(createConvertGPUToSPIRVPass());
+
+  OpPassManager &nestedPM = passManager.nest<spirv::ModuleOp>();
+  nestedPM.addPass(spirv::createLowerABIAttributesPass());
+  nestedPM.addPass(spirv::createUpdateVersionCapabilityExtensionPass());
+  nestedPM.addPass(spirv::createEncodeDescriptorSetsPass());
+
+  passManager.addPass(createGPULaunchFuncToStandardCallPass());
+  LowerToLLVMOptions options = {
+      /*useBarePtrCallConv=*/false,
+      /*emitCWrappers=*/true,
+      /*indexBitwidth=*/kDeriveIndexBitwidthFromDataLayout};
+  passManager.addPass(createLowerToLLVMPass(options));
+  passManager.addPass(createConvertSPIRVToLLVMPass());
+  passManager.addPass(createEmulateKernelCallInLLVMPass());
+  return passManager.run(module);
+}
+
+int main(int argc, char **argv) {
+  mlir::registerAllDialects();
+  llvm::InitLLVM y(argc, argv);
+  llvm::InitializeNativeTarget();
+  llvm::InitializeNativeTargetAsmPrinter();
+  mlir::initializeLLVMPasses();
+
+  return mlir::JitRunnerMain(argc, argv, &runMLIRPasses, &convertMLIRModule);
+}