MLIR  22.0.0git
GPUToSPIRVPass.cpp
Go to the documentation of this file.
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 
15 
26 #include "mlir/IR/PatternMatch.h"
27 
28 namespace mlir {
29 #define GEN_PASS_DEF_CONVERTGPUTOSPIRV
30 #include "mlir/Conversion/Passes.h.inc"
31 } // namespace mlir
32 
33 using namespace mlir;
34 
35 namespace {
36 /// Pass to lower GPU Dialect to SPIR-V. The pass only converts the gpu.func ops
37 /// inside gpu.module ops. i.e., the function that are referenced in
38 /// gpu.launch_func ops. For each such function
39 ///
40 /// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp
41 /// (the original function is still needed by the gpu::LaunchKernelOp, so cannot
42 /// replace it).
43 ///
44 /// 2) Lower the body of the spirv::ModuleOp.
45 struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> {
46  explicit GPUToSPIRVPass(bool mapMemorySpace)
47  : mapMemorySpace(mapMemorySpace) {}
48  void runOnOperation() override;
49 
50 private:
51  /// Queries the target environment from 'targets' attribute of the given
52  /// `moduleOp`.
53  spirv::TargetEnvAttr lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp);
54 
55  /// Queries the target environment from 'targets' attribute of the given
56  /// `moduleOp` or returns target environment as returned by
57  /// `spirv::lookupTargetEnvOrDefault` if not provided by 'targets'.
58  spirv::TargetEnvAttr lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp);
59  bool mapMemorySpace;
60 };
61 
63 GPUToSPIRVPass::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 
74 GPUToSPIRVPass::lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp) {
75  if (spirv::TargetEnvAttr targetEnvAttr = lookupTargetEnvInTargets(moduleOp))
76  return targetEnvAttr;
77 
78  return spirv::lookupTargetEnvOrDefault(moduleOp);
79 }
80 
81 void 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 GPU
96  // launch op still needs the original GPU kernel module.
97  // For Vulkan Shader capabilities, we insert the newly converted SPIR-V
98  // module right after the original GPU module, as that's the expectation of
99  // 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-V
101  // module inside the original GPU module, as that's the expectaion of the
102  // 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 different
112  // 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))
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 =
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 
144  options.use64bitIndex = this->use64bitIndex;
145  SPIRVTypeConverter typeConverter(targetAttr, options);
147 
148  RewritePatternSet patterns(context);
149  populateGPUToSPIRVPatterns(typeConverter, patterns);
151  patterns);
152 
153  // TODO: Change SPIR-V conversion to be progressive and remove the following
154  // patterns.
155  ScfToSPIRVContext scfContext;
156  populateSCFToSPIRVPatterns(typeConverter, scfContext, 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 be
167  // replaced with an empty func.func op with the same arguments as the gpu.func
168  // 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 } // namespace
188 
189 std::unique_ptr<OperationPass<ModuleOp>>
190 mlir::createConvertGPUToSPIRVPass(bool mapMemorySpace) {
191  return std::make_unique<GPUToSPIRVPass>(mapMemorySpace);
192 }
static MLIRContext * getContext(OpFoldResult val)
static llvm::ManagedStatic< PassManagerOptions > options
Attributes are known-constant values of operations.
Definition: Attributes.h:25
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:63
This class helps build Operations.
Definition: Builders.h:207
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
static std::unique_ptr< SPIRVConversionTarget > get(spirv::TargetEnvAttr targetAttr)
Creates a SPIR-V conversion target for the given target environment.
Type conversion from builtin types to SPIR-V types for shader interface.
static WalkResult advance()
Definition: WalkResult.h:47
Type converter for converting numeric MemRef memory spaces into SPIR-V symbolic ones.
Definition: MemRefToSPIRV.h:48
An attribute that specifies the target version, allowed extensions and capabilities,...
A wrapper class around a spirv::TargetEnvAttr to provide query methods for allowed version/capabiliti...
Definition: TargetAndABI.h:29
void populateArithToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition: Remarks.h:491
std::unique_ptr< ConversionTarget > getMemorySpaceToStorageClassTarget(MLIRContext &)
Creates the target that populates legality of ops with MemRef types.
TargetEnvAttr lookupTargetEnvOrDefault(Operation *op)
Queries the target environment recursively from enclosing symbol table ops containing the given op or...
std::function< std::optional< spirv::StorageClass >(Attribute)> MemorySpaceToStorageClassMap
Mapping from numeric MemRef memory spaces into SPIR-V symbolic ones.
Definition: MemRefToSPIRV.h:26
void convertMemRefTypesAndAttrs(Operation *op, MemorySpaceToStorageClassConverter &typeConverter)
Converts all MemRef types and attributes in the op, as decided by the typeConverter.
std::optional< spirv::StorageClass > mapMemorySpaceToOpenCLStorageClass(Attribute)
Maps MemRef memory spaces to storage classes for OpenCL-flavored SPIR-V using the default rule.
std::optional< spirv::StorageClass > mapMemorySpaceToVulkanStorageClass(Attribute)
Maps MemRef memory spaces to storage classes for Vulkan-flavored SPIR-V using the default rule.
Include the generated interface declarations.
std::unique_ptr< OperationPass< ModuleOp > > createConvertGPUToSPIRVPass(bool mapMemorySpace=true)
Creates a pass to convert GPU kernel ops to corresponding SPIR-V ops.
LogicalResult applyFullConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, ConversionConfig config=ConversionConfig())
Apply a complete conversion on the given operations, and all nested operations.
void populateFuncToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating Func ops to SPIR-V ops.
Definition: FuncToSPIRV.cpp:87
void populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Collect a set of patterns to convert WMMA ops from GPU dialect to SPIRV, using the KHR Cooperative Ma...
const FrozenRewritePatternSet & patterns
void populateSCFToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, ScfToSPIRVContext &scfToSPIRVContext, RewritePatternSet &patterns)
Collects a set of patterns to lower from scf.for, scf.if, and loop.terminator to CFG operations withi...
Definition: SCFToSPIRV.cpp:448
void populateGPUToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating GPU Ops to SPIR-V ops.
Definition: GPUToSPIRV.cpp:830
void populateMMAToSPIRVCoopMatrixTypeConversion(SPIRVTypeConverter &typeConverter)
Adds MMAMatrixType conversions to SPIR-V cooperative matrix KHR type conversion to the type converter...
void populateMemRefToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating MemRef ops to SPIR-V ops.
void populateVectorToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating Vector Ops to SPIR-V ops.