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

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

@@ -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 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

@@ -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/lib/Pass/TestSPIRVCPURunnerPipeline.cpp

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