18bfedb3cSKazuaki Ishizaki //===- KernelOutlining.cpp - Implementation of GPU kernel outlining -------===// 260965b46SAlex Zinenko // 330857107SMehdi Amini // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 456222a06SMehdi Amini // See https://llvm.org/LICENSE.txt for license information. 556222a06SMehdi Amini // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 660965b46SAlex Zinenko // 756222a06SMehdi Amini //===----------------------------------------------------------------------===// 860965b46SAlex Zinenko // 960965b46SAlex Zinenko // This file implements the GPU dialect kernel outlining pass. 1060965b46SAlex Zinenko // 1160965b46SAlex Zinenko //===----------------------------------------------------------------------===// 1260965b46SAlex Zinenko 1367d0d7acSMichele Scuttari #include "mlir/Dialect/GPU/Transforms/Passes.h" 1467d0d7acSMichele Scuttari 15c60b897dSRiver Riddle #include "mlir/AsmParser/AsmParser.h" 16abc362a1SJakub Kuderski #include "mlir/Dialect/Arith/IR/Arith.h" 17ace01605SRiver Riddle #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" 1832fe1a8aSDiego Caballero #include "mlir/Dialect/DLTI/DLTI.h" 1936550692SRiver Riddle #include "mlir/Dialect/Func/IR/FuncOps.h" 20d7ef488bSMogball #include "mlir/Dialect/GPU/IR/GPUDialect.h" 21*bc29fc93SPetr Kurapov #include "mlir/Dialect/GPU/Utils/GPUUtils.h" 22e2310704SJulian Gross #include "mlir/Dialect/MemRef/IR/MemRef.h" 2360965b46SAlex Zinenko #include "mlir/IR/Builders.h" 24be575c5dSKrzysztof Drewniak #include "mlir/IR/BuiltinAttributes.h" 254d67b278SJeff Niu #include "mlir/IR/IRMapping.h" 261f971e23SRiver Riddle #include "mlir/IR/Matchers.h" 27b8cd0c14STres Popp #include "mlir/IR/SymbolTable.h" 28edeff6e6SStephan Herhut #include "mlir/Support/LLVM.h" 29283b5e73SStephan Herhut #include "mlir/Transforms/RegionUtils.h" 30be575c5dSKrzysztof Drewniak #include <limits> 3160965b46SAlex Zinenko 3267d0d7acSMichele Scuttari namespace mlir { 3367d0d7acSMichele Scuttari #define GEN_PASS_DEF_GPULAUNCHSINKINDEXCOMPUTATIONS 3467d0d7acSMichele Scuttari #define GEN_PASS_DEF_GPUKERNELOUTLINING 3567d0d7acSMichele Scuttari #include "mlir/Dialect/GPU/Transforms/Passes.h.inc" 3667d0d7acSMichele Scuttari } // namespace mlir 3767d0d7acSMichele Scuttari 3860965b46SAlex Zinenko using namespace mlir; 3960965b46SAlex Zinenko 4060965b46SAlex Zinenko template <typename OpTy> 4160965b46SAlex Zinenko static void createForAllDimensions(OpBuilder &builder, Location loc, 42e62a6956SRiver Riddle SmallVectorImpl<Value> &values) { 43aae51255SMogball for (auto dim : {gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z}) 44aae51255SMogball values.push_back(builder.create<OpTy>(loc, builder.getIndexType(), dim)); 4560965b46SAlex Zinenko } 4660965b46SAlex Zinenko 47edeff6e6SStephan Herhut /// Adds operations generating block/thread ids and grid/block dimensions at the 48edeff6e6SStephan Herhut /// beginning of the `launchFuncOpBody` region. Add mapping from argument in 49edeff6e6SStephan Herhut /// entry block of `launchOpBody`, to the corresponding result value of the 50edeff6e6SStephan Herhut /// added operations. 513f44495dSMaheshRavishankar static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody, 525b33cff3SGuray Ozen Region &launchOpBody, IRMapping &map, 535b33cff3SGuray Ozen bool hasCluster = false) { 546273fa0cSAlex Zinenko OpBuilder builder(loc->getContext()); 553f44495dSMaheshRavishankar Block &firstBlock = launchOpBody.front(); 563f44495dSMaheshRavishankar builder.setInsertionPointToStart(&launchFuncOpBody.front()); 575b33cff3SGuray Ozen SmallVector<Value> indexOps; 585b33cff3SGuray Ozen // The order is important here, as it must match the order of the arguments 596273fa0cSAlex Zinenko createForAllDimensions<gpu::BlockIdOp>(builder, loc, indexOps); 606273fa0cSAlex Zinenko createForAllDimensions<gpu::ThreadIdOp>(builder, loc, indexOps); 616273fa0cSAlex Zinenko createForAllDimensions<gpu::GridDimOp>(builder, loc, indexOps); 626273fa0cSAlex Zinenko createForAllDimensions<gpu::BlockDimOp>(builder, loc, indexOps); 635b33cff3SGuray Ozen if (hasCluster) { 645b33cff3SGuray Ozen createForAllDimensions<gpu::ClusterIdOp>(builder, loc, indexOps); 655b33cff3SGuray Ozen createForAllDimensions<gpu::ClusterDimOp>(builder, loc, indexOps); 665b33cff3SGuray Ozen } 6760965b46SAlex Zinenko // Replace the leading 12 function args with the respective thread/block index 6860965b46SAlex Zinenko // operations. Iterate backwards since args are erased and indices change. 69e4853be2SMehdi Amini for (const auto &indexOp : enumerate(indexOps)) 703f44495dSMaheshRavishankar map.map(firstBlock.getArgument(indexOp.index()), indexOp.value()); 7160965b46SAlex Zinenko } 7260965b46SAlex Zinenko 73edeff6e6SStephan Herhut /// Identifies operations that are beneficial to sink into kernels. These 74edeff6e6SStephan Herhut /// operations may not have side-effects, as otherwise sinking (and hence 75edeff6e6SStephan Herhut /// duplicating them) is not legal. 76d271fc04SIvan Butygin static bool isLikelyAnIndexComputation(Operation *op) { 771f971e23SRiver Riddle return matchPattern(op, m_Constant()) || 781f971e23SRiver Riddle isa<memref::DimOp, arith::SelectOp, arith::CmpIOp>(op); 79edeff6e6SStephan Herhut } 80edeff6e6SStephan Herhut 81edeff6e6SStephan Herhut /// For a given operation `op`, computes whether it is beneficial to sink the 82edeff6e6SStephan Herhut /// operation into the kernel. An operation can be sunk if doing so does not 83edeff6e6SStephan Herhut /// introduce new kernel arguments. Whether a value is already available in the 84edeff6e6SStephan Herhut /// kernel (and hence does not introduce new arguments) is checked by 85366d8435SStephan Herhut /// querying `existingDependencies` and `availableValues`. 86edeff6e6SStephan Herhut /// If an operand is not yet available, we recursively check whether it can be 87edeff6e6SStephan Herhut /// made available by siking its defining op. 88edeff6e6SStephan Herhut /// Operations that are indentified for sinking are added to `beneficiaryOps` in 89366d8435SStephan Herhut /// the order they should appear in the kernel. Furthermore, `availableValues` 90366d8435SStephan Herhut /// is updated with results that will be available after sinking the identified 91edeff6e6SStephan Herhut /// ops. 92a2e2fbbaSIvan Butygin static bool extractBeneficiaryOps( 93a2e2fbbaSIvan Butygin Operation *op, const SetVector<Value> &existingDependencies, 944efb7754SRiver Riddle SetVector<Operation *> &beneficiaryOps, 95a2e2fbbaSIvan Butygin llvm::SmallPtrSetImpl<Value> &availableValues, 96a2e2fbbaSIvan Butygin llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) { 97edeff6e6SStephan Herhut if (beneficiaryOps.count(op)) 98edeff6e6SStephan Herhut return true; 99edeff6e6SStephan Herhut 100edeff6e6SStephan Herhut if (!isSinkingBeneficiary(op)) 101edeff6e6SStephan Herhut return false; 102edeff6e6SStephan Herhut 103edeff6e6SStephan Herhut for (Value operand : op->getOperands()) { 10441b09f4eSKazuaki Ishizaki // It is already visible in the kernel, keep going. 105edeff6e6SStephan Herhut if (availableValues.count(operand)) 106edeff6e6SStephan Herhut continue; 107366d8435SStephan Herhut // Else check whether it can be made available via sinking or already is a 108366d8435SStephan Herhut // dependency. 109edeff6e6SStephan Herhut Operation *definingOp = operand.getDefiningOp(); 110a2e2fbbaSIvan Butygin if ((!definingOp || !extractBeneficiaryOps(definingOp, existingDependencies, 111a2e2fbbaSIvan Butygin beneficiaryOps, availableValues, 112a2e2fbbaSIvan Butygin isSinkingBeneficiary)) && 113366d8435SStephan Herhut !existingDependencies.count(operand)) 114edeff6e6SStephan Herhut return false; 115edeff6e6SStephan Herhut } 116edeff6e6SStephan Herhut // We will sink the operation, mark its results as now available. 117edeff6e6SStephan Herhut beneficiaryOps.insert(op); 118edeff6e6SStephan Herhut for (Value result : op->getResults()) 119edeff6e6SStephan Herhut availableValues.insert(result); 120edeff6e6SStephan Herhut return true; 121abb62668SStephan Herhut } 122abb62668SStephan Herhut 123a2e2fbbaSIvan Butygin LogicalResult mlir::sinkOperationsIntoLaunchOp( 124a2e2fbbaSIvan Butygin gpu::LaunchOp launchOp, 125a2e2fbbaSIvan Butygin llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) { 126a2e2fbbaSIvan Butygin assert(isSinkingBeneficiary); 12710c04f46SRiver Riddle Region &launchOpBody = launchOp.getBody(); 128318ff019SStephan Herhut 1293f44495dSMaheshRavishankar // Identify uses from values defined outside of the scope of the launch 1303f44495dSMaheshRavishankar // operation. 1314efb7754SRiver Riddle SetVector<Value> sinkCandidates; 1323f44495dSMaheshRavishankar getUsedValuesDefinedAbove(launchOpBody, sinkCandidates); 1333f44495dSMaheshRavishankar 1344efb7754SRiver Riddle SetVector<Operation *> toBeSunk; 135366d8435SStephan Herhut llvm::SmallPtrSet<Value, 4> availableValues; 136366d8435SStephan Herhut for (Value operand : sinkCandidates) { 1373f44495dSMaheshRavishankar Operation *operandOp = operand.getDefiningOp(); 138edeff6e6SStephan Herhut if (!operandOp) 1393f44495dSMaheshRavishankar continue; 140a2e2fbbaSIvan Butygin extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues, 141a2e2fbbaSIvan Butygin isSinkingBeneficiary); 142dfd06af5SStephan Herhut } 1433f44495dSMaheshRavishankar 1443f44495dSMaheshRavishankar // Insert operations so that the defs get cloned before uses. 1454d67b278SJeff Niu IRMapping map; 1463f44495dSMaheshRavishankar OpBuilder builder(launchOpBody); 147edeff6e6SStephan Herhut for (Operation *op : toBeSunk) { 148edeff6e6SStephan Herhut Operation *clonedOp = builder.clone(*op, map); 1493f44495dSMaheshRavishankar // Only replace uses within the launch op. 150edeff6e6SStephan Herhut for (auto pair : llvm::zip(op->getResults(), clonedOp->getResults())) 151edeff6e6SStephan Herhut replaceAllUsesInRegionWith(std::get<0>(pair), std::get<1>(pair), 15210c04f46SRiver Riddle launchOp.getBody()); 1533f44495dSMaheshRavishankar } 1543f44495dSMaheshRavishankar return success(); 155dfd06af5SStephan Herhut } 156dfd06af5SStephan Herhut 157be575c5dSKrzysztof Drewniak /// Return the provided KernelDim3 as an array of i32 constants if possible. 158be575c5dSKrzysztof Drewniak static DenseI32ArrayAttr maybeConstantDimsAttr(gpu::KernelDim3 dims) { 159be575c5dSKrzysztof Drewniak SmallVector<int32_t, 3> constants; 160be575c5dSKrzysztof Drewniak MLIRContext *ctx = dims.x.getContext(); 161be575c5dSKrzysztof Drewniak for (Value v : {dims.x, dims.y, dims.z}) { 162be575c5dSKrzysztof Drewniak APInt constValue; 163be575c5dSKrzysztof Drewniak if (!matchPattern(v, m_ConstantInt(&constValue))) 164be575c5dSKrzysztof Drewniak return nullptr; 165be575c5dSKrzysztof Drewniak // In the event someone called for a too-large block or grid dimension, 166be575c5dSKrzysztof Drewniak // don't set bounds as it is likely to cause more confusing behavior. 167be575c5dSKrzysztof Drewniak if (constValue.ugt(std::numeric_limits<uint32_t>::max())) 168be575c5dSKrzysztof Drewniak return nullptr; 169be575c5dSKrzysztof Drewniak constants.push_back( 170be575c5dSKrzysztof Drewniak constValue.getLimitedValue(std::numeric_limits<uint32_t>::max())); 171be575c5dSKrzysztof Drewniak } 172be575c5dSKrzysztof Drewniak return DenseI32ArrayAttr::get(ctx, constants); 173be575c5dSKrzysztof Drewniak } 174be575c5dSKrzysztof Drewniak 175edeff6e6SStephan Herhut /// Outline the `gpu.launch` operation body into a kernel function. Replace 176edeff6e6SStephan Herhut /// `gpu.terminator` operations by `gpu.return` in the generated function. 177be575c5dSKrzysztof Drewniak /// Set block and grid size bounds if known. 1783f44495dSMaheshRavishankar static gpu::GPUFuncOp outlineKernelFuncImpl(gpu::LaunchOp launchOp, 1793f44495dSMaheshRavishankar StringRef kernelFnName, 1804efb7754SRiver Riddle SetVector<Value> &operands) { 18160965b46SAlex Zinenko Location loc = launchOp.getLoc(); 1826273fa0cSAlex Zinenko // Create a builder with no insertion point, insertion will happen separately 1836273fa0cSAlex Zinenko // due to symbol table manipulation. 1846273fa0cSAlex Zinenko OpBuilder builder(launchOp.getContext()); 18510c04f46SRiver Riddle Region &launchOpBody = launchOp.getBody(); 1866273fa0cSAlex Zinenko 187283b5e73SStephan Herhut // Identify uses from values defined outside of the scope of the launch 188283b5e73SStephan Herhut // operation. 1893f44495dSMaheshRavishankar getUsedValuesDefinedAbove(launchOpBody, operands); 190283b5e73SStephan Herhut 1913f44495dSMaheshRavishankar // Create the gpu.func operation. 192283b5e73SStephan Herhut SmallVector<Type, 4> kernelOperandTypes; 193283b5e73SStephan Herhut kernelOperandTypes.reserve(operands.size()); 194283b5e73SStephan Herhut for (Value operand : operands) { 195283b5e73SStephan Herhut kernelOperandTypes.push_back(operand.getType()); 196283b5e73SStephan Herhut } 19760965b46SAlex Zinenko FunctionType type = 1981b97cdf8SRiver Riddle FunctionType::get(launchOp.getContext(), kernelOperandTypes, {}); 19954e96f4fSFabian Mora auto outlinedFunc = builder.create<gpu::GPUFuncOp>( 20054e96f4fSFabian Mora loc, kernelFnName, type, 20154e96f4fSFabian Mora TypeRange(ValueRange(launchOp.getWorkgroupAttributions())), 20254e96f4fSFabian Mora TypeRange(ValueRange(launchOp.getPrivateAttributions()))); 2031ffc1aaaSChristian Sigg outlinedFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), 20460965b46SAlex Zinenko builder.getUnitAttr()); 205be575c5dSKrzysztof Drewniak 206be575c5dSKrzysztof Drewniak // If we can infer bounds on the grid and/or block sizes from the arguments 207be575c5dSKrzysztof Drewniak // to the launch op, propagate them to the generated kernel. This is safe 208be575c5dSKrzysztof Drewniak // because multiple launches with the same body are not deduplicated. 209be575c5dSKrzysztof Drewniak if (auto blockBounds = 210be575c5dSKrzysztof Drewniak maybeConstantDimsAttr(launchOp.getBlockSizeOperandValues())) 21143fd4c49SKrzysztof Drewniak outlinedFunc.setKnownBlockSizeAttr(blockBounds); 212be575c5dSKrzysztof Drewniak if (auto gridBounds = 213be575c5dSKrzysztof Drewniak maybeConstantDimsAttr(launchOp.getGridSizeOperandValues())) 21443fd4c49SKrzysztof Drewniak outlinedFunc.setKnownGridSizeAttr(gridBounds); 215be575c5dSKrzysztof Drewniak 2164d67b278SJeff Niu IRMapping map; 2173f44495dSMaheshRavishankar 2183f44495dSMaheshRavishankar // Map the arguments corresponding to the launch parameters like blockIdx, 2195b33cff3SGuray Ozen // threadIdx, etc. If cluster is present, then we also generate clusterIdx and 2205b33cff3SGuray Ozen // clusterDim. 22110c04f46SRiver Riddle Region &outlinedFuncBody = outlinedFunc.getBody(); 2225b33cff3SGuray Ozen injectGpuIndexOperations(loc, outlinedFuncBody, launchOpBody, map, 2235b33cff3SGuray Ozen launchOp.hasClusterSize()); 2243f44495dSMaheshRavishankar 22554e96f4fSFabian Mora // Map memory attributions from the LaunOp op to the GPUFuncOp attributions. 22654e96f4fSFabian Mora for (const auto &[launchArg, funcArg] : 22754e96f4fSFabian Mora llvm::zip(launchOp.getWorkgroupAttributions(), 22854e96f4fSFabian Mora outlinedFunc.getWorkgroupAttributions())) 22954e96f4fSFabian Mora map.map(launchArg, funcArg); 23054e96f4fSFabian Mora for (const auto &[launchArg, funcArg] : 23154e96f4fSFabian Mora llvm::zip(launchOp.getPrivateAttributions(), 23254e96f4fSFabian Mora outlinedFunc.getPrivateAttributions())) 23354e96f4fSFabian Mora map.map(launchArg, funcArg); 23454e96f4fSFabian Mora 2353f44495dSMaheshRavishankar // Map arguments from gpu.launch region to the arguments of the gpu.func 2363f44495dSMaheshRavishankar // operation. 2373f44495dSMaheshRavishankar Block &entryBlock = outlinedFuncBody.front(); 238e4853be2SMehdi Amini for (const auto &operand : enumerate(operands)) 2393f44495dSMaheshRavishankar map.map(operand.value(), entryBlock.getArgument(operand.index())); 2403f44495dSMaheshRavishankar 2413f44495dSMaheshRavishankar // Clone the region of the gpu.launch operation into the gpu.func operation. 2423f44495dSMaheshRavishankar launchOpBody.cloneInto(&outlinedFuncBody, map); 2433f44495dSMaheshRavishankar 244d566a5cdSMehdi Amini // Replace the terminator op with returns. 245d566a5cdSMehdi Amini for (Block &block : launchOpBody) { 246d566a5cdSMehdi Amini Block *clonedBlock = map.lookup(&block); 247d566a5cdSMehdi Amini auto terminator = dyn_cast<gpu::TerminatorOp>(clonedBlock->getTerminator()); 248d566a5cdSMehdi Amini if (!terminator) 249d566a5cdSMehdi Amini continue; 250d566a5cdSMehdi Amini OpBuilder replacer(terminator); 251d566a5cdSMehdi Amini replacer.create<gpu::ReturnOp>(terminator->getLoc()); 252d566a5cdSMehdi Amini terminator->erase(); 253d566a5cdSMehdi Amini } 2543f44495dSMaheshRavishankar 255d566a5cdSMehdi Amini // Splice now the entry block of the gpu.launch operation at the end of the 256d566a5cdSMehdi Amini // gpu.func entry block and erase the redundant block. 257d566a5cdSMehdi Amini Block *clonedLaunchOpEntry = map.lookup(&launchOpBody.front()); 258d566a5cdSMehdi Amini entryBlock.getOperations().splice(entryBlock.getOperations().end(), 259d566a5cdSMehdi Amini clonedLaunchOpEntry->getOperations()); 260d566a5cdSMehdi Amini clonedLaunchOpEntry->erase(); 261d566a5cdSMehdi Amini 26260965b46SAlex Zinenko return outlinedFunc; 26360965b46SAlex Zinenko } 26460965b46SAlex Zinenko 2653f44495dSMaheshRavishankar gpu::GPUFuncOp mlir::outlineKernelFunc(gpu::LaunchOp launchOp, 2663f44495dSMaheshRavishankar StringRef kernelFnName, 2673f44495dSMaheshRavishankar llvm::SmallVectorImpl<Value> &operands) { 2683f44495dSMaheshRavishankar DenseSet<Value> inputOperandSet; 2693f44495dSMaheshRavishankar inputOperandSet.insert(operands.begin(), operands.end()); 2704efb7754SRiver Riddle SetVector<Value> operandSet(operands.begin(), operands.end()); 2713f44495dSMaheshRavishankar auto funcOp = outlineKernelFuncImpl(launchOp, kernelFnName, operandSet); 2723f44495dSMaheshRavishankar for (auto operand : operandSet) { 2733f44495dSMaheshRavishankar if (!inputOperandSet.count(operand)) 2743f44495dSMaheshRavishankar operands.push_back(operand); 2753f44495dSMaheshRavishankar } 2763f44495dSMaheshRavishankar return funcOp; 2773f44495dSMaheshRavishankar } 2783f44495dSMaheshRavishankar 279edeff6e6SStephan Herhut /// Replace `gpu.launch` operations with an `gpu.launch_func` operation 280edeff6e6SStephan Herhut /// launching `kernelFunc`. The kernel func contains the body of the 281edeff6e6SStephan Herhut /// `gpu.launch` with constant region arguments inlined. 2823f44495dSMaheshRavishankar static void convertToLaunchFuncOp(gpu::LaunchOp launchOp, 283283b5e73SStephan Herhut gpu::GPUFuncOp kernelFunc, 284283b5e73SStephan Herhut ValueRange operands) { 28560965b46SAlex Zinenko OpBuilder builder(launchOp); 28608b63db8SUday Bondhugula // The launch op has an optional dynamic shared memory size. If it doesn't 28708b63db8SUday Bondhugula // exist, we use zero. 28810c04f46SRiver Riddle Value asyncToken = launchOp.getAsyncToken(); 2895b33cff3SGuray Ozen std::optional<gpu::KernelDim3> clusterSize = 2905b33cff3SGuray Ozen launchOp.getClusterSizeOperandValues(); 291f47a38f5SUday Bondhugula auto launchFunc = builder.create<gpu::LaunchFuncOp>( 29260965b46SAlex Zinenko launchOp.getLoc(), kernelFunc, launchOp.getGridSizeOperandValues(), 29310c04f46SRiver Riddle launchOp.getBlockSizeOperandValues(), 29410c04f46SRiver Riddle launchOp.getDynamicSharedMemorySize(), operands, 29510c04f46SRiver Riddle asyncToken ? asyncToken.getType() : nullptr, 2965b33cff3SGuray Ozen launchOp.getAsyncDependencies(), clusterSize); 297f47a38f5SUday Bondhugula launchOp.replaceAllUsesWith(launchFunc); 29860965b46SAlex Zinenko launchOp.erase(); 29960965b46SAlex Zinenko } 30060965b46SAlex Zinenko 30160965b46SAlex Zinenko namespace { 302d271fc04SIvan Butygin /// Pass that moves ops which are likely an index computation into gpu.launch 303d271fc04SIvan Butygin /// body. 304d271fc04SIvan Butygin class GpuLaunchSinkIndexComputationsPass 30567d0d7acSMichele Scuttari : public impl::GpuLaunchSinkIndexComputationsBase< 306d271fc04SIvan Butygin GpuLaunchSinkIndexComputationsPass> { 307d271fc04SIvan Butygin public: 308d271fc04SIvan Butygin void runOnOperation() override { 309d271fc04SIvan Butygin Operation *op = getOperation(); 310d271fc04SIvan Butygin if (op->walk([](gpu::LaunchOp launch) { 311d271fc04SIvan Butygin // Pull in instructions that can be sunk 312d271fc04SIvan Butygin if (failed(sinkOperationsIntoLaunchOp(launch, 313d271fc04SIvan Butygin isLikelyAnIndexComputation))) 314d271fc04SIvan Butygin return WalkResult::interrupt(); 315d271fc04SIvan Butygin 316d271fc04SIvan Butygin return WalkResult::advance(); 317d271fc04SIvan Butygin }).wasInterrupted()) 318d271fc04SIvan Butygin signalPassFailure(); 319d271fc04SIvan Butygin } 320d271fc04SIvan Butygin }; 321d271fc04SIvan Butygin 322b8676da1SChristian Sigg /// Pass that moves the kernel of each LaunchOp into its separate nested module. 323b8676da1SChristian Sigg /// 324b8676da1SChristian Sigg /// This pass moves the kernel code of each LaunchOp into a function created 325b8676da1SChristian Sigg /// inside a nested module. It also creates an external function of the same 326b8676da1SChristian Sigg /// name in the parent module. 327b8676da1SChristian Sigg /// 3289a52ea5cSTres Popp /// The gpu.modules are intended to be compiled to a cubin blob independently in 3299a52ea5cSTres Popp /// a separate pass. The external functions can then be annotated with the 330b8676da1SChristian Sigg /// symbol of the cubin accessor function. 331722f909fSRiver Riddle class GpuKernelOutliningPass 33267d0d7acSMichele Scuttari : public impl::GpuKernelOutliningBase<GpuKernelOutliningPass> { 33360965b46SAlex Zinenko public: 33432fe1a8aSDiego Caballero GpuKernelOutliningPass(StringRef dlStr) { 33532fe1a8aSDiego Caballero if (!dlStr.empty() && !dataLayoutStr.hasValue()) 33632fe1a8aSDiego Caballero dataLayoutStr = dlStr.str(); 33732fe1a8aSDiego Caballero } 33832fe1a8aSDiego Caballero 33932fe1a8aSDiego Caballero GpuKernelOutliningPass(const GpuKernelOutliningPass &other) 340039b969bSMichele Scuttari : GpuKernelOutliningBase(other), dataLayoutSpec(other.dataLayoutSpec) { 34170c463efSDaniil Dudkin dataLayoutStr = other.dataLayoutStr.getValue(); 34232fe1a8aSDiego Caballero } 34332fe1a8aSDiego Caballero 34432fe1a8aSDiego Caballero LogicalResult initialize(MLIRContext *context) override { 34532fe1a8aSDiego Caballero // Initialize the data layout specification from the data layout string. 34632fe1a8aSDiego Caballero if (!dataLayoutStr.empty()) { 34732fe1a8aSDiego Caballero Attribute resultAttr = mlir::parseAttribute(dataLayoutStr, context); 34832fe1a8aSDiego Caballero if (!resultAttr) 34932fe1a8aSDiego Caballero return failure(); 35032fe1a8aSDiego Caballero 3515550c821STres Popp dataLayoutSpec = dyn_cast<DataLayoutSpecInterface>(resultAttr); 35232fe1a8aSDiego Caballero if (!dataLayoutSpec) 35332fe1a8aSDiego Caballero return failure(); 35432fe1a8aSDiego Caballero } 35532fe1a8aSDiego Caballero 35632fe1a8aSDiego Caballero return success(); 35732fe1a8aSDiego Caballero } 35832fe1a8aSDiego Caballero 359722f909fSRiver Riddle void runOnOperation() override { 360722f909fSRiver Riddle SymbolTable symbolTable(getOperation()); 36190d65d32SAlex Zinenko bool modified = false; 3629a3d3c70Sdrazi for (auto func : getOperation().getOps<SymbolOpInterface>()) { 363b8676da1SChristian Sigg // Insert just after the function. 364c4a04059SChristian Sigg Block::iterator insertPt(func->getNextNode()); 3653f44495dSMaheshRavishankar auto funcWalkResult = func.walk([&](gpu::LaunchOp op) { 3664efb7754SRiver Riddle SetVector<Value> operands; 367516d6edeSZhen Wang std::string kernelFnName; 368516d6edeSZhen Wang if (op.getKernelFunc()) { 369516d6edeSZhen Wang kernelFnName = op.getKernelFunc()->getRootReference().str(); 370516d6edeSZhen Wang } else { 371516d6edeSZhen Wang kernelFnName = 372516d6edeSZhen Wang Twine(op->getParentOfType<SymbolOpInterface>().getName(), 373516d6edeSZhen Wang "_kernel") 37458ceae95SRiver Riddle .str(); 375516d6edeSZhen Wang } 3763f44495dSMaheshRavishankar 3773f44495dSMaheshRavishankar gpu::GPUFuncOp outlinedFunc = 3783f44495dSMaheshRavishankar outlineKernelFuncImpl(op, kernelFnName, operands); 379b8676da1SChristian Sigg 38090d65d32SAlex Zinenko // Create nested module and insert outlinedFunc. The module will 38190d65d32SAlex Zinenko // originally get the same name as the function, but may be renamed on 38290d65d32SAlex Zinenko // insertion into the parent module. 383516d6edeSZhen Wang auto kernelModule = createKernelModule(op, outlinedFunc, symbolTable); 384b8cd0c14STres Popp symbolTable.insert(kernelModule, insertPt); 385b8676da1SChristian Sigg 386b8676da1SChristian Sigg // Potentially changes signature, pulling in constants. 387283b5e73SStephan Herhut convertToLaunchFuncOp(op, outlinedFunc, operands.getArrayRef()); 38890d65d32SAlex Zinenko modified = true; 3893f44495dSMaheshRavishankar return WalkResult::advance(); 39060965b46SAlex Zinenko }); 3913f44495dSMaheshRavishankar if (funcWalkResult.wasInterrupted()) 3923f44495dSMaheshRavishankar return signalPassFailure(); 39360965b46SAlex Zinenko } 39490d65d32SAlex Zinenko 39590d65d32SAlex Zinenko // If any new module was inserted in this module, annotate this module as 39690d65d32SAlex Zinenko // a container module. 39790d65d32SAlex Zinenko if (modified) 3981ffc1aaaSChristian Sigg getOperation()->setAttr(gpu::GPUDialect::getContainerModuleAttrName(), 39990d65d32SAlex Zinenko UnitAttr::get(&getContext())); 40060965b46SAlex Zinenko } 40174cdbf59SChristian Sigg 40274cdbf59SChristian Sigg private: 403edeff6e6SStephan Herhut /// Returns a gpu.module containing kernelFunc and all callees (recursive). 404516d6edeSZhen Wang gpu::GPUModuleOp createKernelModule(gpu::LaunchOp gpuLaunchOp, 405516d6edeSZhen Wang gpu::GPUFuncOp kernelFunc, 406b8cd0c14STres Popp const SymbolTable &parentSymbolTable) { 4079a52ea5cSTres Popp // TODO: This code cannot use an OpBuilder because it must be inserted into 4089a52ea5cSTres Popp // a SymbolTable by the caller. SymbolTable needs to be refactored to 4099a52ea5cSTres Popp // prevent manual building of Ops with symbols in code using SymbolTables 4109a52ea5cSTres Popp // and then this needs to use the OpBuilder. 41102b6fb21SMehdi Amini auto *context = getOperation().getContext(); 412bb1d976fSAlex Zinenko OpBuilder builder(context); 413516d6edeSZhen Wang std::string kernelModuleName; 414516d6edeSZhen Wang gpu::GPUModuleOp kernelModule; 415516d6edeSZhen Wang if (gpuLaunchOp.getKernelModule()) { 416516d6edeSZhen Wang kernelModuleName = 417516d6edeSZhen Wang gpuLaunchOp.getKernelModule()->getRootReference().str(); 418516d6edeSZhen Wang kernelModule = 419516d6edeSZhen Wang parentSymbolTable.lookup<gpu::GPUModuleOp>(kernelModuleName); 420516d6edeSZhen Wang } else { 421516d6edeSZhen Wang kernelModuleName = kernelFunc.getName(); 422516d6edeSZhen Wang } 423516d6edeSZhen Wang 424516d6edeSZhen Wang // Check if the module already exists in the symbol table 425516d6edeSZhen Wang if (!kernelModule) { 426516d6edeSZhen Wang // If not found, create a new GPU module 427516d6edeSZhen Wang kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(), 428516d6edeSZhen Wang kernelModuleName); 429516d6edeSZhen Wang } 43032fe1a8aSDiego Caballero 43132fe1a8aSDiego Caballero // If a valid data layout spec was provided, attach it to the kernel module. 43232fe1a8aSDiego Caballero // Otherwise, the default data layout will be used. 43332fe1a8aSDiego Caballero if (dataLayoutSpec) 434e2b658cdSDiego Caballero kernelModule->setAttr(DLTIDialect::kDataLayoutAttrName, dataLayoutSpec); 43532fe1a8aSDiego Caballero 436b8cd0c14STres Popp SymbolTable symbolTable(kernelModule); 437b8cd0c14STres Popp symbolTable.insert(kernelFunc); 43874cdbf59SChristian Sigg 4394562e389SRiver Riddle SmallVector<Operation *, 8> symbolDefWorklist = {kernelFunc}; 4409fbf52e3SMLIR Team while (!symbolDefWorklist.empty()) { 441e8bcc37fSRamkumar Ramachandra if (std::optional<SymbolTable::UseRange> symbolUses = 4429fbf52e3SMLIR Team SymbolTable::getSymbolUses(symbolDefWorklist.pop_back_val())) { 4439fbf52e3SMLIR Team for (SymbolTable::SymbolUse symbolUse : *symbolUses) { 4449b9c647cSRiver Riddle StringRef symbolName = 4455550c821STres Popp cast<FlatSymbolRefAttr>(symbolUse.getSymbolRef()).getValue(); 446b8cd0c14STres Popp if (symbolTable.lookup(symbolName)) 4479fbf52e3SMLIR Team continue; 44874cdbf59SChristian Sigg 4499fbf52e3SMLIR Team Operation *symbolDefClone = 450b8cd0c14STres Popp parentSymbolTable.lookup(symbolName)->clone(); 4519fbf52e3SMLIR Team symbolDefWorklist.push_back(symbolDefClone); 452b8cd0c14STres Popp symbolTable.insert(symbolDefClone); 4539fbf52e3SMLIR Team } 4549fbf52e3SMLIR Team } 45574cdbf59SChristian Sigg } 45674cdbf59SChristian Sigg 45774cdbf59SChristian Sigg return kernelModule; 45874cdbf59SChristian Sigg } 45932fe1a8aSDiego Caballero 46032fe1a8aSDiego Caballero Option<std::string> dataLayoutStr{ 46132fe1a8aSDiego Caballero *this, "data-layout-str", 46232fe1a8aSDiego Caballero llvm::cl::desc("String containing the data layout specification to be " 46332fe1a8aSDiego Caballero "attached to the GPU kernel module")}; 46432fe1a8aSDiego Caballero 46532fe1a8aSDiego Caballero DataLayoutSpecInterface dataLayoutSpec; 46660965b46SAlex Zinenko }; 46760965b46SAlex Zinenko 46860965b46SAlex Zinenko } // namespace 46960965b46SAlex Zinenko 470d271fc04SIvan Butygin std::unique_ptr<Pass> mlir::createGpuLauchSinkIndexComputationsPass() { 471d271fc04SIvan Butygin return std::make_unique<GpuLaunchSinkIndexComputationsPass>(); 472d271fc04SIvan Butygin } 473d271fc04SIvan Butygin 47432fe1a8aSDiego Caballero std::unique_ptr<OperationPass<ModuleOp>> 47532fe1a8aSDiego Caballero mlir::createGpuKernelOutliningPass(StringRef dataLayoutStr) { 47632fe1a8aSDiego Caballero return std::make_unique<GpuKernelOutliningPass>(dataLayoutStr); 47760965b46SAlex Zinenko } 478