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/ConvertToSPIRV/ConvertToSPIRVPass.h"
  16 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
  17 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
  18 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
  19 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
  20 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
  21 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
  22 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.h"
  23 #include "mlir/Dialect/Arith/IR/Arith.h"
  24 #include "mlir/Dialect/Func/IR/FuncOps.h"
  25 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
  26 #include "mlir/Dialect/GPU/Transforms/Passes.h"
  27 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
  28 #include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
  29 #include "mlir/Dialect/MemRef/IR/MemRef.h"
  30 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
  31 #include "mlir/Dialect/SCF/IR/SCF.h"
  32 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
  33 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
  34 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
  35 #include "mlir/Dialect/Vector/IR/VectorOps.h"
  36 #include "mlir/ExecutionEngine/JitRunner.h"
  37 #include "mlir/Pass/Pass.h"
  38 #include "mlir/Pass/PassManager.h"
  39 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
  40 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
  41 #include "llvm/Support/InitLLVM.h"
  42 #include "llvm/Support/TargetSelect.h"
  43
  44 using namespace mlir;
  45
  46 namespace {
  47 struct VulkanRunnerOptions {
  48   llvm::cl::OptionCategory category{"mlir-vulkan-runner options"};
  49   llvm::cl::opt<bool> spirvWebGPUPrepare{
  50       "vulkan-runner-spirv-webgpu-prepare",
  51       llvm::cl::desc("Run MLIR transforms used when targetting WebGPU"),
  52       llvm::cl::cat(category)};
  53 };
  54 } // namespace
  55
  56 static LogicalResult runMLIRPasses(Operation *op,
  57                                    VulkanRunnerOptions &options) {
  58   auto module = dyn_cast<ModuleOp>(op);
  59   if (!module)
  60     return op->emitOpError("expected a 'builtin.module' op");
  61   PassManager passManager(module.getContext());
  62   if (failed(applyPassManagerCLOptions(passManager)))
  63     return failure();
  64
  65   passManager.addPass(createGpuKernelOutliningPass());
  66   passManager.addPass(memref::createFoldMemRefAliasOpsPass());
  67
  68   ConvertToSPIRVPassOptions convertToSPIRVOptions{};
  69   convertToSPIRVOptions.convertGPUModules = true;
  70   passManager.addPass(createConvertToSPIRVPass(convertToSPIRVOptions));
  71   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
  72   modulePM.addPass(spirv::createSPIRVLowerABIAttributesPass());
  73   modulePM.addPass(spirv::createSPIRVUpdateVCEPass());
  74   if (options.spirvWebGPUPrepare)
  75     modulePM.addPass(spirv::createSPIRVWebGPUPreparePass());
  76
  77   passManager.addPass(createConvertGpuLaunchFuncToVulkanLaunchFuncPass());
  78   passManager.addPass(createFinalizeMemRefToLLVMConversionPass());
  79   passManager.addPass(createConvertVectorToLLVMPass());
  80   passManager.nest<func::FuncOp>().addPass(LLVM::createRequestCWrappersPass());
  81   ConvertFuncToLLVMPassOptions funcToLLVMOptions{};
  82   funcToLLVMOptions.indexBitwidth =
  83       DataLayout(module).getTypeSizeInBits(IndexType::get(module.getContext()));
  84   passManager.addPass(createConvertFuncToLLVMPass(funcToLLVMOptions));
  85   passManager.addPass(createReconcileUnrealizedCastsPass());
  86   passManager.addPass(createConvertVulkanLaunchFuncToVulkanCallsPass());
  87
  88   return passManager.run(module);
  89 }
  90
  91 int main(int argc, char **argv) {
  92   llvm::llvm_shutdown_obj x;
  93   registerPassManagerCLOptions();
  94
  95   llvm::InitLLVM y(argc, argv);
  96   llvm::InitializeNativeTarget();
  97   llvm::InitializeNativeTargetAsmPrinter();
  98
  99   // Initialize runner-specific CLI options. These will be parsed and
 100   // initialzied in `JitRunnerMain`.
 101   VulkanRunnerOptions options;
 102   auto runPassesWithOptions = [&options](Operation *op, JitRunnerOptions &) {
 103     return runMLIRPasses(op, options);
 104   };
 105
 106   mlir::JitRunnerConfig jitRunnerConfig;
 107   jitRunnerConfig.mlirTransformer = runPassesWithOptions;
 108
 109   mlir::DialectRegistry registry;
 110   registry.insert<mlir::arith::ArithDialect, mlir::LLVM::LLVMDialect,
 111                   mlir::gpu::GPUDialect, mlir::spirv::SPIRVDialect,
 112                   mlir::scf::SCFDialect, mlir::func::FuncDialect,
 113                   mlir::memref::MemRefDialect, mlir::vector::VectorDialect>();
 114   mlir::registerBuiltinDialectTranslation(registry);
 115   mlir::registerLLVMDialectTranslation(registry);
 116
 117   return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 118 }