[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 remove the runner component of the mlir-spirv-cpu-runner.

The runtime libraries and the tests are moved and renamed to fit into
the Execution Engine and Integration tests, following the model of the
similar migration done for the CUDA Runner in D97463.

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/lib/ExecutionEngine/CMakeLists.txt

@@ -14,6 +14,7 @@ set(LLVM_OPTIONAL_SOURCES
   RunnerUtils.cpp
   OptUtils.cpp
   JitRunner.cpp
+  SpirvCpuRuntimeWrappers.cpp
   SyclRuntimeWrappers.cpp
   )
 
@@ -401,4 +402,17 @@ if(LLVM_ENABLE_PIC)
 
     set_property(TARGET mlir_sycl_runtime APPEND PROPERTY BUILD_RPATH "${LevelZero_LIBRARIES_DIR}" "${SyclRuntime_LIBRARIES_DIR}")
   endif()
+
+  if(MLIR_ENABLE_SPIRV_CPU_RUNNER)
+    add_mlir_library(mlir_spirv_cpu_runtime
+      SHARED
+      SpirvCpuRuntimeWrappers.cpp
+
+      EXCLUDE_FROM_LIBMLIR
+    )
+
+    target_compile_definitions(mlir_spirv_cpu_runtime
+      PRIVATE
+      mlir_spirv_cpu_runtime_EXPORTS)
+  endif()
 endif()

mlir/test/mlir-spirv-cpu-runner/mlir_test_spirv_cpu_runner_c_wrappers.cpp renamed to mlir/lib/ExecutionEngine/SpirvCpuRuntimeWrappers.cpp

@@ -1,4 +1,4 @@
-//===- mlir_test_spirv_cpu_runner_c_wrappers.cpp - Runner testing library -===//
+//===- SpirvCpuRuntimeWrappers.cpp - Runner testing library -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

mlir/test/CMakeLists.txt

@@ -197,10 +197,8 @@ if(LLVM_BUILD_EXAMPLES)
 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
+    mlir_spirv_cpu_runtime
   )
 endif()
 

mlir/test/mlir-spirv-cpu-runner/double.mlir renamed to mlir/test/Integration/GPU/SPIRV/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_spirv_cpu_runtime --link-nested-modules \
 // RUN: | FileCheck %s
 
 // CHECK: [8,  8,  8,  8,  8,  8]

mlir/test/mlir-spirv-cpu-runner/lit.local.cfg renamed to mlir/test/Integration/GPU/SPIRV/lit.local.cfg

@@ -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_spirv_cpu_runtime --link-nested-modules \
 // RUN: | FileCheck %s
 
 // CHECK: data =

mlir/test/mlir-spirv-cpu-runner/simple_add.mlir renamed to mlir/test/Integration/GPU/SPIRV/simple_add.mlir

@@ -125,11 +125,6 @@ def add_runtime(name):
     "not",
 ]
 
-if config.enable_spirv_cpu_runner:
-    tools.extend(
-        ["mlir-spirv-cpu-runner", add_runtime("mlir_test_spirv_cpu_runner_c_wrappers")]
-    )
-
 if config.enable_vulkan_runner:
     tools.extend([add_runtime("vulkan-runtime-wrappers")])
 
@@ -142,6 +137,9 @@ def add_runtime(name):
 if config.enable_sycl_runner:
     tools.extend([add_runtime("mlir_sycl_runtime")])
 
+if config.enable_spirv_cpu_runner:
+    tools.extend([add_runtime("mlir_spirv_cpu_runtime")])
+
 if config.mlir_run_arm_sve_tests or config.mlir_run_arm_sme_tests:
     tools.extend([add_runtime("mlir_arm_runner_utils")])
 

mlir/test/lit.cfg.py

@@ -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/test/mlir-spirv-cpu-runner/CMakeLists.txt

@@ -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);
 }