xref: /llvm-project/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp (revision bc29fc937c6cb4a210f80c93c79fc6ed97c801f8)
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