mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp

   1 //===- ConvertGPUToSPIRVPass.cpp - GPU to SPIR-V dialect lowering passes --===//
   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 file implements a pass to convert a kernel function in the GPU Dialect
  10 // into a spv.module operation
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
  15 #include "../PassDetail.h"
  16 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.h"
  17 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
  18 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.h"
  19 #include "mlir/Dialect/GPU/GPUDialect.h"
  20 #include "mlir/Dialect/SCF/SCF.h"
  21 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
  22 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
  23 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
  24
  25 using namespace mlir;
  26
  27 namespace {
  28 /// Pass to lower GPU Dialect to SPIR-V. The pass only converts the gpu.func ops
  29 /// inside gpu.module ops. i.e., the function that are referenced in
  30 /// gpu.launch_func ops. For each such function
  31 ///
  32 /// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp
  33 /// (the original function is still needed by the gpu::LaunchKernelOp, so cannot
  34 /// replace it).
  35 ///
  36 /// 2) Lower the body of the spirv::ModuleOp.
  37 struct GPUToSPIRVPass : public ConvertGPUToSPIRVBase<GPUToSPIRVPass> {
  38   void runOnOperation() override;
  39 };
  40 } // namespace
  41
  42 void GPUToSPIRVPass::runOnOperation() {
  43   MLIRContext *context = &getContext();
  44   ModuleOp module = getOperation();
  45
  46   SmallVector<Operation *, 1> kernelModules;
  47   OpBuilder builder(context);
  48   module.walk([&builder, &kernelModules](gpu::GPUModuleOp moduleOp) {
  49     // For each kernel module (should be only 1 for now, but that is not a
  50     // requirement here), clone the module for conversion because the
  51     // gpu.launch function still needs the kernel module.
  52     builder.setInsertionPoint(moduleOp.getOperation());
  53     kernelModules.push_back(builder.clone(*moduleOp.getOperation()));
  54   });
  55
  56   auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
  57   std::unique_ptr<ConversionTarget> target =
  58       spirv::SPIRVConversionTarget::get(targetAttr);
  59
  60   SPIRVTypeConverter typeConverter(targetAttr);
  61   ScfToSPIRVContext scfContext;
  62   OwningRewritePatternList patterns;
  63   populateGPUToSPIRVPatterns(context, typeConverter, patterns);
  64   populateSCFToSPIRVPatterns(context, typeConverter,scfContext, patterns);
  65   populateStandardToSPIRVPatterns(context, typeConverter, patterns);
  66
  67   if (failed(applyFullConversion(kernelModules, *target, std::move(patterns))))
  68     return signalPassFailure();
  69 }
  70
  71 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertGPUToSPIRVPass() {
  72   return std::make_unique<GPUToSPIRVPass>();
  73 }