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

mlir/docs/SPIRVToLLVMDialectConversion.md

@@ -501,7 +501,7 @@ The SPIR-V to LLVM conversion does not involve modelling of workgroups. Hence,
 we say that only current invocation is in conversion's scope. This means that
 global variables with pointers of `Input`, `Output`, and `Private` storage
 classes are supported. Also, `StorageBuffer` storage class is allowed for
-executing [`mlir-spirv-cpu-runner`](#mlir-spirv-cpu-runner).
+executing [SPIR-V CPU Runner tests](#spir-v-cpu-runner-tests).
 
 Moreover, `bind` that specifies the descriptor set and the binding number and
 `built_in` that specifies SPIR-V `BuiltIn` decoration have no conversion into
@@ -815,14 +815,15 @@ Module in SPIR-V has one region that contains one block. It is defined via
 `spirv.module` is converted into `ModuleOp`. This plays a role of enclosing scope
 to LLVM ops. At the moment, SPIR-V module attributes are ignored.
 
-## `mlir-spirv-cpu-runner`
+## SPIR-V CPU Runner Tests
 
-`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
-supported.
+The `mlir-cpu-runner` has support for executing a `gpu` dialect kernel on the
+CPU via SPIR-V to LLVM dialect conversion. This is referred to as the "SPIR-V
+CPU Runner". The `--link-nested-modules` flag needs to be passed for this.
+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
 
@@ -846,7 +847,7 @@ executed using `ExecutionEngine`.
 ### Walk-through
 
 This section gives a detailed overview of the IR changes while running
-`mlir-spirv-cpu-runner`. First, consider that we have the following IR. (For
+SPIR-V CPU Runner tests. First, consider that we have the following IR. (For
 simplicity some type annotations and function implementations have been
 omitted).
 

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 → 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 → 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/simple_add.mlir → mlir/test/Integration/GPU/SPIRV/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_spirv_cpu_runtime --link-nested-modules \
 // RUN: | FileCheck %s
 
 // CHECK: data =

mlir/test/lib/Pass/TestSPIRVCPURunnerPipeline.cpp

@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Implements a pipeline for use by mlir-spirv-cpu-runner tests.
+// Implements a pipeline for use by SPIR-V CPU Runner tests.
 //
 //===----------------------------------------------------------------------===//
 
@@ -40,7 +40,7 @@ void registerTestSPIRVCPURunnerPipeline() {
   PassPipelineRegistration<>(
       "test-spirv-cpu-runner-pipeline",
       "Runs a series of passes for lowering SPIR-V-dialect MLIR to "
-      "LLVM-dialect MLIR intended for mlir-spirv-cpu-runner.",
+      "LLVM-dialect MLIR intended for SPIR-V CPU Runner tests.",
       buildTestSPIRVCPURunnerPipeline);
 }
 } // namespace test

mlir/test/lit.cfg.py

@@ -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/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,68 @@
 #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;
+
+// TODO: Consider removing this linking functionality from the SPIR-V CPU Runner
+//       flow in favour of a more proper host/device split like other runners.
+//       https://github.com/llvm/llvm-project/issues/115348
+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 SPIR-V 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 +88,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);
 }