9a52ea5cf9
Summary: This is based on the use of code constantly checking for an attribute on a model and instead represents the distinct operaion with a different op. Instead, this op can be used to provide better filtering. Reverts "Revert "[mlir] Create a gpu.module operation for the GPU Dialect."" This reverts commit ac446302ca4145cdc89f377c0c364c29ee303be5 after fixing internal Google issues. This additionally updates ROCDL lowering to use the new gpu.module. Reviewers: herhut, mravishankar, antiagainst, nicolasvasilache Subscribers: jholewinski, mgorny, mehdi_amini, jpienaar, burmako, shauheen, csigg, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, llvm-commits, mravishankar, rriddle, antiagainst, bkramer Tags: #llvm Differential Revision: https://reviews.llvm.org/D72921
758 lines
34 KiB
C++
758 lines
34 KiB
C++
//===- LowerGpuOpsToNVVMOps.cpp - MLIR GPU to NVVM lowering passes --------===//
|
|
//
|
|
// Part of the MLIR 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 a pass to generate NVVMIR operations for higher-level
|
|
// GPU operations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
|
|
|
|
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
|
|
#include "mlir/Dialect/GPU/GPUDialect.h"
|
|
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
|
|
#include "mlir/IR/BlockAndValueMapping.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
#include "llvm/Support/FormatVariadic.h"
|
|
|
|
#include "../GPUCommon/IndexIntrinsicsOpLowering.h"
|
|
#include "../GPUCommon/OpToFuncCallLowering.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
/// Derived type converter for GPU to NVVM lowering. The GPU dialect uses memory
|
|
/// space 5 for private memory attributions, but NVVM represents private
|
|
/// memory allocations as local `alloca`s in the default address space. This
|
|
/// converter drops the private memory space to support the use case above.
|
|
class NVVMTypeConverter : public LLVMTypeConverter {
|
|
public:
|
|
using LLVMTypeConverter::LLVMTypeConverter;
|
|
|
|
Type convertType(Type type) override {
|
|
auto memref = type.dyn_cast<MemRefType>();
|
|
if (memref &&
|
|
memref.getMemorySpace() == gpu::GPUDialect::getPrivateAddressSpace()) {
|
|
type = MemRefType::get(memref.getShape(), memref.getElementType(),
|
|
memref.getAffineMaps());
|
|
}
|
|
|
|
return LLVMTypeConverter::convertType(type);
|
|
}
|
|
};
|
|
|
|
/// Converts all_reduce op to LLVM/NVVM ops.
|
|
struct GPUAllReduceOpLowering : public LLVMOpLowering {
|
|
using AccumulatorFactory =
|
|
std::function<Value(Location, Value, Value, ConversionPatternRewriter &)>;
|
|
|
|
explicit GPUAllReduceOpLowering(LLVMTypeConverter &lowering_)
|
|
: LLVMOpLowering(gpu::AllReduceOp::getOperationName(),
|
|
lowering_.getDialect()->getContext(), lowering_),
|
|
int32Type(LLVM::LLVMType::getInt32Ty(lowering_.getDialect())) {}
|
|
|
|
PatternMatchResult
|
|
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Location loc = op->getLoc();
|
|
Value operand = operands.front();
|
|
|
|
// TODO(csigg): Generalize to other types of accumulation.
|
|
assert(op->getOperand(0).getType().isIntOrFloat());
|
|
|
|
// Create the reduction using an accumulator factory.
|
|
AccumulatorFactory factory =
|
|
getFactory(cast<gpu::AllReduceOp>(op), operand);
|
|
assert(factory && "failed to create accumulator factory");
|
|
Value result = createBlockReduce(loc, operand, factory, rewriter);
|
|
|
|
rewriter.replaceOp(op, {result});
|
|
return matchSuccess();
|
|
}
|
|
|
|
private:
|
|
/// Returns an accumulator factory using either the op attribute or the body
|
|
/// region.
|
|
AccumulatorFactory getFactory(gpu::AllReduceOp allReduce,
|
|
Value operand) const {
|
|
if (!allReduce.body().empty()) {
|
|
return getFactory(allReduce.body());
|
|
}
|
|
if (allReduce.op()) {
|
|
auto type = operand.getType().cast<LLVM::LLVMType>();
|
|
return getFactory(*allReduce.op(), type.getUnderlyingType());
|
|
}
|
|
return AccumulatorFactory();
|
|
}
|
|
|
|
/// Returns an accumulator factory that clones the body. The body's entry
|
|
/// block is expected to have 2 arguments. The gpu.yield return the
|
|
/// accumulated value of the same type.
|
|
AccumulatorFactory getFactory(Region &body) const {
|
|
return AccumulatorFactory([&](Location loc, Value lhs, Value rhs,
|
|
ConversionPatternRewriter &rewriter) {
|
|
Block *block = rewriter.getInsertionBlock();
|
|
Block *split = rewriter.splitBlock(block, rewriter.getInsertionPoint());
|
|
|
|
// Insert accumulator body between split block.
|
|
BlockAndValueMapping mapping;
|
|
mapping.map(body.front().getArgument(0), lhs);
|
|
mapping.map(body.front().getArgument(1), rhs);
|
|
rewriter.cloneRegionBefore(body, *split->getParent(),
|
|
split->getIterator(), mapping);
|
|
|
|
// Add branch before inserted body, into body.
|
|
block = block->getNextNode();
|
|
rewriter.create<LLVM::BrOp>(loc, ArrayRef<Value>{},
|
|
llvm::makeArrayRef(block), ValueRange());
|
|
|
|
// Replace all gpu.yield ops with branch out of body.
|
|
for (; block != split; block = block->getNextNode()) {
|
|
Operation *terminator = block->getTerminator();
|
|
if (!llvm::isa<gpu::YieldOp>(terminator))
|
|
continue;
|
|
rewriter.setInsertionPointToEnd(block);
|
|
rewriter.replaceOpWithNewOp<LLVM::BrOp>(
|
|
terminator, ArrayRef<Value>{}, llvm::makeArrayRef(split),
|
|
ValueRange(terminator->getOperand(0)));
|
|
}
|
|
|
|
// Return accumulator result.
|
|
rewriter.setInsertionPointToStart(split);
|
|
return split->addArgument(lhs.getType());
|
|
});
|
|
}
|
|
|
|
/// Returns an accumulator factory that creates an op specified by opName.
|
|
AccumulatorFactory getFactory(StringRef opName, llvm::Type *type) const {
|
|
if (type->isVectorTy() || type->isArrayTy())
|
|
return getFactory(opName, type->getSequentialElementType());
|
|
|
|
bool isFloatingPoint = type->isFloatingPointTy();
|
|
|
|
if (opName == "add") {
|
|
return isFloatingPoint ? getFactory<LLVM::FAddOp>()
|
|
: getFactory<LLVM::AddOp>();
|
|
}
|
|
if (opName == "mul") {
|
|
return isFloatingPoint ? getFactory<LLVM::FMulOp>()
|
|
: getFactory<LLVM::MulOp>();
|
|
}
|
|
|
|
return AccumulatorFactory();
|
|
}
|
|
|
|
/// Returns an accumulator factory that creates an op of type T.
|
|
template <typename T> AccumulatorFactory getFactory() const {
|
|
return [](Location loc, Value lhs, Value rhs,
|
|
ConversionPatternRewriter &rewriter) {
|
|
return rewriter.create<T>(loc, lhs.getType(), lhs, rhs);
|
|
};
|
|
}
|
|
|
|
/// Creates an all_reduce across the block.
|
|
///
|
|
/// First reduce the elements within a warp. The first thread of each warp
|
|
/// writes the intermediate result to shared memory. After synchronizing the
|
|
/// block, the first warp reduces the values from shared memory. The result
|
|
/// is broadcasted to all threads through shared memory.
|
|
///
|
|
/// %warp_reduce = `createWarpReduce(%operand)`
|
|
/// %shared_mem_ptr = llvm.mlir.addressof @reduce_buffer
|
|
/// %zero = llvm.mlir.constant(0 : i32) : !llvm.i32
|
|
/// %lane_id = nvvm.read.ptx.sreg.laneid : !llvm.i32
|
|
/// %is_first_lane = llvm.icmp "eq" %lane_id, %zero : !llvm.i1
|
|
/// %thread_idx = `getLinearThreadIndex()` : !llvm.i32
|
|
/// llvm.cond_br %is_first_lane, ^then1, ^continue1
|
|
/// ^then1:
|
|
/// %warp_id = `getWarpId()`
|
|
/// %store_dst = llvm.getelementptr %shared_mem_ptr[%zero, %warp_id]
|
|
/// llvm.store %store_dst, %warp_reduce
|
|
/// llvm.br ^continue1
|
|
/// ^continue1:
|
|
/// nvvm.barrier0
|
|
/// %num_warps = `getNumWarps()` : !llvm.i32
|
|
/// %is_valid_warp = llvm.icmp "slt" %thread_idx, %num_warps
|
|
/// %result_ptr = llvm.getelementptr %shared_mem_ptr[%zero, %zero]
|
|
/// llvm.cond_br %is_first_lane, ^then2, ^continue2
|
|
/// ^then2:
|
|
/// %load_src = llvm.getelementptr %shared_mem_ptr[%zero, %thread_idx]
|
|
/// %value = llvm.load %load_src
|
|
/// %result = `createWarpReduce(%value)`
|
|
/// llvm.store %result_ptr, %result
|
|
/// llvm.br ^continue2
|
|
/// ^continue2:
|
|
/// nvvm.barrier0
|
|
/// %result = llvm.load %result_ptr
|
|
/// return %result
|
|
///
|
|
Value createBlockReduce(Location loc, Value operand,
|
|
AccumulatorFactory &accumFactory,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto type = operand.getType().cast<LLVM::LLVMType>();
|
|
|
|
// Create shared memory array to store the warp reduction.
|
|
auto module = operand.getDefiningOp()->getParentOfType<gpu::GPUModuleOp>();
|
|
assert(module && "op must belong to a module");
|
|
Value sharedMemPtr =
|
|
createSharedMemoryArray(loc, module, type, kWarpSize, rewriter);
|
|
|
|
Value zero = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(0u));
|
|
Value laneId = rewriter.create<NVVM::LaneIdOp>(loc, int32Type);
|
|
Value isFirstLane = rewriter.create<LLVM::ICmpOp>(
|
|
loc, LLVM::ICmpPredicate::eq, laneId, zero);
|
|
Value threadIdx = getLinearThreadIndex(loc, rewriter);
|
|
Value blockSize = getBlockSize(loc, rewriter);
|
|
Value activeWidth = getActiveWidth(loc, threadIdx, blockSize, rewriter);
|
|
|
|
// Reduce elements within each warp to produce the intermediate results.
|
|
Value warpReduce = createWarpReduce(loc, activeWidth, laneId, operand,
|
|
accumFactory, rewriter);
|
|
|
|
// Write the intermediate results to shared memory, using the first lane of
|
|
// each warp.
|
|
createPredicatedBlock(loc, rewriter, isFirstLane, [&] {
|
|
Value warpId = getDivideByWarpSize(threadIdx, rewriter);
|
|
Value storeDst = rewriter.create<LLVM::GEPOp>(
|
|
loc, type, sharedMemPtr, ArrayRef<Value>({zero, warpId}));
|
|
rewriter.create<LLVM::StoreOp>(loc, warpReduce, storeDst);
|
|
});
|
|
rewriter.create<NVVM::Barrier0Op>(loc);
|
|
|
|
Value numWarps = getNumWarps(loc, blockSize, rewriter);
|
|
Value isValidWarp = rewriter.create<LLVM::ICmpOp>(
|
|
loc, LLVM::ICmpPredicate::slt, threadIdx, numWarps);
|
|
Value resultPtr = rewriter.create<LLVM::GEPOp>(
|
|
loc, type, sharedMemPtr, ArrayRef<Value>({zero, zero}));
|
|
|
|
// Use the first numWarps threads to reduce the intermediate results from
|
|
// shared memory. The final result is written to shared memory again.
|
|
createPredicatedBlock(loc, rewriter, isValidWarp, [&] {
|
|
Value loadSrc = rewriter.create<LLVM::GEPOp>(
|
|
loc, type, sharedMemPtr, ArrayRef<Value>({zero, threadIdx}));
|
|
Value value = rewriter.create<LLVM::LoadOp>(loc, type, loadSrc);
|
|
Value result = createWarpReduce(loc, numWarps, laneId, value,
|
|
accumFactory, rewriter);
|
|
rewriter.create<LLVM::StoreOp>(loc, result, resultPtr);
|
|
});
|
|
rewriter.create<NVVM::Barrier0Op>(loc);
|
|
|
|
// Load and return result from shared memory.
|
|
Value result = rewriter.create<LLVM::LoadOp>(loc, type, resultPtr);
|
|
return result;
|
|
}
|
|
|
|
/// Creates an if-block skeleton and calls the two factories to generate the
|
|
/// ops in the `then` and `else` block..
|
|
///
|
|
/// llvm.cond_br %condition, ^then, ^continue
|
|
/// ^then:
|
|
/// %then_operands = `thenOpsFactory()`
|
|
/// llvm.br ^continue(%then_operands)
|
|
/// ^else:
|
|
/// %else_operands = `elseOpsFactory()`
|
|
/// llvm.br ^continue(%else_operands)
|
|
/// ^continue(%block_operands):
|
|
///
|
|
template <typename ThenOpsFactory, typename ElseOpsFactory>
|
|
void createIf(Location loc, ConversionPatternRewriter &rewriter,
|
|
Value condition, ThenOpsFactory &&thenOpsFactory,
|
|
ElseOpsFactory &&elseOpsFactory) const {
|
|
Block *currentBlock = rewriter.getInsertionBlock();
|
|
auto currentPoint = rewriter.getInsertionPoint();
|
|
|
|
Block *thenBlock = rewriter.splitBlock(currentBlock, currentPoint);
|
|
Block *elseBlock = rewriter.splitBlock(thenBlock, thenBlock->begin());
|
|
Block *continueBlock = rewriter.splitBlock(elseBlock, elseBlock->begin());
|
|
|
|
rewriter.setInsertionPointToEnd(currentBlock);
|
|
rewriter.create<LLVM::CondBrOp>(loc, llvm::makeArrayRef(condition),
|
|
ArrayRef<Block *>{thenBlock, elseBlock});
|
|
|
|
auto addBranch = [&](ValueRange operands) {
|
|
rewriter.create<LLVM::BrOp>(loc, ArrayRef<Value>{},
|
|
llvm::makeArrayRef(continueBlock),
|
|
llvm::makeArrayRef(operands));
|
|
};
|
|
|
|
rewriter.setInsertionPointToStart(thenBlock);
|
|
auto thenOperands = thenOpsFactory();
|
|
addBranch(thenOperands);
|
|
|
|
rewriter.setInsertionPointToStart(elseBlock);
|
|
auto elseOperands = elseOpsFactory();
|
|
addBranch(elseOperands);
|
|
|
|
assert(thenOperands.size() == elseOperands.size());
|
|
rewriter.setInsertionPointToStart(continueBlock);
|
|
for (auto operand : thenOperands)
|
|
continueBlock->addArgument(operand.getType());
|
|
}
|
|
|
|
/// Shortcut for createIf with empty else block and no block operands.
|
|
template <typename Factory>
|
|
void createPredicatedBlock(Location loc, ConversionPatternRewriter &rewriter,
|
|
Value condition,
|
|
Factory &&predicatedOpsFactory) const {
|
|
createIf(
|
|
loc, rewriter, condition,
|
|
[&] {
|
|
predicatedOpsFactory();
|
|
return ArrayRef<Value>();
|
|
},
|
|
[&] { return ArrayRef<Value>(); });
|
|
}
|
|
|
|
/// Creates a reduction across the first activeWidth lanes of a warp.
|
|
/// The first lane returns the result, all others return values are undefined.
|
|
Value createWarpReduce(Location loc, Value activeWidth, Value laneId,
|
|
Value operand, AccumulatorFactory accumFactory,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value warpSize = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(kWarpSize));
|
|
Value isPartialWarp = rewriter.create<LLVM::ICmpOp>(
|
|
loc, LLVM::ICmpPredicate::slt, activeWidth, warpSize);
|
|
auto type = operand.getType().cast<LLVM::LLVMType>();
|
|
|
|
createIf(
|
|
loc, rewriter, isPartialWarp,
|
|
// Generate reduction over a (potentially) partial warp.
|
|
[&] {
|
|
Value value = operand;
|
|
Value one = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(1));
|
|
// Bit mask of active lanes: `(1 << activeWidth) - 1`.
|
|
Value activeMask = rewriter.create<LLVM::SubOp>(
|
|
loc, int32Type,
|
|
rewriter.create<LLVM::ShlOp>(loc, int32Type, one, activeWidth),
|
|
one);
|
|
// Clamp lane: `activeWidth - 1`
|
|
Value maskAndClamp =
|
|
rewriter.create<LLVM::SubOp>(loc, int32Type, activeWidth, one);
|
|
auto dialect = lowering.getDialect();
|
|
auto predTy = LLVM::LLVMType::getInt1Ty(dialect);
|
|
auto shflTy = LLVM::LLVMType::getStructTy(dialect, {type, predTy});
|
|
auto returnValueAndIsValidAttr = rewriter.getUnitAttr();
|
|
|
|
// Repeatedly shuffle value from 'laneId ^ i' and accumulate if source
|
|
// lane is within the active range. All lanes contain the final
|
|
// result, but only the first lane's result is used.
|
|
for (int i = 1; i < kWarpSize; i <<= 1) {
|
|
Value offset = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(i));
|
|
Value shfl = rewriter.create<NVVM::ShflBflyOp>(
|
|
loc, shflTy, activeMask, value, offset, maskAndClamp,
|
|
returnValueAndIsValidAttr);
|
|
Value isActiveSrcLane = rewriter.create<LLVM::ExtractValueOp>(
|
|
loc, predTy, shfl, rewriter.getIndexArrayAttr(1));
|
|
// Skip the accumulation if the shuffle op read from a lane outside
|
|
// of the active range.
|
|
createIf(
|
|
loc, rewriter, isActiveSrcLane,
|
|
[&] {
|
|
Value shflValue = rewriter.create<LLVM::ExtractValueOp>(
|
|
loc, type, shfl, rewriter.getIndexArrayAttr(0));
|
|
return SmallVector<Value, 1>{
|
|
accumFactory(loc, value, shflValue, rewriter)};
|
|
},
|
|
[&] { return llvm::makeArrayRef(value); });
|
|
value = rewriter.getInsertionBlock()->getArgument(0);
|
|
}
|
|
return SmallVector<Value, 1>{value};
|
|
},
|
|
// Generate a reduction over the entire warp. This is a specialization
|
|
// of the above reduction with unconditional accumulation.
|
|
[&] {
|
|
Value value = operand;
|
|
Value activeMask = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(~0u));
|
|
Value maskAndClamp = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(kWarpSize - 1));
|
|
for (int i = 1; i < kWarpSize; i <<= 1) {
|
|
Value offset = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(i));
|
|
Value shflValue = rewriter.create<NVVM::ShflBflyOp>(
|
|
loc, type, activeMask, value, offset, maskAndClamp,
|
|
/*return_value_and_is_valid=*/UnitAttr());
|
|
value = accumFactory(loc, value, shflValue, rewriter);
|
|
}
|
|
return SmallVector<Value, 1>{value};
|
|
});
|
|
return rewriter.getInsertionBlock()->getArgument(0);
|
|
}
|
|
|
|
/// Creates a global array stored in shared memory.
|
|
Value createSharedMemoryArray(Location loc, gpu::GPUModuleOp module,
|
|
LLVM::LLVMType elementType, int numElements,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
OpBuilder builder(module.body());
|
|
|
|
auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements);
|
|
StringRef name = "reduce_buffer";
|
|
auto globalOp = builder.create<LLVM::GlobalOp>(
|
|
loc, arrayType.cast<LLVM::LLVMType>(),
|
|
/*isConstant=*/false, LLVM::Linkage::Internal, name,
|
|
/*value=*/Attribute(), gpu::GPUDialect::getWorkgroupAddressSpace());
|
|
|
|
return rewriter.create<LLVM::AddressOfOp>(loc, globalOp);
|
|
}
|
|
|
|
/// Returns the index of the thread within the block.
|
|
Value getLinearThreadIndex(Location loc,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value dimX = rewriter.create<NVVM::BlockDimXOp>(loc, int32Type);
|
|
Value dimY = rewriter.create<NVVM::BlockDimYOp>(loc, int32Type);
|
|
Value idX = rewriter.create<NVVM::ThreadIdXOp>(loc, int32Type);
|
|
Value idY = rewriter.create<NVVM::ThreadIdYOp>(loc, int32Type);
|
|
Value idZ = rewriter.create<NVVM::ThreadIdZOp>(loc, int32Type);
|
|
Value tmp1 = rewriter.create<LLVM::MulOp>(loc, int32Type, idZ, dimY);
|
|
Value tmp2 = rewriter.create<LLVM::AddOp>(loc, int32Type, tmp1, idY);
|
|
Value tmp3 = rewriter.create<LLVM::MulOp>(loc, int32Type, tmp2, dimX);
|
|
return rewriter.create<LLVM::AddOp>(loc, int32Type, tmp3, idX);
|
|
}
|
|
|
|
/// Returns the number of threads in the block.
|
|
Value getBlockSize(Location loc, ConversionPatternRewriter &rewriter) const {
|
|
Value dimX = rewriter.create<NVVM::BlockDimXOp>(loc, int32Type);
|
|
Value dimY = rewriter.create<NVVM::BlockDimYOp>(loc, int32Type);
|
|
Value dimZ = rewriter.create<NVVM::BlockDimZOp>(loc, int32Type);
|
|
Value dimXY = rewriter.create<LLVM::MulOp>(loc, int32Type, dimX, dimY);
|
|
return rewriter.create<LLVM::MulOp>(loc, int32Type, dimXY, dimZ);
|
|
}
|
|
|
|
/// Returns the number of warps in the block.
|
|
Value getNumWarps(Location loc, Value blockSize,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto warpSizeMinusOne = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(kWarpSize - 1));
|
|
auto biasedBlockSize = rewriter.create<LLVM::AddOp>(
|
|
loc, int32Type, blockSize, warpSizeMinusOne);
|
|
return getDivideByWarpSize(biasedBlockSize, rewriter);
|
|
}
|
|
|
|
/// Returns the number of active threads in the warp, not clamped to 32.
|
|
Value getActiveWidth(Location loc, Value threadIdx, Value blockSize,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value threadIdxMask = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(~(kWarpSize - 1)));
|
|
Value numThreadsWithSmallerWarpId =
|
|
rewriter.create<LLVM::AndOp>(loc, threadIdx, threadIdxMask);
|
|
return rewriter.create<LLVM::SubOp>(loc, blockSize,
|
|
numThreadsWithSmallerWarpId);
|
|
}
|
|
|
|
/// Returns value divided by the warp size (i.e. 32).
|
|
Value getDivideByWarpSize(Value value,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto loc = value.getLoc();
|
|
auto warpSize = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(kWarpSize));
|
|
return rewriter.create<LLVM::SDivOp>(loc, int32Type, value, warpSize);
|
|
}
|
|
|
|
LLVM::LLVMType int32Type;
|
|
|
|
static constexpr int kWarpSize = 32;
|
|
};
|
|
|
|
struct GPUShuffleOpLowering : public LLVMOpLowering {
|
|
explicit GPUShuffleOpLowering(LLVMTypeConverter &lowering_)
|
|
: LLVMOpLowering(gpu::ShuffleOp::getOperationName(),
|
|
lowering_.getDialect()->getContext(), lowering_) {}
|
|
|
|
/// Lowers a shuffle to the corresponding NVVM op.
|
|
///
|
|
/// Convert the `width` argument into an activeMask (a bitmask which specifies
|
|
/// which threads participate in the shuffle) and a maskAndClamp (specifying
|
|
/// the highest lane which participates in the shuffle).
|
|
///
|
|
/// %one = llvm.constant(1 : i32) : !llvm.i32
|
|
/// %shl = llvm.shl %one, %width : !llvm.i32
|
|
/// %active_mask = llvm.sub %shl, %one : !llvm.i32
|
|
/// %mask_and_clamp = llvm.sub %width, %one : !llvm.i32
|
|
/// %shfl = nvvm.shfl.sync.bfly %active_mask, %value, %offset,
|
|
/// %mask_and_clamp : !llvm<"{ float, i1 }">
|
|
/// %shfl_value = llvm.extractvalue %shfl[0 : index] :
|
|
/// !llvm<"{ float, i1 }">
|
|
/// %shfl_pred = llvm.extractvalue %shfl[1 : index] :
|
|
/// !llvm<"{ float, i1 }">
|
|
PatternMatchResult
|
|
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Location loc = op->getLoc();
|
|
gpu::ShuffleOpOperandAdaptor adaptor(operands);
|
|
|
|
auto dialect = lowering.getDialect();
|
|
auto valueTy = adaptor.value().getType().cast<LLVM::LLVMType>();
|
|
auto int32Type = LLVM::LLVMType::getInt32Ty(dialect);
|
|
auto predTy = LLVM::LLVMType::getInt1Ty(dialect);
|
|
auto resultTy = LLVM::LLVMType::getStructTy(dialect, {valueTy, predTy});
|
|
|
|
Value one = rewriter.create<LLVM::ConstantOp>(
|
|
loc, int32Type, rewriter.getI32IntegerAttr(1));
|
|
// Bit mask of active lanes: `(1 << activeWidth) - 1`.
|
|
Value activeMask = rewriter.create<LLVM::SubOp>(
|
|
loc, int32Type,
|
|
rewriter.create<LLVM::ShlOp>(loc, int32Type, one, adaptor.width()),
|
|
one);
|
|
// Clamp lane: `activeWidth - 1`
|
|
Value maskAndClamp =
|
|
rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.width(), one);
|
|
|
|
auto returnValueAndIsValidAttr = rewriter.getUnitAttr();
|
|
Value shfl = rewriter.create<NVVM::ShflBflyOp>(
|
|
loc, resultTy, activeMask, adaptor.value(), adaptor.offset(),
|
|
maskAndClamp, returnValueAndIsValidAttr);
|
|
Value shflValue = rewriter.create<LLVM::ExtractValueOp>(
|
|
loc, valueTy, shfl, rewriter.getIndexArrayAttr(0));
|
|
Value isActiveSrcLane = rewriter.create<LLVM::ExtractValueOp>(
|
|
loc, predTy, shfl, rewriter.getIndexArrayAttr(1));
|
|
|
|
rewriter.replaceOp(op, {shflValue, isActiveSrcLane});
|
|
return matchSuccess();
|
|
}
|
|
};
|
|
|
|
struct GPUFuncOpLowering : LLVMOpLowering {
|
|
explicit GPUFuncOpLowering(LLVMTypeConverter &typeConverter)
|
|
: LLVMOpLowering(gpu::GPUFuncOp::getOperationName(),
|
|
typeConverter.getDialect()->getContext(),
|
|
typeConverter) {}
|
|
|
|
PatternMatchResult
|
|
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
assert(operands.empty() && "func op is not expected to have operands");
|
|
auto gpuFuncOp = cast<gpu::GPUFuncOp>(op);
|
|
Location loc = gpuFuncOp.getLoc();
|
|
|
|
SmallVector<LLVM::GlobalOp, 3> workgroupBuffers;
|
|
workgroupBuffers.reserve(gpuFuncOp.getNumWorkgroupAttributions());
|
|
for (auto en : llvm::enumerate(gpuFuncOp.getWorkgroupAttributions())) {
|
|
Value attribution = en.value();
|
|
|
|
auto type = attribution.getType().dyn_cast<MemRefType>();
|
|
assert(type && type.hasStaticShape() && "unexpected type in attribution");
|
|
|
|
uint64_t numElements = type.getNumElements();
|
|
|
|
auto elementType =
|
|
lowering.convertType(type.getElementType()).cast<LLVM::LLVMType>();
|
|
auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements);
|
|
std::string name =
|
|
llvm::formatv("__wg_{0}_{1}", gpuFuncOp.getName(), en.index());
|
|
auto globalOp = rewriter.create<LLVM::GlobalOp>(
|
|
gpuFuncOp.getLoc(), arrayType, /*isConstant=*/false,
|
|
LLVM::Linkage::Internal, name, /*value=*/Attribute(),
|
|
gpu::GPUDialect::getWorkgroupAddressSpace());
|
|
workgroupBuffers.push_back(globalOp);
|
|
}
|
|
|
|
// Rewrite the original GPU function to an LLVM function.
|
|
auto funcType = lowering.convertType(gpuFuncOp.getType())
|
|
.cast<LLVM::LLVMType>()
|
|
.getPointerElementTy();
|
|
|
|
// Remap proper input types.
|
|
TypeConverter::SignatureConversion signatureConversion(
|
|
gpuFuncOp.front().getNumArguments());
|
|
for (unsigned i = 0, e = funcType.getFunctionNumParams(); i < e; ++i)
|
|
signatureConversion.addInputs(i, funcType.getFunctionParamType(i));
|
|
|
|
// Create the new function operation. Only copy those attributes that are
|
|
// not specific to function modeling.
|
|
SmallVector<NamedAttribute, 4> attributes;
|
|
for (const auto &attr : gpuFuncOp.getAttrs()) {
|
|
if (attr.first.is(SymbolTable::getSymbolAttrName()) ||
|
|
attr.first.is(impl::getTypeAttrName()) ||
|
|
attr.first.is(gpu::GPUFuncOp::getNumWorkgroupAttributionsAttrName()))
|
|
continue;
|
|
attributes.push_back(attr);
|
|
}
|
|
auto llvmFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
|
|
gpuFuncOp.getLoc(), gpuFuncOp.getName(), funcType,
|
|
LLVM::Linkage::External, attributes);
|
|
|
|
{
|
|
// Insert operations that correspond to converted workgroup and private
|
|
// memory attributions to the body of the function. This must operate on
|
|
// the original function, before the body region is inlined in the new
|
|
// function to maintain the relation between block arguments and the
|
|
// parent operation that assigns their semantics.
|
|
OpBuilder::InsertionGuard guard(rewriter);
|
|
|
|
// Rewrite workgroup memory attributions to addresses of global buffers.
|
|
rewriter.setInsertionPointToStart(&gpuFuncOp.front());
|
|
unsigned numProperArguments = gpuFuncOp.getNumArguments();
|
|
auto i32Type = LLVM::LLVMType::getInt32Ty(lowering.getDialect());
|
|
|
|
Value zero = nullptr;
|
|
if (!workgroupBuffers.empty())
|
|
zero = rewriter.create<LLVM::ConstantOp>(loc, i32Type,
|
|
rewriter.getI32IntegerAttr(0));
|
|
for (auto en : llvm::enumerate(workgroupBuffers)) {
|
|
LLVM::GlobalOp global = en.value();
|
|
Value address = rewriter.create<LLVM::AddressOfOp>(loc, global);
|
|
auto elementType = global.getType().getArrayElementType();
|
|
Value memory = rewriter.create<LLVM::GEPOp>(
|
|
loc, elementType.getPointerTo(global.addr_space().getZExtValue()),
|
|
address, ArrayRef<Value>{zero, zero});
|
|
|
|
// Build a memref descriptor pointing to the buffer to plug with the
|
|
// existing memref infrastructure. This may use more registers than
|
|
// otherwise necessary given that memref sizes are fixed, but we can try
|
|
// and canonicalize that away later.
|
|
Value attribution = gpuFuncOp.getWorkgroupAttributions()[en.index()];
|
|
auto type = attribution.getType().cast<MemRefType>();
|
|
auto descr = MemRefDescriptor::fromStaticShape(rewriter, loc, lowering,
|
|
type, memory);
|
|
signatureConversion.remapInput(numProperArguments + en.index(), descr);
|
|
}
|
|
|
|
// Rewrite private memory attributions to alloca'ed buffers.
|
|
unsigned numWorkgroupAttributions =
|
|
gpuFuncOp.getNumWorkgroupAttributions();
|
|
auto int64Ty = LLVM::LLVMType::getInt64Ty(lowering.getDialect());
|
|
for (auto en : llvm::enumerate(gpuFuncOp.getPrivateAttributions())) {
|
|
Value attribution = en.value();
|
|
auto type = attribution.getType().cast<MemRefType>();
|
|
assert(type && type.hasStaticShape() &&
|
|
"unexpected type in attribution");
|
|
|
|
// Explicitly drop memory space when lowering private memory
|
|
// attributions since NVVM models it as `alloca`s in the default
|
|
// memory space and does not support `alloca`s with addrspace(5).
|
|
auto ptrType = lowering.convertType(type.getElementType())
|
|
.cast<LLVM::LLVMType>()
|
|
.getPointerTo();
|
|
Value numElements = rewriter.create<LLVM::ConstantOp>(
|
|
gpuFuncOp.getLoc(), int64Ty,
|
|
rewriter.getI64IntegerAttr(type.getNumElements()));
|
|
Value allocated = rewriter.create<LLVM::AllocaOp>(
|
|
gpuFuncOp.getLoc(), ptrType, numElements, /*alignment=*/0);
|
|
auto descr = MemRefDescriptor::fromStaticShape(rewriter, loc, lowering,
|
|
type, allocated);
|
|
signatureConversion.remapInput(
|
|
numProperArguments + numWorkgroupAttributions + en.index(), descr);
|
|
}
|
|
}
|
|
|
|
// Move the region to the new function, update the entry block signature.
|
|
rewriter.inlineRegionBefore(gpuFuncOp.getBody(), llvmFuncOp.getBody(),
|
|
llvmFuncOp.end());
|
|
rewriter.applySignatureConversion(&llvmFuncOp.getBody(),
|
|
signatureConversion);
|
|
|
|
{
|
|
// For memref-typed arguments, insert the relevant loads in the beginning
|
|
// of the block to comply with the LLVM dialect calling convention. This
|
|
// needs to be done after signature conversion to get the right types.
|
|
OpBuilder::InsertionGuard guard(rewriter);
|
|
Block &block = llvmFuncOp.front();
|
|
rewriter.setInsertionPointToStart(&block);
|
|
|
|
for (auto en : llvm::enumerate(gpuFuncOp.getType().getInputs())) {
|
|
if (!en.value().isa<MemRefType>() &&
|
|
!en.value().isa<UnrankedMemRefType>())
|
|
continue;
|
|
|
|
BlockArgument arg = block.getArgument(en.index());
|
|
Value loaded = rewriter.create<LLVM::LoadOp>(loc, arg);
|
|
rewriter.replaceUsesOfBlockArgument(arg, loaded);
|
|
}
|
|
}
|
|
|
|
rewriter.eraseOp(gpuFuncOp);
|
|
return matchSuccess();
|
|
}
|
|
};
|
|
|
|
struct GPUReturnOpLowering : public LLVMOpLowering {
|
|
GPUReturnOpLowering(LLVMTypeConverter &typeConverter)
|
|
: LLVMOpLowering(gpu::ReturnOp::getOperationName(),
|
|
typeConverter.getDialect()->getContext(),
|
|
typeConverter) {}
|
|
|
|
PatternMatchResult
|
|
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, operands,
|
|
ArrayRef<Block *>());
|
|
return matchSuccess();
|
|
}
|
|
};
|
|
|
|
/// Import the GPU Ops to NVVM Patterns.
|
|
#include "GPUToNVVM.cpp.inc"
|
|
|
|
/// A pass that replaces all occurrences of GPU device operations with their
|
|
/// corresponding NVVM equivalent.
|
|
///
|
|
/// This pass only handles device code and is not meant to be run on GPU host
|
|
/// code.
|
|
class LowerGpuOpsToNVVMOpsPass
|
|
: public OperationPass<LowerGpuOpsToNVVMOpsPass, gpu::GPUModuleOp> {
|
|
public:
|
|
void runOnOperation() override {
|
|
gpu::GPUModuleOp m = getOperation();
|
|
OwningRewritePatternList patterns;
|
|
NVVMTypeConverter converter(m.getContext());
|
|
populateStdToLLVMConversionPatterns(converter, patterns);
|
|
populateGpuToNVVMConversionPatterns(converter, patterns);
|
|
ConversionTarget target(getContext());
|
|
target.addIllegalDialect<gpu::GPUDialect>();
|
|
target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
|
|
LLVM::ExpOp>();
|
|
target.addIllegalOp<FuncOp>();
|
|
target.addLegalDialect<LLVM::LLVMDialect>();
|
|
target.addLegalDialect<NVVM::NVVMDialect>();
|
|
// TODO(csigg): Remove once we support replacing non-root ops.
|
|
target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
|
|
if (failed(applyPartialConversion(m, target, patterns, &converter)))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
|
|
} // anonymous namespace
|
|
|
|
void mlir::populateGpuToNVVMConversionPatterns(
|
|
LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
|
|
populateWithGenerated(converter.getDialect()->getContext(), &patterns);
|
|
patterns
|
|
.insert<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp,
|
|
NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>,
|
|
GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp,
|
|
NVVM::BlockDimYOp, NVVM::BlockDimZOp>,
|
|
GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp,
|
|
NVVM::BlockIdYOp, NVVM::BlockIdZOp>,
|
|
GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp,
|
|
NVVM::GridDimYOp, NVVM::GridDimZOp>,
|
|
GPUAllReduceOpLowering, GPUShuffleOpLowering, GPUFuncOpLowering,
|
|
GPUReturnOpLowering>(converter);
|
|
patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__nv_fabsf",
|
|
"__nv_fabs");
|
|
patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf",
|
|
"__nv_ceil");
|
|
patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__nv_cosf",
|
|
"__nv_cos");
|
|
patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__nv_expf",
|
|
"__nv_exp");
|
|
}
|
|
|
|
std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
|
|
mlir::createLowerGpuOpsToNVVMOpsPass() {
|
|
return std::make_unique<LowerGpuOpsToNVVMOpsPass>();
|
|
}
|
|
|
|
static PassRegistration<LowerGpuOpsToNVVMOpsPass>
|
|
pass("convert-gpu-to-nvvm", "Generate NVVM operations for gpu operations");
|