c60b897d22
The current Parser library is solely focused on providing API for the textual MLIR format, but MLIR will soon also provide a binary format. This commit renames the current Parser library to AsmParser to better correspond to what the library is actually intended for. A new Parser library is added which will act as a unified parser interface between both text and binary formats. Most parser clients are unaffected, given that the unified interface is essentially the same as the current interface. Only clients that rely on utilizing the AsmParserState, or those that want to parse Attributes/Types need to be updated to point to the AsmParser library. Differential Revision: https://reviews.llvm.org/D129605
393 lines
16 KiB
C++
393 lines
16 KiB
C++
//===- KernelOutlining.cpp - Implementation of GPU kernel outlining -------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements the GPU dialect kernel outlining pass.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "PassDetail.h"
|
|
#include "mlir/AsmParser/AsmParser.h"
|
|
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
|
|
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
|
|
#include "mlir/Dialect/DLTI/DLTI.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
|
|
#include "mlir/Dialect/GPU/Transforms/Passes.h"
|
|
#include "mlir/Dialect/GPU/Transforms/Utils.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/IR/BlockAndValueMapping.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/IR/SymbolTable.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "mlir/Transforms/RegionUtils.h"
|
|
|
|
using namespace mlir;
|
|
|
|
template <typename OpTy>
|
|
static void createForAllDimensions(OpBuilder &builder, Location loc,
|
|
SmallVectorImpl<Value> &values) {
|
|
for (auto dim : {gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z})
|
|
values.push_back(builder.create<OpTy>(loc, builder.getIndexType(), dim));
|
|
}
|
|
|
|
/// Adds operations generating block/thread ids and grid/block dimensions at the
|
|
/// beginning of the `launchFuncOpBody` region. Add mapping from argument in
|
|
/// entry block of `launchOpBody`, to the corresponding result value of the
|
|
/// added operations.
|
|
static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody,
|
|
Region &launchOpBody,
|
|
BlockAndValueMapping &map) {
|
|
OpBuilder builder(loc->getContext());
|
|
Block &firstBlock = launchOpBody.front();
|
|
builder.setInsertionPointToStart(&launchFuncOpBody.front());
|
|
SmallVector<Value, 12> indexOps;
|
|
createForAllDimensions<gpu::BlockIdOp>(builder, loc, indexOps);
|
|
createForAllDimensions<gpu::ThreadIdOp>(builder, loc, indexOps);
|
|
createForAllDimensions<gpu::GridDimOp>(builder, loc, indexOps);
|
|
createForAllDimensions<gpu::BlockDimOp>(builder, loc, indexOps);
|
|
// Replace the leading 12 function args with the respective thread/block index
|
|
// operations. Iterate backwards since args are erased and indices change.
|
|
for (const auto &indexOp : enumerate(indexOps))
|
|
map.map(firstBlock.getArgument(indexOp.index()), indexOp.value());
|
|
}
|
|
|
|
/// Identifies operations that are beneficial to sink into kernels. These
|
|
/// operations may not have side-effects, as otherwise sinking (and hence
|
|
/// duplicating them) is not legal.
|
|
static bool isLikelyAnIndexComputation(Operation *op) {
|
|
return matchPattern(op, m_Constant()) ||
|
|
isa<memref::DimOp, arith::SelectOp, arith::CmpIOp>(op);
|
|
}
|
|
|
|
/// For a given operation `op`, computes whether it is beneficial to sink the
|
|
/// operation into the kernel. An operation can be sunk if doing so does not
|
|
/// introduce new kernel arguments. Whether a value is already available in the
|
|
/// kernel (and hence does not introduce new arguments) is checked by
|
|
/// querying `existingDependencies` and `availableValues`.
|
|
/// If an operand is not yet available, we recursively check whether it can be
|
|
/// made available by siking its defining op.
|
|
/// Operations that are indentified for sinking are added to `beneficiaryOps` in
|
|
/// the order they should appear in the kernel. Furthermore, `availableValues`
|
|
/// is updated with results that will be available after sinking the identified
|
|
/// ops.
|
|
static bool extractBeneficiaryOps(
|
|
Operation *op, const SetVector<Value> &existingDependencies,
|
|
SetVector<Operation *> &beneficiaryOps,
|
|
llvm::SmallPtrSetImpl<Value> &availableValues,
|
|
llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) {
|
|
if (beneficiaryOps.count(op))
|
|
return true;
|
|
|
|
if (!isSinkingBeneficiary(op))
|
|
return false;
|
|
|
|
for (Value operand : op->getOperands()) {
|
|
// It is already visible in the kernel, keep going.
|
|
if (availableValues.count(operand))
|
|
continue;
|
|
// Else check whether it can be made available via sinking or already is a
|
|
// dependency.
|
|
Operation *definingOp = operand.getDefiningOp();
|
|
if ((!definingOp || !extractBeneficiaryOps(definingOp, existingDependencies,
|
|
beneficiaryOps, availableValues,
|
|
isSinkingBeneficiary)) &&
|
|
!existingDependencies.count(operand))
|
|
return false;
|
|
}
|
|
// We will sink the operation, mark its results as now available.
|
|
beneficiaryOps.insert(op);
|
|
for (Value result : op->getResults())
|
|
availableValues.insert(result);
|
|
return true;
|
|
}
|
|
|
|
LogicalResult mlir::sinkOperationsIntoLaunchOp(
|
|
gpu::LaunchOp launchOp,
|
|
llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) {
|
|
assert(isSinkingBeneficiary);
|
|
Region &launchOpBody = launchOp.body();
|
|
|
|
// Identify uses from values defined outside of the scope of the launch
|
|
// operation.
|
|
SetVector<Value> sinkCandidates;
|
|
getUsedValuesDefinedAbove(launchOpBody, sinkCandidates);
|
|
|
|
SetVector<Operation *> toBeSunk;
|
|
llvm::SmallPtrSet<Value, 4> availableValues;
|
|
for (Value operand : sinkCandidates) {
|
|
Operation *operandOp = operand.getDefiningOp();
|
|
if (!operandOp)
|
|
continue;
|
|
extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues,
|
|
isSinkingBeneficiary);
|
|
}
|
|
|
|
// Insert operations so that the defs get cloned before uses.
|
|
BlockAndValueMapping map;
|
|
OpBuilder builder(launchOpBody);
|
|
for (Operation *op : toBeSunk) {
|
|
Operation *clonedOp = builder.clone(*op, map);
|
|
// Only replace uses within the launch op.
|
|
for (auto pair : llvm::zip(op->getResults(), clonedOp->getResults()))
|
|
replaceAllUsesInRegionWith(std::get<0>(pair), std::get<1>(pair),
|
|
launchOp.body());
|
|
}
|
|
return success();
|
|
}
|
|
|
|
/// Outline the `gpu.launch` operation body into a kernel function. Replace
|
|
/// `gpu.terminator` operations by `gpu.return` in the generated function.
|
|
static gpu::GPUFuncOp outlineKernelFuncImpl(gpu::LaunchOp launchOp,
|
|
StringRef kernelFnName,
|
|
SetVector<Value> &operands) {
|
|
Location loc = launchOp.getLoc();
|
|
// Create a builder with no insertion point, insertion will happen separately
|
|
// due to symbol table manipulation.
|
|
OpBuilder builder(launchOp.getContext());
|
|
Region &launchOpBody = launchOp.body();
|
|
|
|
// Identify uses from values defined outside of the scope of the launch
|
|
// operation.
|
|
getUsedValuesDefinedAbove(launchOpBody, operands);
|
|
|
|
// Create the gpu.func operation.
|
|
SmallVector<Type, 4> kernelOperandTypes;
|
|
kernelOperandTypes.reserve(operands.size());
|
|
for (Value operand : operands) {
|
|
kernelOperandTypes.push_back(operand.getType());
|
|
}
|
|
FunctionType type =
|
|
FunctionType::get(launchOp.getContext(), kernelOperandTypes, {});
|
|
auto outlinedFunc = builder.create<gpu::GPUFuncOp>(loc, kernelFnName, type);
|
|
outlinedFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
|
|
builder.getUnitAttr());
|
|
BlockAndValueMapping map;
|
|
|
|
// Map the arguments corresponding to the launch parameters like blockIdx,
|
|
// threadIdx, etc.
|
|
Region &outlinedFuncBody = outlinedFunc.body();
|
|
injectGpuIndexOperations(loc, outlinedFuncBody, launchOpBody, map);
|
|
|
|
// Map arguments from gpu.launch region to the arguments of the gpu.func
|
|
// operation.
|
|
Block &entryBlock = outlinedFuncBody.front();
|
|
for (const auto &operand : enumerate(operands))
|
|
map.map(operand.value(), entryBlock.getArgument(operand.index()));
|
|
|
|
// Clone the region of the gpu.launch operation into the gpu.func operation.
|
|
// TODO: If cloneInto can be modified such that if a mapping for
|
|
// a block exists, that block will be used to clone operations into (at the
|
|
// end of the block), instead of creating a new block, this would be much
|
|
// cleaner.
|
|
launchOpBody.cloneInto(&outlinedFuncBody, map);
|
|
|
|
// Branch from entry of the gpu.func operation to the block that is cloned
|
|
// from the entry block of the gpu.launch operation.
|
|
Block &launchOpEntry = launchOpBody.front();
|
|
Block *clonedLaunchOpEntry = map.lookup(&launchOpEntry);
|
|
builder.setInsertionPointToEnd(&entryBlock);
|
|
builder.create<cf::BranchOp>(loc, clonedLaunchOpEntry);
|
|
|
|
outlinedFunc.walk([](gpu::TerminatorOp op) {
|
|
OpBuilder replacer(op);
|
|
replacer.create<gpu::ReturnOp>(op.getLoc());
|
|
op.erase();
|
|
});
|
|
return outlinedFunc;
|
|
}
|
|
|
|
gpu::GPUFuncOp mlir::outlineKernelFunc(gpu::LaunchOp launchOp,
|
|
StringRef kernelFnName,
|
|
llvm::SmallVectorImpl<Value> &operands) {
|
|
DenseSet<Value> inputOperandSet;
|
|
inputOperandSet.insert(operands.begin(), operands.end());
|
|
SetVector<Value> operandSet(operands.begin(), operands.end());
|
|
auto funcOp = outlineKernelFuncImpl(launchOp, kernelFnName, operandSet);
|
|
for (auto operand : operandSet) {
|
|
if (!inputOperandSet.count(operand))
|
|
operands.push_back(operand);
|
|
}
|
|
return funcOp;
|
|
}
|
|
|
|
/// Replace `gpu.launch` operations with an `gpu.launch_func` operation
|
|
/// launching `kernelFunc`. The kernel func contains the body of the
|
|
/// `gpu.launch` with constant region arguments inlined.
|
|
static void convertToLaunchFuncOp(gpu::LaunchOp launchOp,
|
|
gpu::GPUFuncOp kernelFunc,
|
|
ValueRange operands) {
|
|
OpBuilder builder(launchOp);
|
|
// The launch op has an optional dynamic shared memory size. If it doesn't
|
|
// exist, we use zero.
|
|
Value asyncToken = launchOp.asyncToken();
|
|
auto launchFunc = builder.create<gpu::LaunchFuncOp>(
|
|
launchOp.getLoc(), kernelFunc, launchOp.getGridSizeOperandValues(),
|
|
launchOp.getBlockSizeOperandValues(), launchOp.dynamicSharedMemorySize(),
|
|
operands, asyncToken ? asyncToken.getType() : nullptr,
|
|
launchOp.asyncDependencies());
|
|
launchOp.replaceAllUsesWith(launchFunc);
|
|
launchOp.erase();
|
|
}
|
|
|
|
namespace {
|
|
/// Pass that moves ops which are likely an index computation into gpu.launch
|
|
/// body.
|
|
class GpuLaunchSinkIndexComputationsPass
|
|
: public GpuLaunchSinkIndexComputationsBase<
|
|
GpuLaunchSinkIndexComputationsPass> {
|
|
public:
|
|
void runOnOperation() override {
|
|
Operation *op = getOperation();
|
|
if (op->walk([](gpu::LaunchOp launch) {
|
|
// Pull in instructions that can be sunk
|
|
if (failed(sinkOperationsIntoLaunchOp(launch,
|
|
isLikelyAnIndexComputation)))
|
|
return WalkResult::interrupt();
|
|
|
|
return WalkResult::advance();
|
|
}).wasInterrupted())
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
|
|
/// Pass that moves the kernel of each LaunchOp into its separate nested module.
|
|
///
|
|
/// This pass moves the kernel code of each LaunchOp into a function created
|
|
/// inside a nested module. It also creates an external function of the same
|
|
/// name in the parent module.
|
|
///
|
|
/// The gpu.modules are intended to be compiled to a cubin blob independently in
|
|
/// a separate pass. The external functions can then be annotated with the
|
|
/// symbol of the cubin accessor function.
|
|
class GpuKernelOutliningPass
|
|
: public GpuKernelOutliningBase<GpuKernelOutliningPass> {
|
|
public:
|
|
GpuKernelOutliningPass(StringRef dlStr) {
|
|
if (!dlStr.empty() && !dataLayoutStr.hasValue())
|
|
dataLayoutStr = dlStr.str();
|
|
}
|
|
|
|
GpuKernelOutliningPass(const GpuKernelOutliningPass &other)
|
|
: GpuKernelOutliningBase(other), dataLayoutSpec(other.dataLayoutSpec) {
|
|
dataLayoutStr = other.dataLayoutStr.getValue();
|
|
}
|
|
|
|
LogicalResult initialize(MLIRContext *context) override {
|
|
// Initialize the data layout specification from the data layout string.
|
|
if (!dataLayoutStr.empty()) {
|
|
Attribute resultAttr = mlir::parseAttribute(dataLayoutStr, context);
|
|
if (!resultAttr)
|
|
return failure();
|
|
|
|
dataLayoutSpec = resultAttr.dyn_cast<DataLayoutSpecInterface>();
|
|
if (!dataLayoutSpec)
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void runOnOperation() override {
|
|
SymbolTable symbolTable(getOperation());
|
|
bool modified = false;
|
|
for (auto func : getOperation().getOps<func::FuncOp>()) {
|
|
// Insert just after the function.
|
|
Block::iterator insertPt(func->getNextNode());
|
|
auto funcWalkResult = func.walk([&](gpu::LaunchOp op) {
|
|
SetVector<Value> operands;
|
|
std::string kernelFnName =
|
|
Twine(op->getParentOfType<func::FuncOp>().getName(), "_kernel")
|
|
.str();
|
|
|
|
gpu::GPUFuncOp outlinedFunc =
|
|
outlineKernelFuncImpl(op, kernelFnName, operands);
|
|
|
|
// Create nested module and insert outlinedFunc. The module will
|
|
// originally get the same name as the function, but may be renamed on
|
|
// insertion into the parent module.
|
|
auto kernelModule = createKernelModule(outlinedFunc, symbolTable);
|
|
symbolTable.insert(kernelModule, insertPt);
|
|
|
|
// Potentially changes signature, pulling in constants.
|
|
convertToLaunchFuncOp(op, outlinedFunc, operands.getArrayRef());
|
|
modified = true;
|
|
return WalkResult::advance();
|
|
});
|
|
if (funcWalkResult.wasInterrupted())
|
|
return signalPassFailure();
|
|
}
|
|
|
|
// If any new module was inserted in this module, annotate this module as
|
|
// a container module.
|
|
if (modified)
|
|
getOperation()->setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
|
|
UnitAttr::get(&getContext()));
|
|
}
|
|
|
|
private:
|
|
/// Returns a gpu.module containing kernelFunc and all callees (recursive).
|
|
gpu::GPUModuleOp createKernelModule(gpu::GPUFuncOp kernelFunc,
|
|
const SymbolTable &parentSymbolTable) {
|
|
// TODO: This code cannot use an OpBuilder because it must be inserted into
|
|
// a SymbolTable by the caller. SymbolTable needs to be refactored to
|
|
// prevent manual building of Ops with symbols in code using SymbolTables
|
|
// and then this needs to use the OpBuilder.
|
|
auto *context = getOperation().getContext();
|
|
OpBuilder builder(context);
|
|
auto kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(),
|
|
kernelFunc.getName());
|
|
|
|
// If a valid data layout spec was provided, attach it to the kernel module.
|
|
// Otherwise, the default data layout will be used.
|
|
if (dataLayoutSpec)
|
|
kernelModule->setAttr(DLTIDialect::kDataLayoutAttrName, dataLayoutSpec);
|
|
|
|
SymbolTable symbolTable(kernelModule);
|
|
symbolTable.insert(kernelFunc);
|
|
|
|
SmallVector<Operation *, 8> symbolDefWorklist = {kernelFunc};
|
|
while (!symbolDefWorklist.empty()) {
|
|
if (Optional<SymbolTable::UseRange> symbolUses =
|
|
SymbolTable::getSymbolUses(symbolDefWorklist.pop_back_val())) {
|
|
for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
|
|
StringRef symbolName =
|
|
symbolUse.getSymbolRef().cast<FlatSymbolRefAttr>().getValue();
|
|
if (symbolTable.lookup(symbolName))
|
|
continue;
|
|
|
|
Operation *symbolDefClone =
|
|
parentSymbolTable.lookup(symbolName)->clone();
|
|
symbolDefWorklist.push_back(symbolDefClone);
|
|
symbolTable.insert(symbolDefClone);
|
|
}
|
|
}
|
|
}
|
|
|
|
return kernelModule;
|
|
}
|
|
|
|
Option<std::string> dataLayoutStr{
|
|
*this, "data-layout-str",
|
|
llvm::cl::desc("String containing the data layout specification to be "
|
|
"attached to the GPU kernel module")};
|
|
|
|
DataLayoutSpecInterface dataLayoutSpec;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<Pass> mlir::createGpuLauchSinkIndexComputationsPass() {
|
|
return std::make_unique<GpuLaunchSinkIndexComputationsPass>();
|
|
}
|
|
|
|
std::unique_ptr<OperationPass<ModuleOp>>
|
|
mlir::createGpuKernelOutliningPass(StringRef dataLayoutStr) {
|
|
return std::make_unique<GpuKernelOutliningPass>(dataLayoutStr);
|
|
}
|