1 //===- KernelOutlining.cpp - Implementation of GPU kernel outlining -------===// 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 the GPU dialect kernel outlining pass. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "mlir/Dialect/GPU/Transforms/Passes.h" 14 15 #include "mlir/AsmParser/AsmParser.h" 16 #include "mlir/Dialect/Arith/IR/Arith.h" 17 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" 18 #include "mlir/Dialect/DLTI/DLTI.h" 19 #include "mlir/Dialect/Func/IR/FuncOps.h" 20 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 21 #include "mlir/Dialect/GPU/Transforms/Utils.h" 22 #include "mlir/Dialect/MemRef/IR/MemRef.h" 23 #include "mlir/IR/Builders.h" 24 #include "mlir/IR/BuiltinAttributes.h" 25 #include "mlir/IR/IRMapping.h" 26 #include "mlir/IR/Matchers.h" 27 #include "mlir/IR/SymbolTable.h" 28 #include "mlir/Support/LLVM.h" 29 #include "mlir/Transforms/RegionUtils.h" 30 #include <limits> 31 32 namespace mlir { 33 #define GEN_PASS_DEF_GPULAUNCHSINKINDEXCOMPUTATIONS 34 #define GEN_PASS_DEF_GPUKERNELOUTLINING 35 #include "mlir/Dialect/GPU/Transforms/Passes.h.inc" 36 } // namespace mlir 37 38 using namespace mlir; 39 40 template <typename OpTy> 41 static void createForAllDimensions(OpBuilder &builder, Location loc, 42 SmallVectorImpl<Value> &values) { 43 for (auto dim : {gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z}) 44 values.push_back(builder.create<OpTy>(loc, builder.getIndexType(), dim)); 45 } 46 47 /// Adds operations generating block/thread ids and grid/block dimensions at the 48 /// beginning of the `launchFuncOpBody` region. Add mapping from argument in 49 /// entry block of `launchOpBody`, to the corresponding result value of the 50 /// added operations. 51 static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody, 52 Region &launchOpBody, IRMapping &map, 53 bool hasCluster = false) { 54 OpBuilder builder(loc->getContext()); 55 Block &firstBlock = launchOpBody.front(); 56 builder.setInsertionPointToStart(&launchFuncOpBody.front()); 57 SmallVector<Value> indexOps; 58 // The order is important here, as it must match the order of the arguments 59 createForAllDimensions<gpu::BlockIdOp>(builder, loc, indexOps); 60 createForAllDimensions<gpu::ThreadIdOp>(builder, loc, indexOps); 61 createForAllDimensions<gpu::GridDimOp>(builder, loc, indexOps); 62 createForAllDimensions<gpu::BlockDimOp>(builder, loc, indexOps); 63 if (hasCluster) { 64 createForAllDimensions<gpu::ClusterIdOp>(builder, loc, indexOps); 65 createForAllDimensions<gpu::ClusterDimOp>(builder, loc, indexOps); 66 } 67 // Replace the leading 12 function args with the respective thread/block index 68 // operations. Iterate backwards since args are erased and indices change. 69 for (const auto &indexOp : enumerate(indexOps)) 70 map.map(firstBlock.getArgument(indexOp.index()), indexOp.value()); 71 } 72 73 /// Identifies operations that are beneficial to sink into kernels. These 74 /// operations may not have side-effects, as otherwise sinking (and hence 75 /// duplicating them) is not legal. 76 static bool isLikelyAnIndexComputation(Operation *op) { 77 return matchPattern(op, m_Constant()) || 78 isa<memref::DimOp, arith::SelectOp, arith::CmpIOp>(op); 79 } 80 81 /// For a given operation `op`, computes whether it is beneficial to sink the 82 /// operation into the kernel. An operation can be sunk if doing so does not 83 /// introduce new kernel arguments. Whether a value is already available in the 84 /// kernel (and hence does not introduce new arguments) is checked by 85 /// querying `existingDependencies` and `availableValues`. 86 /// If an operand is not yet available, we recursively check whether it can be 87 /// made available by siking its defining op. 88 /// Operations that are indentified for sinking are added to `beneficiaryOps` in 89 /// the order they should appear in the kernel. Furthermore, `availableValues` 90 /// is updated with results that will be available after sinking the identified 91 /// ops. 92 static bool extractBeneficiaryOps( 93 Operation *op, const SetVector<Value> &existingDependencies, 94 SetVector<Operation *> &beneficiaryOps, 95 llvm::SmallPtrSetImpl<Value> &availableValues, 96 llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) { 97 if (beneficiaryOps.count(op)) 98 return true; 99 100 if (!isSinkingBeneficiary(op)) 101 return false; 102 103 for (Value operand : op->getOperands()) { 104 // It is already visible in the kernel, keep going. 105 if (availableValues.count(operand)) 106 continue; 107 // Else check whether it can be made available via sinking or already is a 108 // dependency. 109 Operation *definingOp = operand.getDefiningOp(); 110 if ((!definingOp || !extractBeneficiaryOps(definingOp, existingDependencies, 111 beneficiaryOps, availableValues, 112 isSinkingBeneficiary)) && 113 !existingDependencies.count(operand)) 114 return false; 115 } 116 // We will sink the operation, mark its results as now available. 117 beneficiaryOps.insert(op); 118 for (Value result : op->getResults()) 119 availableValues.insert(result); 120 return true; 121 } 122 123 LogicalResult mlir::sinkOperationsIntoLaunchOp( 124 gpu::LaunchOp launchOp, 125 llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) { 126 assert(isSinkingBeneficiary); 127 Region &launchOpBody = launchOp.getBody(); 128 129 // Identify uses from values defined outside of the scope of the launch 130 // operation. 131 SetVector<Value> sinkCandidates; 132 getUsedValuesDefinedAbove(launchOpBody, sinkCandidates); 133 134 SetVector<Operation *> toBeSunk; 135 llvm::SmallPtrSet<Value, 4> availableValues; 136 for (Value operand : sinkCandidates) { 137 Operation *operandOp = operand.getDefiningOp(); 138 if (!operandOp) 139 continue; 140 extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues, 141 isSinkingBeneficiary); 142 } 143 144 // Insert operations so that the defs get cloned before uses. 145 IRMapping map; 146 OpBuilder builder(launchOpBody); 147 for (Operation *op : toBeSunk) { 148 Operation *clonedOp = builder.clone(*op, map); 149 // Only replace uses within the launch op. 150 for (auto pair : llvm::zip(op->getResults(), clonedOp->getResults())) 151 replaceAllUsesInRegionWith(std::get<0>(pair), std::get<1>(pair), 152 launchOp.getBody()); 153 } 154 return success(); 155 } 156 157 /// Return the provided KernelDim3 as an array of i32 constants if possible. 158 static DenseI32ArrayAttr maybeConstantDimsAttr(gpu::KernelDim3 dims) { 159 SmallVector<int32_t, 3> constants; 160 MLIRContext *ctx = dims.x.getContext(); 161 for (Value v : {dims.x, dims.y, dims.z}) { 162 APInt constValue; 163 if (!matchPattern(v, m_ConstantInt(&constValue))) 164 return nullptr; 165 // In the event someone called for a too-large block or grid dimension, 166 // don't set bounds as it is likely to cause more confusing behavior. 167 if (constValue.ugt(std::numeric_limits<uint32_t>::max())) 168 return nullptr; 169 constants.push_back( 170 constValue.getLimitedValue(std::numeric_limits<uint32_t>::max())); 171 } 172 return DenseI32ArrayAttr::get(ctx, constants); 173 } 174 175 /// Outline the `gpu.launch` operation body into a kernel function. Replace 176 /// `gpu.terminator` operations by `gpu.return` in the generated function. 177 /// Set block and grid size bounds if known. 178 static gpu::GPUFuncOp outlineKernelFuncImpl(gpu::LaunchOp launchOp, 179 StringRef kernelFnName, 180 SetVector<Value> &operands) { 181 Location loc = launchOp.getLoc(); 182 // Create a builder with no insertion point, insertion will happen separately 183 // due to symbol table manipulation. 184 OpBuilder builder(launchOp.getContext()); 185 Region &launchOpBody = launchOp.getBody(); 186 187 // Identify uses from values defined outside of the scope of the launch 188 // operation. 189 getUsedValuesDefinedAbove(launchOpBody, operands); 190 191 // Create the gpu.func operation. 192 SmallVector<Type, 4> kernelOperandTypes; 193 kernelOperandTypes.reserve(operands.size()); 194 for (Value operand : operands) { 195 kernelOperandTypes.push_back(operand.getType()); 196 } 197 FunctionType type = 198 FunctionType::get(launchOp.getContext(), kernelOperandTypes, {}); 199 auto outlinedFunc = builder.create<gpu::GPUFuncOp>( 200 loc, kernelFnName, type, 201 TypeRange(ValueRange(launchOp.getWorkgroupAttributions())), 202 TypeRange(ValueRange(launchOp.getPrivateAttributions()))); 203 outlinedFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), 204 builder.getUnitAttr()); 205 206 // If we can infer bounds on the grid and/or block sizes from the arguments 207 // to the launch op, propagate them to the generated kernel. This is safe 208 // because multiple launches with the same body are not deduplicated. 209 if (auto blockBounds = 210 maybeConstantDimsAttr(launchOp.getBlockSizeOperandValues())) 211 outlinedFunc->setAttr(gpu::GPUFuncOp::getKnownBlockSizeAttrName(), 212 blockBounds); 213 if (auto gridBounds = 214 maybeConstantDimsAttr(launchOp.getGridSizeOperandValues())) 215 outlinedFunc->setAttr(gpu::GPUFuncOp::getKnownGridSizeAttrName(), 216 gridBounds); 217 218 IRMapping map; 219 220 // Map the arguments corresponding to the launch parameters like blockIdx, 221 // threadIdx, etc. If cluster is present, then we also generate clusterIdx and 222 // clusterDim. 223 Region &outlinedFuncBody = outlinedFunc.getBody(); 224 injectGpuIndexOperations(loc, outlinedFuncBody, launchOpBody, map, 225 launchOp.hasClusterSize()); 226 227 // Map memory attributions from the LaunOp op to the GPUFuncOp attributions. 228 for (const auto &[launchArg, funcArg] : 229 llvm::zip(launchOp.getWorkgroupAttributions(), 230 outlinedFunc.getWorkgroupAttributions())) 231 map.map(launchArg, funcArg); 232 for (const auto &[launchArg, funcArg] : 233 llvm::zip(launchOp.getPrivateAttributions(), 234 outlinedFunc.getPrivateAttributions())) 235 map.map(launchArg, funcArg); 236 237 // Map arguments from gpu.launch region to the arguments of the gpu.func 238 // operation. 239 Block &entryBlock = outlinedFuncBody.front(); 240 for (const auto &operand : enumerate(operands)) 241 map.map(operand.value(), entryBlock.getArgument(operand.index())); 242 243 // Clone the region of the gpu.launch operation into the gpu.func operation. 244 launchOpBody.cloneInto(&outlinedFuncBody, map); 245 246 // Replace the terminator op with returns. 247 for (Block &block : launchOpBody) { 248 Block *clonedBlock = map.lookup(&block); 249 auto terminator = dyn_cast<gpu::TerminatorOp>(clonedBlock->getTerminator()); 250 if (!terminator) 251 continue; 252 OpBuilder replacer(terminator); 253 replacer.create<gpu::ReturnOp>(terminator->getLoc()); 254 terminator->erase(); 255 } 256 257 // Splice now the entry block of the gpu.launch operation at the end of the 258 // gpu.func entry block and erase the redundant block. 259 Block *clonedLaunchOpEntry = map.lookup(&launchOpBody.front()); 260 entryBlock.getOperations().splice(entryBlock.getOperations().end(), 261 clonedLaunchOpEntry->getOperations()); 262 clonedLaunchOpEntry->erase(); 263 264 return outlinedFunc; 265 } 266 267 gpu::GPUFuncOp mlir::outlineKernelFunc(gpu::LaunchOp launchOp, 268 StringRef kernelFnName, 269 llvm::SmallVectorImpl<Value> &operands) { 270 DenseSet<Value> inputOperandSet; 271 inputOperandSet.insert(operands.begin(), operands.end()); 272 SetVector<Value> operandSet(operands.begin(), operands.end()); 273 auto funcOp = outlineKernelFuncImpl(launchOp, kernelFnName, operandSet); 274 for (auto operand : operandSet) { 275 if (!inputOperandSet.count(operand)) 276 operands.push_back(operand); 277 } 278 return funcOp; 279 } 280 281 /// Replace `gpu.launch` operations with an `gpu.launch_func` operation 282 /// launching `kernelFunc`. The kernel func contains the body of the 283 /// `gpu.launch` with constant region arguments inlined. 284 static void convertToLaunchFuncOp(gpu::LaunchOp launchOp, 285 gpu::GPUFuncOp kernelFunc, 286 ValueRange operands) { 287 OpBuilder builder(launchOp); 288 // The launch op has an optional dynamic shared memory size. If it doesn't 289 // exist, we use zero. 290 Value asyncToken = launchOp.getAsyncToken(); 291 std::optional<gpu::KernelDim3> clusterSize = 292 launchOp.getClusterSizeOperandValues(); 293 auto launchFunc = builder.create<gpu::LaunchFuncOp>( 294 launchOp.getLoc(), kernelFunc, launchOp.getGridSizeOperandValues(), 295 launchOp.getBlockSizeOperandValues(), 296 launchOp.getDynamicSharedMemorySize(), operands, 297 asyncToken ? asyncToken.getType() : nullptr, 298 launchOp.getAsyncDependencies(), clusterSize); 299 launchOp.replaceAllUsesWith(launchFunc); 300 launchOp.erase(); 301 } 302 303 namespace { 304 /// Pass that moves ops which are likely an index computation into gpu.launch 305 /// body. 306 class GpuLaunchSinkIndexComputationsPass 307 : public impl::GpuLaunchSinkIndexComputationsBase< 308 GpuLaunchSinkIndexComputationsPass> { 309 public: 310 void runOnOperation() override { 311 Operation *op = getOperation(); 312 if (op->walk([](gpu::LaunchOp launch) { 313 // Pull in instructions that can be sunk 314 if (failed(sinkOperationsIntoLaunchOp(launch, 315 isLikelyAnIndexComputation))) 316 return WalkResult::interrupt(); 317 318 return WalkResult::advance(); 319 }).wasInterrupted()) 320 signalPassFailure(); 321 } 322 }; 323 324 /// Pass that moves the kernel of each LaunchOp into its separate nested module. 325 /// 326 /// This pass moves the kernel code of each LaunchOp into a function created 327 /// inside a nested module. It also creates an external function of the same 328 /// name in the parent module. 329 /// 330 /// The gpu.modules are intended to be compiled to a cubin blob independently in 331 /// a separate pass. The external functions can then be annotated with the 332 /// symbol of the cubin accessor function. 333 class GpuKernelOutliningPass 334 : public impl::GpuKernelOutliningBase<GpuKernelOutliningPass> { 335 public: 336 GpuKernelOutliningPass(StringRef dlStr) { 337 if (!dlStr.empty() && !dataLayoutStr.hasValue()) 338 dataLayoutStr = dlStr.str(); 339 } 340 341 GpuKernelOutliningPass(const GpuKernelOutliningPass &other) 342 : GpuKernelOutliningBase(other), dataLayoutSpec(other.dataLayoutSpec) { 343 dataLayoutStr = other.dataLayoutStr.getValue(); 344 } 345 346 LogicalResult initialize(MLIRContext *context) override { 347 // Initialize the data layout specification from the data layout string. 348 if (!dataLayoutStr.empty()) { 349 Attribute resultAttr = mlir::parseAttribute(dataLayoutStr, context); 350 if (!resultAttr) 351 return failure(); 352 353 dataLayoutSpec = dyn_cast<DataLayoutSpecInterface>(resultAttr); 354 if (!dataLayoutSpec) 355 return failure(); 356 } 357 358 return success(); 359 } 360 361 void runOnOperation() override { 362 SymbolTable symbolTable(getOperation()); 363 bool modified = false; 364 for (auto func : getOperation().getOps<SymbolOpInterface>()) { 365 // Insert just after the function. 366 Block::iterator insertPt(func->getNextNode()); 367 auto funcWalkResult = func.walk([&](gpu::LaunchOp op) { 368 SetVector<Value> operands; 369 std::string kernelFnName = 370 Twine(op->getParentOfType<SymbolOpInterface>().getName(), "_kernel") 371 .str(); 372 373 gpu::GPUFuncOp outlinedFunc = 374 outlineKernelFuncImpl(op, kernelFnName, operands); 375 376 // Create nested module and insert outlinedFunc. The module will 377 // originally get the same name as the function, but may be renamed on 378 // insertion into the parent module. 379 auto kernelModule = createKernelModule(outlinedFunc, symbolTable); 380 symbolTable.insert(kernelModule, insertPt); 381 382 // Potentially changes signature, pulling in constants. 383 convertToLaunchFuncOp(op, outlinedFunc, operands.getArrayRef()); 384 modified = true; 385 return WalkResult::advance(); 386 }); 387 if (funcWalkResult.wasInterrupted()) 388 return signalPassFailure(); 389 } 390 391 // If any new module was inserted in this module, annotate this module as 392 // a container module. 393 if (modified) 394 getOperation()->setAttr(gpu::GPUDialect::getContainerModuleAttrName(), 395 UnitAttr::get(&getContext())); 396 } 397 398 private: 399 /// Returns a gpu.module containing kernelFunc and all callees (recursive). 400 gpu::GPUModuleOp createKernelModule(gpu::GPUFuncOp kernelFunc, 401 const SymbolTable &parentSymbolTable) { 402 // TODO: This code cannot use an OpBuilder because it must be inserted into 403 // a SymbolTable by the caller. SymbolTable needs to be refactored to 404 // prevent manual building of Ops with symbols in code using SymbolTables 405 // and then this needs to use the OpBuilder. 406 auto *context = getOperation().getContext(); 407 OpBuilder builder(context); 408 auto kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(), 409 kernelFunc.getName()); 410 411 // If a valid data layout spec was provided, attach it to the kernel module. 412 // Otherwise, the default data layout will be used. 413 if (dataLayoutSpec) 414 kernelModule->setAttr(DLTIDialect::kDataLayoutAttrName, dataLayoutSpec); 415 416 SymbolTable symbolTable(kernelModule); 417 symbolTable.insert(kernelFunc); 418 419 SmallVector<Operation *, 8> symbolDefWorklist = {kernelFunc}; 420 while (!symbolDefWorklist.empty()) { 421 if (std::optional<SymbolTable::UseRange> symbolUses = 422 SymbolTable::getSymbolUses(symbolDefWorklist.pop_back_val())) { 423 for (SymbolTable::SymbolUse symbolUse : *symbolUses) { 424 StringRef symbolName = 425 cast<FlatSymbolRefAttr>(symbolUse.getSymbolRef()).getValue(); 426 if (symbolTable.lookup(symbolName)) 427 continue; 428 429 Operation *symbolDefClone = 430 parentSymbolTable.lookup(symbolName)->clone(); 431 symbolDefWorklist.push_back(symbolDefClone); 432 symbolTable.insert(symbolDefClone); 433 } 434 } 435 } 436 437 return kernelModule; 438 } 439 440 Option<std::string> dataLayoutStr{ 441 *this, "data-layout-str", 442 llvm::cl::desc("String containing the data layout specification to be " 443 "attached to the GPU kernel module")}; 444 445 DataLayoutSpecInterface dataLayoutSpec; 446 }; 447 448 } // namespace 449 450 std::unique_ptr<Pass> mlir::createGpuLauchSinkIndexComputationsPass() { 451 return std::make_unique<GpuLaunchSinkIndexComputationsPass>(); 452 } 453 454 std::unique_ptr<OperationPass<ModuleOp>> 455 mlir::createGpuKernelOutliningPass(StringRef dataLayoutStr) { 456 return std::make_unique<GpuKernelOutliningPass>(dataLayoutStr); 457 } 458