[mlir] Remove the mlir-spirv-cpu-runner (move to mlir-cpu-runner)

This commit builds on and completes the work done in
9f6c632 to eliminate the need for a
separate mlir-spirv-cpu-runner binary. Since the MLIR processing is
already done outside this runner, the only real difference between it
and the mlir-cpu-runner is the final linking step between the nested
LLVM IR modules. By moving this step into mlir-cpu-runner behind a new
command-line flag (`--link-nested-modules`), this commit is able to
completely remvoe the runner component of the mlir-spirv-cpu-runner.
The runtime libraries and associated build infrastructure are left
undisturbed.

Author: andfau-amd

mlir/docs/SPIRVToLLVMDialectConversion.md

@@ -817,12 +817,14 @@ to LLVM ops. At the moment, SPIR-V module attributes are ignored.
 
 ## `mlir-spirv-cpu-runner`
 
-`mlir-spirv-cpu-runner` allows to execute `gpu` dialect kernel on the CPU via
-SPIR-V to LLVM dialect conversion. Currently, only single-threaded kernel is
+`mlir-spirv-cpu-runner` is the name of a formerly separate tool that allows to
+execute `gpu` dialect kernel on the CPU via SPIR-V to LLVM dialect conversion.
+This runner is now merged with the `mlir-cpu-runner`; pass the
+`--link-nested-modules` flag. Currently, only single-threaded kernels are
 supported.
 
-To build the runner, add the following option to `cmake`: `bash
--DMLIR_ENABLE_SPIRV_CPU_RUNNER=1`
+To build the required runtime libaries, add the following option to `cmake`:
+`-DMLIR_ENABLE_SPIRV_CPU_RUNNER=1`
 
 ### Pipeline
 

mlir/test/CMakeLists.txt

@@ -199,7 +199,6 @@ endif()
 if(MLIR_ENABLE_SPIRV_CPU_RUNNER)
   add_subdirectory(mlir-spirv-cpu-runner)
   list(APPEND MLIR_TEST_DEPENDS
-    mlir-spirv-cpu-runner
     mlir_test_spirv_cpu_runner_c_wrappers
   )
 endif()

mlir/test/mlir-spirv-cpu-runner/double.mlir

@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s -test-spirv-cpu-runner-pipeline \
-// RUN: | mlir-spirv-cpu-runner - -e main --entry-point-result=void --shared-libs=%mlir_runner_utils,%mlir_test_spirv_cpu_runner_c_wrappers \
+// RUN: | mlir-cpu-runner - -e main --entry-point-result=void --shared-libs=%mlir_runner_utils,%mlir_test_spirv_cpu_runner_c_wrappers --link-nested-modules \
 // RUN: | FileCheck %s
 
 // CHECK: [8,  8,  8,  8,  8,  8]

mlir/test/mlir-spirv-cpu-runner/simple_add.mlir

@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s -test-spirv-cpu-runner-pipeline \
-// RUN: | mlir-spirv-cpu-runner - -e main --entry-point-result=void --shared-libs=%mlir_runner_utils,%mlir_test_spirv_cpu_runner_c_wrappers \
+// RUN: | mlir-cpu-runner - -e main --entry-point-result=void --shared-libs=%mlir_runner_utils,%mlir_test_spirv_cpu_runner_c_wrappers --link-nested-modules \
 // RUN: | FileCheck %s
 
 // CHECK: data =

mlir/tools/CMakeLists.txt

@@ -6,7 +6,6 @@ add_subdirectory(mlir-query)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-rewrite)
 add_subdirectory(mlir-shlib)
-add_subdirectory(mlir-spirv-cpu-runner)
 add_subdirectory(mlir-translate)
 add_subdirectory(mlir-vulkan-runner)
 add_subdirectory(tblgen-lsp-server)

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

@@ -17,10 +17,65 @@
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Export.h"
 
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
+using namespace mlir;
+
+llvm::cl::opt<bool> LinkNestedModules(
+    "link-nested-modules",
+    llvm::cl::desc("Link two nested MLIR modules into a single LLVM IR module. "
+                   "Useful if both the host and device code can be run on the "
+                   "same CPU, as in mlir-spirv-cpu-runner tests."));
+
+/// 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(Operation *op, llvm::LLVMContext &context) {
+  auto module = dyn_cast<ModuleOp>(op);
+  if (!module)
+    return op->emitError("op must be a 'builtin.module"), nullptr;
+
+  std::unique_ptr<llvm::Module> kernelModule;
+  if (LinkNestedModules) {
+    // Verify that there is only one nested module.
+    auto modules = module.getOps<ModuleOp>();
+    if (!llvm::hasSingleElement(modules)) {
+      module.emitError("The module must contain exactly one nested module");
+      return nullptr;
+    }
+
+    // Translate nested module and erase it.
+    ModuleOp nested = *modules.begin();
+    kernelModule = translateModuleToLLVMIR(nested, context);
+    nested.erase();
+  }
+
+  std::unique_ptr<llvm::Module> mainModule =
+      translateModuleToLLVMIR(module, context);
+
+  if (LinkNestedModules)
+    llvm::Linker::linkModules(*mainModule, std::move(kernelModule));
+
+  return mainModule;
+}
+
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
@@ -30,5 +85,7 @@ int main(int argc, char **argv) {
   mlir::DialectRegistry registry;
   mlir::registerAllToLLVMIRTranslations(registry);
 
-  return mlir::JitRunnerMain(argc, argv, registry);
+  mlir::JitRunnerConfig jitRunnerConfig;
+  jitRunnerConfig.llvmModuleBuilder = convertMLIRModule;
+  return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 }