33 #define GEN_PASS_DEF_GPULAUNCHSINKINDEXCOMPUTATIONS
34 #define GEN_PASS_DEF_GPUKERNELOUTLINING
35 #include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
40 template <
typename OpTy>
43 for (
auto dim : {gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z})
53 bool hasCluster =
false) {
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);
64 createForAllDimensions<gpu::ClusterIdOp>(builder, loc, indexOps);
65 createForAllDimensions<gpu::ClusterDimOp>(builder, loc, indexOps);
69 for (
const auto &indexOp :
enumerate(indexOps))
78 isa<memref::DimOp, arith::SelectOp, arith::CmpIOp>(op);
97 if (beneficiaryOps.count(op))
100 if (!isSinkingBeneficiary(op))
105 if (availableValues.count(operand))
109 Operation *definingOp = operand.getDefiningOp();
111 beneficiaryOps, availableValues,
112 isSinkingBeneficiary)) &&
113 !existingDependencies.count(operand))
117 beneficiaryOps.insert(op);
119 availableValues.insert(result);
124 gpu::LaunchOp launchOp,
126 assert(isSinkingBeneficiary);
127 Region &launchOpBody = launchOp.getBody();
136 for (
Value operand : sinkCandidates) {
137 Operation *operandOp = operand.getDefiningOp();
141 isSinkingBeneficiary);
161 for (
Value v : {dims.
x, dims.
y, dims.
z}) {
179 StringRef kernelFnName,
184 OpBuilder builder(launchOp.getContext());
185 Region &launchOpBody = launchOp.getBody();
193 kernelOperandTypes.reserve(operands.size());
194 for (
Value operand : operands) {
195 kernelOperandTypes.push_back(operand.getType());
199 auto outlinedFunc = builder.
create<gpu::GPUFuncOp>(
200 loc, kernelFnName, type,
203 outlinedFunc->
setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
209 if (
auto blockBounds =
211 outlinedFunc.setKnownBlockSizeAttr(blockBounds);
212 if (
auto gridBounds =
214 outlinedFunc.setKnownGridSizeAttr(gridBounds);
221 Region &outlinedFuncBody = outlinedFunc.getBody();
223 launchOp.hasClusterSize());
226 for (
const auto &[launchArg, funcArg] :
227 llvm::zip(launchOp.getWorkgroupAttributions(),
228 outlinedFunc.getWorkgroupAttributions()))
229 map.
map(launchArg, funcArg);
230 for (
const auto &[launchArg, funcArg] :
231 llvm::zip(launchOp.getPrivateAttributions(),
232 outlinedFunc.getPrivateAttributions()))
233 map.
map(launchArg, funcArg);
238 for (
const auto &operand :
enumerate(operands))
239 map.
map(operand.value(), entryBlock.
getArgument(operand.index()));
242 launchOpBody.
cloneInto(&outlinedFuncBody, map);
245 for (
Block &block : launchOpBody) {
247 auto terminator = dyn_cast<gpu::TerminatorOp>(clonedBlock->
getTerminator());
251 replacer.
create<gpu::ReturnOp>(terminator->getLoc());
260 clonedLaunchOpEntry->
erase();
266 StringRef kernelFnName,
269 inputOperandSet.insert(operands.begin(), operands.end());
272 for (
auto operand : operandSet) {
273 if (!inputOperandSet.count(operand))
274 operands.push_back(operand);
283 gpu::GPUFuncOp kernelFunc,
288 Value asyncToken = launchOp.getAsyncToken();
289 std::optional<gpu::KernelDim3> clusterSize =
290 launchOp.getClusterSizeOperandValues();
291 auto launchFunc = builder.
create<gpu::LaunchFuncOp>(
292 launchOp.getLoc(), kernelFunc, launchOp.getGridSizeOperandValues(),
293 launchOp.getBlockSizeOperandValues(),
294 launchOp.getDynamicSharedMemorySize(), operands,
295 asyncToken ? asyncToken.
getType() :
nullptr,
296 launchOp.getAsyncDependencies(), clusterSize);
304 class GpuLaunchSinkIndexComputationsPass
305 :
public impl::GpuLaunchSinkIndexComputationsBase<
306 GpuLaunchSinkIndexComputationsPass> {
308 void runOnOperation()
override {
310 if (op->
walk([](gpu::LaunchOp launch) {
312 if (failed(sinkOperationsIntoLaunchOp(launch,
313 isLikelyAnIndexComputation)))
314 return WalkResult::interrupt();
316 return WalkResult::advance();
331 class GpuKernelOutliningPass
332 :
public impl::GpuKernelOutliningBase<GpuKernelOutliningPass> {
334 GpuKernelOutliningPass(StringRef dlStr) {
335 if (!dlStr.empty() && !dataLayoutStr.hasValue())
336 dataLayoutStr = dlStr.str();
339 GpuKernelOutliningPass(
const GpuKernelOutliningPass &other)
340 : GpuKernelOutliningBase(other), dataLayoutSpec(other.dataLayoutSpec) {
341 dataLayoutStr = other.dataLayoutStr.getValue();
344 LogicalResult initialize(
MLIRContext *context)
override {
346 if (!dataLayoutStr.empty()) {
351 dataLayoutSpec = dyn_cast<DataLayoutSpecInterface>(resultAttr);
359 void runOnOperation()
override {
361 bool modified =
false;
362 for (
auto func : getOperation().getOps<SymbolOpInterface>()) {
365 auto funcWalkResult = func.walk([&](gpu::LaunchOp op) {
367 std::string kernelFnName =
371 gpu::GPUFuncOp outlinedFunc =
377 auto kernelModule = createKernelModule(outlinedFunc, symbolTable);
378 symbolTable.insert(kernelModule, insertPt);
385 if (funcWalkResult.wasInterrupted())
386 return signalPassFailure();
392 getOperation()->setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
398 gpu::GPUModuleOp createKernelModule(gpu::GPUFuncOp kernelFunc,
404 auto *context = getOperation().getContext();
406 auto kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(),
407 kernelFunc.getName());
412 kernelModule->setAttr(DLTIDialect::kDataLayoutAttrName, dataLayoutSpec);
415 symbolTable.insert(kernelFunc);
418 while (!symbolDefWorklist.empty()) {
419 if (std::optional<SymbolTable::UseRange> symbolUses =
422 StringRef symbolName =
423 cast<FlatSymbolRefAttr>(symbolUse.getSymbolRef()).getValue();
424 if (symbolTable.lookup(symbolName))
429 symbolDefWorklist.push_back(symbolDefClone);
430 symbolTable.insert(symbolDefClone);
438 Option<std::string> dataLayoutStr{
439 *
this,
"data-layout-str",
440 llvm::cl::desc(
"String containing the data layout specification to be "
441 "attached to the GPU kernel module")};
443 DataLayoutSpecInterface dataLayoutSpec;
449 return std::make_unique<GpuLaunchSinkIndexComputationsPass>();
452 std::unique_ptr<OperationPass<ModuleOp>>
454 return std::make_unique<GpuKernelOutliningPass>(dataLayoutStr);
static MLIRContext * getContext(OpFoldResult val)
static DenseI32ArrayAttr maybeConstantDimsAttr(gpu::KernelDim3 dims)
Return the provided KernelDim3 as an array of i32 constants if possible.
static gpu::GPUFuncOp outlineKernelFuncImpl(gpu::LaunchOp launchOp, StringRef kernelFnName, SetVector< Value > &operands)
Outline the gpu.launch operation body into a kernel function.
static bool isLikelyAnIndexComputation(Operation *op)
Identifies operations that are beneficial to sink into kernels.
static void convertToLaunchFuncOp(gpu::LaunchOp launchOp, gpu::GPUFuncOp kernelFunc, ValueRange operands)
Replace gpu.launch operations with an gpu.launch_func operation launching kernelFunc.
static void createForAllDimensions(OpBuilder &builder, Location loc, SmallVectorImpl< Value > &values)
static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody, Region &launchOpBody, IRMapping &map, bool hasCluster=false)
Adds operations generating block/thread ids and grid/block dimensions at the beginning of the launchF...
static bool extractBeneficiaryOps(Operation *op, const SetVector< Value > &existingDependencies, SetVector< Operation * > &beneficiaryOps, llvm::SmallPtrSetImpl< Value > &availableValues, llvm::function_ref< bool(Operation *)> isSinkingBeneficiary)
For a given operation op, computes whether it is beneficial to sink the operation into the kernel.
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
Attributes are known-constant values of operations.
MLIRContext * getContext() const
Return the context this attribute belongs to.
Block represents an ordered list of Operations.
OpListType::iterator iterator
BlockArgument getArgument(unsigned i)
void erase()
Unlink this Block from its parent region and delete it.
Operation * getTerminator()
Get the terminator operation of this block.
OpListType & getOperations()
This is a utility class for mapping one set of IR entities to another.
auto lookup(T from) const
Lookup a mapped value within the map.
void map(Value from, Value to)
Inserts a new mapping for 'from' to 'to'.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
This class helps build Operations.
Operation * clone(Operation &op, IRMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Operation is the basic unit of execution within MLIR.
Operation * clone(IRMapping &mapper, CloneOptions options=CloneOptions::all())
Create a deep copy of this operation, remapping any operands that use values outside of the operation...
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
operand_range getOperands()
Returns an iterator on the underlying Value's.
void replaceAllUsesWith(ValuesT &&values)
Replace all uses of results of this operation with the provided 'values'.
result_range getResults()
void erase()
Remove this operation from its parent block and delete it.
This class contains a list of basic blocks and a link to the parent operation it is attached to.
void cloneInto(Region *dest, IRMapping &mapper)
Clone the internal blocks from this region into dest.
This class represents a specific symbol use.
This class allows for representing and managing the symbol table used by operations with the 'SymbolT...
Operation * lookup(StringRef name) const
Look up a symbol with the specified name, returning null if no such name exists.
static std::optional< UseRange > getSymbolUses(Operation *from)
Get an iterator range for all of the uses, for any symbol, that are nested within the given operation...
This class provides an abstraction over the various different ranges of value types.
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
MLIRContext * getContext() const
Utility to get the associated MLIRContext that this value is defined in.
Type getType() const
Return the type of this value.
static WalkResult advance()
static DenseArrayAttrImpl get(MLIRContext *context, ArrayRef< int32_t > content)
Builder from ArrayRef<T>.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
detail::constant_int_value_binder m_ConstantInt(IntegerAttr::ValueType *bind_value)
Matches a constant holding a scalar/vector/tensor integer (splat) and writes the integer value to bin...
void replaceAllUsesInRegionWith(Value orig, Value replacement, Region ®ion)
Replace all uses of orig within the given region with replacement.
std::unique_ptr< Pass > createGpuLauchSinkIndexComputationsPass()
Pass that moves ops which are likely an index computation into gpu.launch body.
Attribute parseAttribute(llvm::StringRef attrStr, MLIRContext *context, Type type={}, size_t *numRead=nullptr, bool isKnownNullTerminated=false)
This parses a single MLIR attribute to an MLIR context if it was valid.
void getUsedValuesDefinedAbove(Region ®ion, Region &limit, SetVector< Value > &values)
Fill values with a list of values defined at the ancestors of the limit region and used within region...
LogicalResult sinkOperationsIntoLaunchOp(gpu::LaunchOp launchOp, llvm::function_ref< bool(Operation *)> isSinkingBeneficiary)
Sink operations into the launchOp to reduce the number of values that are used within the region of t...
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
std::unique_ptr< OperationPass< ModuleOp > > createGpuKernelOutliningPass(StringRef dataLayoutStr=StringRef())
Replaces gpu.launch with gpu.launch_func by moving the region into a separate kernel function.
gpu::GPUFuncOp outlineKernelFunc(gpu::LaunchOp launchOp, StringRef kernelFnName, SmallVectorImpl< Value > &operands)
Get a gpu.func created from outlining the region of a gpu.launch op with the given kernelFnName.
Utility class for the GPU dialect to represent triples of Values accessible through ....