193 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a pass to convert a kernel function in the GPU Dialect10// 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/SPIRVOps.h"25#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"26#include "mlir/IR/PatternMatch.h"27 28namespace mlir {29#define GEN_PASS_DEF_CONVERTGPUTOSPIRV30#include "mlir/Conversion/Passes.h.inc"31} // namespace mlir32 33using namespace mlir;34 35namespace {36/// Pass to lower GPU Dialect to SPIR-V. The pass only converts the gpu.func ops37/// inside gpu.module ops. i.e., the function that are referenced in38/// gpu.launch_func ops. For each such function39///40/// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp41/// (the original function is still needed by the gpu::LaunchKernelOp, so cannot42/// replace it).43///44/// 2) Lower the body of the spirv::ModuleOp.45struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> {46 explicit GPUToSPIRVPass(bool mapMemorySpace)47 : mapMemorySpace(mapMemorySpace) {}48 void runOnOperation() override;49 50private:51 /// Queries the target environment from 'targets' attribute of the given52 /// `moduleOp`.53 spirv::TargetEnvAttr lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp);54 55 /// Queries the target environment from 'targets' attribute of the given56 /// `moduleOp` or returns target environment as returned by57 /// `spirv::lookupTargetEnvOrDefault` if not provided by 'targets'.58 spirv::TargetEnvAttr lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp);59 bool mapMemorySpace;60};61 62spirv::TargetEnvAttr63GPUToSPIRVPass::lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp) {64 if (ArrayAttr targets = moduleOp.getTargetsAttr()) {65 for (Attribute targetAttr : targets)66 if (auto spirvTargetEnvAttr = dyn_cast<spirv::TargetEnvAttr>(targetAttr))67 return spirvTargetEnvAttr;68 }69 70 return {};71}72 73spirv::TargetEnvAttr74GPUToSPIRVPass::lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp) {75 if (spirv::TargetEnvAttr targetEnvAttr = lookupTargetEnvInTargets(moduleOp))76 return targetEnvAttr;77 78 return spirv::lookupTargetEnvOrDefault(moduleOp);79}80 81void GPUToSPIRVPass::runOnOperation() {82 MLIRContext *context = &getContext();83 ModuleOp module = getOperation();84 85 SmallVector<Operation *, 1> gpuModules;86 OpBuilder builder(context);87 88 auto targetEnvSupportsKernelCapability = [this](gpu::GPUModuleOp moduleOp) {89 auto targetAttr = lookupTargetEnvOrDefault(moduleOp);90 spirv::TargetEnv targetEnv(targetAttr);91 return targetEnv.allows(spirv::Capability::Kernel);92 };93 94 module.walk([&](gpu::GPUModuleOp moduleOp) {95 // Clone each GPU kernel module for conversion, given that the GPU96 // launch op still needs the original GPU kernel module.97 // For Vulkan Shader capabilities, we insert the newly converted SPIR-V98 // module right after the original GPU module, as that's the expectation of99 // the in-tree SPIR-V CPU runner (the Vulkan runner does not use this pass).100 // For OpenCL Kernel capabilities, we insert the newly converted SPIR-V101 // module inside the original GPU module, as that's the expectaion of the102 // normal GPU compilation pipeline.103 if (targetEnvSupportsKernelCapability(moduleOp)) {104 builder.setInsertionPointToStart(moduleOp.getBody());105 } else {106 builder.setInsertionPoint(moduleOp.getOperation());107 }108 gpuModules.push_back(builder.clone(*moduleOp.getOperation()));109 });110 111 // Run conversion for each module independently as they can have different112 // TargetEnv attributes.113 for (Operation *gpuModule : gpuModules) {114 spirv::TargetEnvAttr targetAttr =115 lookupTargetEnvOrDefault(cast<gpu::GPUModuleOp>(gpuModule));116 117 // Map MemRef memory space to SPIR-V storage class first if requested.118 if (mapMemorySpace) {119 spirv::MemorySpaceToStorageClassMap memorySpaceMap =120 targetEnvSupportsKernelCapability(121 dyn_cast<gpu::GPUModuleOp>(gpuModule))122 ? spirv::mapMemorySpaceToOpenCLStorageClass123 : spirv::mapMemorySpaceToVulkanStorageClass;124 spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);125 spirv::convertMemRefTypesAndAttrs(gpuModule, converter);126 127 // Check if there are any illegal ops remaining.128 std::unique_ptr<ConversionTarget> target =129 spirv::getMemorySpaceToStorageClassTarget(*context);130 gpuModule->walk([&target, this](Operation *childOp) {131 if (target->isIllegal(childOp)) {132 childOp->emitOpError("failed to legalize memory space");133 signalPassFailure();134 return WalkResult::interrupt();135 }136 return WalkResult::advance();137 });138 }139 140 std::unique_ptr<ConversionTarget> target =141 SPIRVConversionTarget::get(targetAttr);142 143 SPIRVConversionOptions options;144 options.use64bitIndex = this->use64bitIndex;145 SPIRVTypeConverter typeConverter(targetAttr, options);146 populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter);147 148 RewritePatternSet patterns(context);149 populateGPUToSPIRVPatterns(typeConverter, patterns);150 populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(typeConverter,151 patterns);152 153 // TODO: Change SPIR-V conversion to be progressive and remove the following154 // patterns.155 ScfToSPIRVContext scfContext;156 populateSCFToSPIRVPatterns(typeConverter, scfContext, patterns);157 mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns);158 populateMemRefToSPIRVPatterns(typeConverter, patterns);159 populateFuncToSPIRVPatterns(typeConverter, patterns);160 populateVectorToSPIRVPatterns(typeConverter, patterns);161 162 if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))163 return signalPassFailure();164 }165 166 // For OpenCL, the gpu.func op in the original gpu.module op needs to be167 // replaced with an empty func.func op with the same arguments as the gpu.func168 // op. The func.func op needs gpu.kernel attribute set.169 module.walk([&](gpu::GPUModuleOp moduleOp) {170 if (targetEnvSupportsKernelCapability(moduleOp)) {171 moduleOp.walk([&](gpu::GPUFuncOp funcOp) {172 builder.setInsertionPoint(funcOp);173 auto newFuncOp =174 func::FuncOp::create(builder, funcOp.getLoc(), funcOp.getName(),175 funcOp.getFunctionType());176 auto entryBlock = newFuncOp.addEntryBlock();177 builder.setInsertionPointToEnd(entryBlock);178 func::ReturnOp::create(builder, funcOp.getLoc());179 newFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(),180 builder.getUnitAttr());181 funcOp.erase();182 });183 }184 });185}186 187} // namespace188 189std::unique_ptr<OperationPass<ModuleOp>>190mlir::createConvertGPUToSPIRVPass(bool mapMemorySpace) {191 return std::make_unique<GPUToSPIRVPass>(mapMemorySpace);192}193