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 spirv.module operation. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h" 15 16 #include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h" 17 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h" 18 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h" 19 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h" 20 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h" 21 #include "mlir/Conversion/VectorToSPIRV/VectorToSPIRV.h" 22 #include "mlir/Dialect/Func/IR/FuncOps.h" 23 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 24 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" 25 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" 26 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h" 27 #include "mlir/IR/PatternMatch.h" 28 29 namespace mlir { 30 #define GEN_PASS_DEF_CONVERTGPUTOSPIRV 31 #include "mlir/Conversion/Passes.h.inc" 32 } // namespace mlir 33 34 using namespace mlir; 35 36 namespace { 37 /// Pass to lower GPU Dialect to SPIR-V. The pass only converts the gpu.func ops 38 /// inside gpu.module ops. i.e., the function that are referenced in 39 /// gpu.launch_func ops. For each such function 40 /// 41 /// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp 42 /// (the original function is still needed by the gpu::LaunchKernelOp, so cannot 43 /// replace it). 44 /// 45 /// 2) Lower the body of the spirv::ModuleOp. 46 struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> { 47 explicit GPUToSPIRVPass(bool mapMemorySpace) 48 : mapMemorySpace(mapMemorySpace) {} 49 void runOnOperation() override; 50 51 private: 52 bool mapMemorySpace; 53 }; 54 55 void GPUToSPIRVPass::runOnOperation() { 56 MLIRContext *context = &getContext(); 57 ModuleOp module = getOperation(); 58 59 SmallVector<Operation *, 1> gpuModules; 60 OpBuilder builder(context); 61 62 auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) { 63 Operation *gpuModule = moduleOp.getOperation(); 64 auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule); 65 spirv::TargetEnv targetEnv(targetAttr); 66 return targetEnv.allows(spirv::Capability::Kernel); 67 }; 68 69 module.walk([&](gpu::GPUModuleOp moduleOp) { 70 // Clone each GPU kernel module for conversion, given that the GPU 71 // launch op still needs the original GPU kernel module. 72 // For Vulkan Shader capabilities, we insert the newly converted SPIR-V 73 // module right after the original GPU module, as that's the expectation of 74 // the in-tree SPIR-V CPU runner (the Vulkan runner does not use this pass). 75 // For OpenCL Kernel capabilities, we insert the newly converted SPIR-V 76 // module inside the original GPU module, as that's the expectaion of the 77 // normal GPU compilation pipeline. 78 if (targetEnvSupportsKernelCapability(moduleOp)) { 79 builder.setInsertionPointToStart(moduleOp.getBody()); 80 } else { 81 builder.setInsertionPoint(moduleOp.getOperation()); 82 } 83 gpuModules.push_back(builder.clone(*moduleOp.getOperation())); 84 }); 85 86 // Run conversion for each module independently as they can have different 87 // TargetEnv attributes. 88 for (Operation *gpuModule : gpuModules) { 89 spirv::TargetEnvAttr targetAttr = 90 spirv::lookupTargetEnvOrDefault(gpuModule); 91 92 // Map MemRef memory space to SPIR-V storage class first if requested. 93 if (mapMemorySpace) { 94 spirv::MemorySpaceToStorageClassMap memorySpaceMap = 95 targetEnvSupportsKernelCapability( 96 dyn_cast<gpu::GPUModuleOp>(gpuModule)) 97 ? spirv::mapMemorySpaceToOpenCLStorageClass 98 : spirv::mapMemorySpaceToVulkanStorageClass; 99 spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap); 100 spirv::convertMemRefTypesAndAttrs(gpuModule, converter); 101 102 // Check if there are any illegal ops remaining. 103 std::unique_ptr<ConversionTarget> target = 104 spirv::getMemorySpaceToStorageClassTarget(*context); 105 gpuModule->walk([&target, this](Operation *childOp) { 106 if (target->isIllegal(childOp)) { 107 childOp->emitOpError("failed to legalize memory space"); 108 signalPassFailure(); 109 return WalkResult::interrupt(); 110 } 111 return WalkResult::advance(); 112 }); 113 } 114 115 std::unique_ptr<ConversionTarget> target = 116 SPIRVConversionTarget::get(targetAttr); 117 118 SPIRVConversionOptions options; 119 options.use64bitIndex = this->use64bitIndex; 120 SPIRVTypeConverter typeConverter(targetAttr, options); 121 populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter); 122 123 RewritePatternSet patterns(context); 124 populateGPUToSPIRVPatterns(typeConverter, patterns); 125 populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(typeConverter, 126 patterns); 127 128 // TODO: Change SPIR-V conversion to be progressive and remove the following 129 // patterns. 130 ScfToSPIRVContext scfContext; 131 populateSCFToSPIRVPatterns(typeConverter, scfContext, patterns); 132 mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns); 133 populateMemRefToSPIRVPatterns(typeConverter, patterns); 134 populateFuncToSPIRVPatterns(typeConverter, patterns); 135 populateVectorToSPIRVPatterns(typeConverter, patterns); 136 137 if (failed(applyFullConversion(gpuModule, *target, std::move(patterns)))) 138 return signalPassFailure(); 139 } 140 141 // For OpenCL, the gpu.func op in the original gpu.module op needs to be 142 // replaced with an empty func.func op with the same arguments as the gpu.func 143 // op. The func.func op needs gpu.kernel attribute set. 144 module.walk([&](gpu::GPUModuleOp moduleOp) { 145 if (targetEnvSupportsKernelCapability(moduleOp)) { 146 moduleOp.walk([&](gpu::GPUFuncOp funcOp) { 147 builder.setInsertionPoint(funcOp); 148 auto newFuncOp = builder.create<func::FuncOp>( 149 funcOp.getLoc(), funcOp.getName(), funcOp.getFunctionType()); 150 auto entryBlock = newFuncOp.addEntryBlock(); 151 builder.setInsertionPointToEnd(entryBlock); 152 builder.create<func::ReturnOp>(funcOp.getLoc()); 153 newFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), 154 builder.getUnitAttr()); 155 funcOp.erase(); 156 }); 157 } 158 }); 159 } 160 161 } // namespace 162 163 std::unique_ptr<OperationPass<ModuleOp>> 164 mlir::createConvertGPUToSPIRVPass(bool mapMemorySpace) { 165 return std::make_unique<GPUToSPIRVPass>(mapMemorySpace); 166 } 167