mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

   1 //===- mlir-vulkan-runner.cpp - MLIR Vulkan Execution Driver --------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This is a command line utility that executes an MLIR file on the Vulkan by
  10 // translating MLIR GPU module to SPIR-V and host part to LLVM IR before
  11 // JIT-compiling and executing the latter.
  12 //
  13 //===----------------------------------------------------------------------===//
  14
  15 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
  16 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
  17 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
  18 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
  19 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
  20 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
  21 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.h"
  22 #include "mlir/Dialect/Arith/IR/Arith.h"
  23 #include "mlir/Dialect/Func/IR/FuncOps.h"
  24 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
  25 #include "mlir/Dialect/GPU/Transforms/Passes.h"
  26 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
  27 #include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
  28 #include "mlir/Dialect/MemRef/IR/MemRef.h"
  29 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
  30 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
  31 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
  32 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
  33 #include "mlir/Dialect/Vector/IR/VectorOps.h"
  34 #include "mlir/ExecutionEngine/JitRunner.h"
  35 #include "mlir/Pass/Pass.h"
  36 #include "mlir/Pass/PassManager.h"
  37 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
  38 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
  39 #include "llvm/Support/InitLLVM.h"
  40 #include "llvm/Support/TargetSelect.h"
  41
  42 using namespace mlir;
  43
  44 namespace {
  45 struct VulkanRunnerOptions {
  46   llvm::cl::OptionCategory category{"mlir-vulkan-runner options"};
  47   llvm::cl::opt<bool> spirvWebGPUPrepare{
  48       "vulkan-runner-spirv-webgpu-prepare",
  49       llvm::cl::desc("Run MLIR transforms used when targetting WebGPU"),
  50       llvm::cl::cat(category)};
  51 };
  52 } // namespace
  53
  54 static LogicalResult runMLIRPasses(Operation *op,
  55                                    VulkanRunnerOptions &options) {
  56   auto module = dyn_cast<ModuleOp>(op);
  57   if (!module)
  58     return op->emitOpError("expected a 'builtin.module' op");
  59   PassManager passManager(module.getContext());
  60   if (failed(applyPassManagerCLOptions(passManager)))
  61     return failure();
  62
  63   passManager.addPass(createGpuKernelOutliningPass());
  64   passManager.addPass(memref::createFoldMemRefAliasOpsPass());
  65
  66   passManager.addPass(createConvertGPUToSPIRVPass(/*mapMemorySpace=*/true));
  67   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
  68   modulePM.addPass(spirv::createSPIRVLowerABIAttributesPass());
  69   modulePM.addPass(spirv::createSPIRVUpdateVCEPass());
  70   if (options.spirvWebGPUPrepare)
  71     modulePM.addPass(spirv::createSPIRVWebGPUPreparePass());
  72
  73   auto enableOpaquePointers = [](auto passOption) {
  74     passOption.useOpaquePointers = true;
  75     return passOption;
  76   };
  77
  78   passManager.addPass(createConvertGpuLaunchFuncToVulkanLaunchFuncPass());
  79   passManager.addPass(createFinalizeMemRefToLLVMConversionPass(
  80       enableOpaquePointers(FinalizeMemRefToLLVMConversionPassOptions{})));
  81   passManager.addPass(createConvertVectorToLLVMPass(
  82       enableOpaquePointers(ConvertVectorToLLVMPassOptions{})));
  83   passManager.nest<func::FuncOp>().addPass(LLVM::createRequestCWrappersPass());
  84   ConvertFuncToLLVMPassOptions funcToLLVMOptions{};
  85   funcToLLVMOptions.indexBitwidth =
  86       DataLayout(module).getTypeSizeInBits(IndexType::get(module.getContext()));
  87   passManager.addPass(
  88       createConvertFuncToLLVMPass(enableOpaquePointers(funcToLLVMOptions)));
  89   passManager.addPass(createReconcileUnrealizedCastsPass());
  90   passManager.addPass(createConvertVulkanLaunchFuncToVulkanCallsPass(
  91       enableOpaquePointers(ConvertVulkanLaunchFuncToVulkanCallsPassOptions{})));
  92
  93   return passManager.run(module);
  94 }
  95
  96 int main(int argc, char **argv) {
  97   llvm::llvm_shutdown_obj x;
  98   registerPassManagerCLOptions();
  99
 100   llvm::InitLLVM y(argc, argv);
 101   llvm::InitializeNativeTarget();
 102   llvm::InitializeNativeTargetAsmPrinter();
 103
 104   // Initialize runner-specific CLI options. These will be parsed and
 105   // initialzied in `JitRunnerMain`.
 106   VulkanRunnerOptions options;
 107   auto runPassesWithOptions = [&options](Operation *op, JitRunnerOptions &) {
 108     return runMLIRPasses(op, options);
 109   };
 110
 111   mlir::JitRunnerConfig jitRunnerConfig;
 112   jitRunnerConfig.mlirTransformer = runPassesWithOptions;
 113
 114   mlir::DialectRegistry registry;
 115   registry.insert<mlir::arith::ArithDialect, mlir::LLVM::LLVMDialect,
 116                   mlir::gpu::GPUDialect, mlir::spirv::SPIRVDialect,
 117                   mlir::func::FuncDialect, mlir::memref::MemRefDialect,
 118                   mlir::vector::VectorDialect>();
 119   mlir::registerBuiltinDialectTranslation(registry);
 120   mlir::registerLLVMDialectTranslation(registry);
 121
 122   return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 123 }