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