1 //===- GPUToSPIRV.cpp - MLIR SPIR-V lowering passes -----------------------===//
2 //
3 // Copyright 2019 The MLIR Authors.
4 //
5 // Licensed under the Apache License, Version 2.0 (the "License");
6 // you may not use this file except in compliance with the License.
7 // You may obtain a copy of the License at
8 //
9 //   http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing, software
12 // distributed under the License is distributed on an "AS IS" BASIS,
13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 // See the License for the specific language governing permissions and
15 // limitations under the License.
16 // =============================================================================
17 //
18 // This file implements a pass to convert a kernel function in the GPU Dialect
19 // into a spv.module operation
20 //
21 //===----------------------------------------------------------------------===//
22 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.h"
23 #include "mlir/Dialect/GPU/GPUDialect.h"
24 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
25 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
26 #include "mlir/Pass/Pass.h"
27 
28 using namespace mlir;
29 
30 namespace {
31 
32 /// Pattern lowering GPU block/thread size/id to loading SPIR-V invocation
33 /// builin variables.
34 template <typename OpTy, spirv::BuiltIn builtin>
35 class LaunchConfigConversion : public SPIRVOpLowering<OpTy> {
36 public:
37   using SPIRVOpLowering<OpTy>::SPIRVOpLowering;
38 
39   PatternMatchResult
40   matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
41                   ConversionPatternRewriter &rewriter) const override;
42 };
43 
44 /// Pattern to convert a kernel function in GPU dialect (a FuncOp with the
45 /// attribute gpu.kernel) within a spv.module.
46 class KernelFnConversion final : public SPIRVOpLowering<FuncOp> {
47 public:
48   using SPIRVOpLowering<FuncOp>::SPIRVOpLowering;
49 
50   PatternMatchResult
51   matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
52                   ConversionPatternRewriter &rewriter) const override;
53 };
54 } // namespace
55 
56 template <typename OpTy, spirv::BuiltIn builtin>
57 PatternMatchResult LaunchConfigConversion<OpTy, builtin>::matchAndRewrite(
58     Operation *op, ArrayRef<Value *> operands,
59     ConversionPatternRewriter &rewriter) const {
60   auto dimAttr = op->getAttrOfType<StringAttr>("dimension");
61   if (!dimAttr) {
62     return this->matchFailure();
63   }
64   int32_t index = 0;
65   if (dimAttr.getValue() == "x") {
66     index = 0;
67   } else if (dimAttr.getValue() == "y") {
68     index = 1;
69   } else if (dimAttr.getValue() == "z") {
70     index = 2;
71   } else {
72     return this->matchFailure();
73   }
74 
75   // SPIR-V invocation builtin variables are a vector of type <3xi32>
76   auto spirvBuiltin = this->loadFromBuiltinVariable(op, builtin, rewriter);
77   rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
78       op, rewriter.getIntegerType(32), spirvBuiltin,
79       rewriter.getI32ArrayAttr({index}));
80   return this->matchSuccess();
81 }
82 
83 PatternMatchResult
84 KernelFnConversion::matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
85                                     ConversionPatternRewriter &rewriter) const {
86   auto funcOp = cast<FuncOp>(op);
87   FuncOp newFuncOp;
88   if (!gpu::GPUDialect::isKernel(funcOp)) {
89     return succeeded(lowerFunction(funcOp, &typeConverter, rewriter, newFuncOp))
90                ? matchSuccess()
91                : matchFailure();
92   }
93 
94   if (failed(
95           lowerAsEntryFunction(funcOp, &typeConverter, rewriter, newFuncOp))) {
96     return matchFailure();
97   }
98   return matchSuccess();
99 }
100 
101 namespace {
102 /// Pass to lower GPU Dialect to SPIR-V. The pass only converts those functions
103 /// that have the "gpu.kernel" attribute, i.e. those functions that are
104 /// referenced in gpu::LaunchKernelOp operations. For each such function
105 ///
106 /// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp
107 /// (the original function is still needed by the gpu::LaunchKernelOp, so cannot
108 /// replace it).
109 ///
110 /// 2) Lower the body of the spirv::ModuleOp.
111 class GPUToSPIRVPass : public ModulePass<GPUToSPIRVPass> {
112   void runOnModule() override;
113 };
114 } // namespace
115 
116 void GPUToSPIRVPass::runOnModule() {
117   auto context = &getContext();
118   auto module = getModule();
119 
120   SmallVector<Operation *, 4> spirvModules;
121   module.walk([&module, &spirvModules](FuncOp funcOp) {
122     if (gpu::GPUDialect::isKernel(funcOp)) {
123       OpBuilder builder(module.getBodyRegion());
124       // Create a new spirv::ModuleOp for this function, and clone the
125       // function into it.
126       // TODO : Generalize this to account for different extensions,
127       // capabilities, extended_instruction_sets, other addressing models
128       // and memory models.
129       auto spvModule = builder.create<spirv::ModuleOp>(
130           funcOp.getLoc(),
131           builder.getI32IntegerAttr(
132               static_cast<int32_t>(spirv::AddressingModel::Logical)),
133           builder.getI32IntegerAttr(
134               static_cast<int32_t>(spirv::MemoryModel::GLSL450)),
135           builder.getStrArrayAttr(
136               spirv::stringifyCapability(spirv::Capability::Shader)),
137           builder.getStrArrayAttr(spirv::stringifyExtension(
138               spirv::Extension::SPV_KHR_storage_buffer_storage_class)));
139       // Hardwire the capability to be Shader.
140       OpBuilder moduleBuilder(spvModule.getOperation()->getRegion(0));
141       moduleBuilder.clone(*funcOp.getOperation());
142       spirvModules.push_back(spvModule);
143     }
144   });
145 
146   /// Dialect conversion to lower the functions with the spirv::ModuleOps.
147   SPIRVBasicTypeConverter basicTypeConverter;
148   SPIRVTypeConverter typeConverter(&basicTypeConverter);
149   OwningRewritePatternList patterns;
150   patterns.insert<
151       KernelFnConversion,
152       LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
153       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
154       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
155       LaunchConfigConversion<gpu::ThreadIdOp,
156                              spirv::BuiltIn::LocalInvocationId>>(context,
157                                                                  typeConverter);
158   populateStandardToSPIRVPatterns(context, patterns);
159 
160   ConversionTarget target(*context);
161   target.addLegalDialect<spirv::SPIRVDialect>();
162   target.addDynamicallyLegalOp<FuncOp>([&](FuncOp Op) {
163     return basicTypeConverter.isSignatureLegal(Op.getType());
164   });
165 
166   if (failed(applyFullConversion(spirvModules, target, patterns,
167                                  &typeConverter))) {
168     return signalPassFailure();
169   }
170 
171   // After the SPIR-V modules have been generated, some finalization is needed
172   // for the entry functions. For example, adding spv.EntryPoint op,
173   // spv.ExecutionMode op, etc.
174   for (auto *spvModule : spirvModules) {
175     for (auto op :
176          cast<spirv::ModuleOp>(spvModule).getBlock().getOps<FuncOp>()) {
177       if (gpu::GPUDialect::isKernel(op)) {
178         OpBuilder builder(op.getContext());
179         builder.setInsertionPointAfter(op);
180         if (failed(finalizeEntryFunction(op, builder))) {
181           return signalPassFailure();
182         }
183         op.getOperation()->removeAttr(Identifier::get(
184             gpu::GPUDialect::getKernelFuncAttrName(), op.getContext()));
185       }
186     }
187   }
188 }
189 
190 OpPassBase<ModuleOp> *createGPUToSPIRVPass() { return new GPUToSPIRVPass(); }
191 
192 static PassRegistration<GPUToSPIRVPass>
193     pass("convert-gpu-to-spirv", "Convert GPU dialect to SPIR-V dialect");
194