0a8e3dd432
Rename interface functions as follows: * `hasTensorSemantics` -> `hasPureTensorSemantics` * `hasBufferSemantics` -> `hasPureBufferSemantics` These two functions return "true" if the op has tensor/buffer operands but not buffer/tensor operands. Also drop the "ranked" part from the interface, i.e., do not distinguish between ranked/unranked types. The new function names describe the functions more accurately. They also align their semantics with the notion of "tensor semantics" with the bufferization framework. (An op is supposed to be bufferized if it has tensor operands, and we don't care if it also has memref operands.) This change is in preparation of #75273, which adds `BufferizableOpInterface::hasTensorSemantics`. By renaming the functions in the `DestinationStyleOpInterface`, we can avoid name clashes between the two interfaces.
278 lines
11 KiB
C++
278 lines
11 KiB
C++
//===- CreateAsyncGroups.cpp - Create async device copies -----------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/NVGPU/Transforms/Transforms.h"
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
|
|
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
|
|
#include "mlir/Dialect/NVGPU/Transforms/Utils.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
|
|
using namespace mlir;
|
|
|
|
/// Return "true" if the given vector transfer op is contiguous and suitable
|
|
/// for replacement with an async copy.
|
|
template <typename OpTy>
|
|
static bool isContiguousXferOp(OpTy op) {
|
|
return op.getPermutationMap().isMinorIdentity() && op.isDimInBounds(0) &&
|
|
op.hasPureBufferSemantics() &&
|
|
isLastMemrefDimUnitStride(
|
|
cast<MemRefType>(nvgpu::getMemrefOperand(op).getType()));
|
|
}
|
|
|
|
/// Return "true" if the given op is a contiguous and suitable
|
|
/// vector.transfer_write or vector.store op.
|
|
static bool isContiguousStore(Operation *write) {
|
|
if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(write))
|
|
return isContiguousXferOp(transferWrite) && !transferWrite.getMask();
|
|
// vector.store are always contiguous.
|
|
return isa<vector::StoreOp>(write);
|
|
}
|
|
|
|
/// Return "true" if the given op is a contiguous and suitable
|
|
/// vector.transfer_read or vector.load op.
|
|
static bool isContiguousRead(Operation *read) {
|
|
if (auto transferRead = dyn_cast<vector::TransferReadOp>(read))
|
|
return isContiguousXferOp(transferRead);
|
|
// vector.load are always contiguous.
|
|
return isa<vector::LoadOp>(read);
|
|
}
|
|
|
|
namespace {
|
|
/// A vector.create_mask op and extract position.
|
|
struct TransferMask {
|
|
vector::CreateMaskOp createMaskOp;
|
|
SmallVector<int64_t> extractPosition;
|
|
};
|
|
} // namespace
|
|
|
|
/// If the given vector load op has a mask that is defined by
|
|
/// vector.create_mask, return that op.
|
|
static FailureOr<TransferMask> getMaskOp(Operation *loadOp) {
|
|
auto transferRead = dyn_cast<vector::TransferReadOp>(loadOp);
|
|
if (!transferRead || !transferRead.getMask())
|
|
return TransferMask{{}, {}};
|
|
assert(transferRead.getMask().getType().getRank() == 1 &&
|
|
"expected 1-D mask");
|
|
|
|
// Case 1: Mask is the result of a vector.create_mask.
|
|
if (auto maskOp =
|
|
transferRead.getMask().getDefiningOp<vector::CreateMaskOp>())
|
|
return TransferMask{maskOp, {}};
|
|
|
|
// Case 2: Mask is the result of a vector.extract(vector.create_mask).
|
|
if (auto extractOp =
|
|
transferRead.getMask().getDefiningOp<vector::ExtractOp>())
|
|
if (auto maskOp =
|
|
extractOp.getVector().getDefiningOp<vector::CreateMaskOp>())
|
|
return TransferMask{maskOp,
|
|
SmallVector<int64_t>(extractOp.getStaticPosition())};
|
|
|
|
// All other cases: not supported.
|
|
return failure();
|
|
}
|
|
|
|
/// Build an SSA value that represents the number of read elements.
|
|
static Value buildNumReadElements(OpBuilder &b, Location loc,
|
|
Operation *readOp) {
|
|
FailureOr<TransferMask> transferMask = getMaskOp(readOp);
|
|
assert(succeeded(transferMask) && "invalid transfer mask");
|
|
|
|
// No mask => no num_read_elements.
|
|
if (!transferMask->createMaskOp)
|
|
return Value();
|
|
|
|
// No extract: return size of "ones" segment in the mask.
|
|
if (transferMask->extractPosition.empty()) {
|
|
assert(transferMask->createMaskOp.getNumOperands() == 1 &&
|
|
"expected single operand");
|
|
return transferMask->createMaskOp.getOperand(0);
|
|
}
|
|
|
|
// vector.extract(vector.create_mask).
|
|
// If extract_pos < num_ones, take number of elements from the least
|
|
// significant dimension. (Do this for all dimensions and bit-AND the
|
|
// conditions.)
|
|
assert(transferMask->createMaskOp.getVectorType().getRank() -
|
|
transferMask->extractPosition.size() ==
|
|
1 &&
|
|
"expected N-D -> (N-1)-D extract");
|
|
Value cond;
|
|
// Note: There is one more `sz` than `pos`. The loop end with the last `pos`.
|
|
for (auto [pos, sz] : llvm::zip(transferMask->extractPosition,
|
|
transferMask->createMaskOp->getOperands())) {
|
|
Value cmp =
|
|
b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt,
|
|
b.create<arith::ConstantIndexOp>(loc, pos), sz);
|
|
if (!cond) {
|
|
cond = cmp;
|
|
continue;
|
|
}
|
|
cond = b.create<arith::AndIOp>(loc, cmp, cond);
|
|
}
|
|
return b.create<arith::SelectOp>(
|
|
loc, cond, transferMask->createMaskOp->getOperands().back(),
|
|
b.create<arith::ConstantIndexOp>(loc, 0));
|
|
}
|
|
|
|
/// Return "true" if the conversion to async copy is supported by "async copy".
|
|
static bool resultsInSupportedAsyncCopy(MemRefType memrefType,
|
|
VectorType vecType) {
|
|
assert(vecType.getRank() == 1 && "expected 1-D vector");
|
|
constexpr int64_t kSupportedCpAsyncAlignmentsInBytes[3] = {4, 8, 16};
|
|
|
|
// Condition 1: the copy size must be supported.
|
|
bool supportedCopySize = false;
|
|
int64_t numElements = vecType.getNumElements();
|
|
Type elementType = vecType.getElementType();
|
|
for (int64_t alignmentInBytes : kSupportedCpAsyncAlignmentsInBytes) {
|
|
if (alignmentInBytes * 8 ==
|
|
numElements * elementType.getIntOrFloatBitWidth()) {
|
|
supportedCopySize = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!supportedCopySize)
|
|
return false;
|
|
|
|
// TODO: Condition 2: the alignments must be supported. For cp.async the
|
|
// NVIDIA doc (section 6.4.1) says: "The address must be naturally aligned to
|
|
// a multiple of the access size. If an address is not properly aligned, the
|
|
// resulting behavior is undefined.".
|
|
return true;
|
|
}
|
|
|
|
void nvgpu::createAsyncGroups(RewriterBase &rewriter, Operation *op,
|
|
bool bypassL1) {
|
|
llvm::SmallSetVector<Operation *, 16> copyToSharedMem;
|
|
|
|
// Look for all the copy that can be converted to async copy ops.
|
|
op->walk([&](Operation *writeOp) {
|
|
// Look for contiguous 1D vector store into shared memory.
|
|
if (!isContiguousStore(writeOp))
|
|
return;
|
|
Value vectorVal = nvgpu::getValueStored(writeOp);
|
|
if (cast<VectorType>(vectorVal.getType()).getRank() != 1)
|
|
return;
|
|
Value storeBase = nvgpu::getMemrefOperand(writeOp);
|
|
if (!nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(
|
|
cast<MemRefType>(storeBase.getType())))
|
|
return;
|
|
|
|
// The stored vector must originate from a contiguous 1D vector load.
|
|
Operation *readOp = vectorVal.getDefiningOp();
|
|
if (readOp == nullptr || !isContiguousRead(readOp))
|
|
return;
|
|
Value loadBase = nvgpu::getMemrefOperand(readOp);
|
|
// Should be reading from global memory (not shared memory).
|
|
if (nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(
|
|
cast<MemRefType>(loadBase.getType())))
|
|
return;
|
|
|
|
// Look for compatible mask and padding.
|
|
if (auto transferRead = dyn_cast<vector::TransferReadOp>(readOp)) {
|
|
if (Value mask = transferRead.getMask()) {
|
|
if (getConstantIntValue(transferRead.getPadding()) ==
|
|
static_cast<int64_t>(0))
|
|
return;
|
|
if (failed(getMaskOp(readOp)))
|
|
return;
|
|
}
|
|
}
|
|
|
|
// Check whether both accesses are supported before we emit: this is
|
|
// necessary to ensure the correctness of DeviceAsyncCopyOp.
|
|
VectorType vecType = cast<VectorType>(vectorVal.getType());
|
|
|
|
if (!resultsInSupportedAsyncCopy(cast<MemRefType>(loadBase.getType()),
|
|
vecType) ||
|
|
!resultsInSupportedAsyncCopy(cast<MemRefType>(storeBase.getType()),
|
|
vecType))
|
|
return;
|
|
|
|
copyToSharedMem.insert(writeOp);
|
|
return;
|
|
});
|
|
|
|
while (!copyToSharedMem.empty()) {
|
|
// Start a group with the first write.
|
|
SmallVector<Operation *> group;
|
|
Operation *writeOp = *copyToSharedMem.begin();
|
|
copyToSharedMem.remove(writeOp);
|
|
group.push_back(writeOp);
|
|
Operation *nextNode = writeOp;
|
|
|
|
// Look in the next nodes for more copies to add to the same group.
|
|
while ((nextNode = nextNode->getNextNode())) {
|
|
// Ignore ops without side effects.
|
|
auto memInterface = dyn_cast<MemoryEffectOpInterface>(nextNode);
|
|
if (memInterface && memInterface.hasNoEffect() &&
|
|
!nextNode->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
|
|
continue;
|
|
// Ignore read from a different address space.
|
|
if (isa<vector::TransferReadOp, vector::LoadOp>(nextNode)) {
|
|
Operation *readOp = nextNode;
|
|
Value memrefOperand = nvgpu::getMemrefOperand(readOp);
|
|
if (!nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(
|
|
cast<MemRefType>(memrefOperand.getType()))) {
|
|
continue;
|
|
}
|
|
}
|
|
if (copyToSharedMem.count(nextNode)) {
|
|
// Found another copy, add it to the group.
|
|
copyToSharedMem.remove(nextNode);
|
|
group.push_back(nextNode);
|
|
continue;
|
|
}
|
|
// If the op is something else stop the accumulating op in the group.
|
|
break;
|
|
}
|
|
|
|
// Emit the group.
|
|
SmallVector<Value> tokens;
|
|
for (Operation *writeOp : group) {
|
|
rewriter.setInsertionPoint(writeOp);
|
|
Value vectorVal = nvgpu::getValueStored(writeOp);
|
|
auto vectorType = cast<VectorType>(vectorVal.getType());
|
|
int64_t numElements = vectorType.getNumElements();
|
|
Operation *readOp = vectorVal.getDefiningOp();
|
|
Value storeBase = nvgpu::getMemrefOperand(writeOp);
|
|
Value loadBase = nvgpu::getMemrefOperand(readOp);
|
|
Value numReadElements =
|
|
buildNumReadElements(rewriter, writeOp->getLoc(), readOp);
|
|
auto dstMemref = cast<MemRefType>(storeBase.getType());
|
|
int64_t sizeInBytes =
|
|
(dstMemref.getElementTypeBitWidth() * numElements) / 8;
|
|
// bypass_l1 only possible with 16 byte transfer.
|
|
Value token = rewriter.create<nvgpu::DeviceAsyncCopyOp>(
|
|
writeOp->getLoc(), nvgpu::DeviceAsyncTokenType::get(op->getContext()),
|
|
/*dst=*/storeBase, /*dstIndices=*/nvgpu::getIndices(writeOp),
|
|
/*src=*/loadBase,
|
|
/*srcIndices=*/nvgpu::getIndices(readOp),
|
|
/*dstElements=*/rewriter.getIndexAttr(numElements),
|
|
/*srcElements=*/numReadElements,
|
|
/*bypassL1=*/bypassL1 && sizeInBytes == 16 ? rewriter.getUnitAttr()
|
|
: UnitAttr());
|
|
tokens.push_back(token);
|
|
}
|
|
|
|
// Create the group and wait for it right after.
|
|
Value groupToken = rewriter.create<nvgpu::DeviceAsyncCreateGroupOp>(
|
|
op->getLoc(), nvgpu::DeviceAsyncTokenType::get(op->getContext()),
|
|
tokens);
|
|
rewriter.create<nvgpu::DeviceAsyncWaitOp>(op->getLoc(), groupToken,
|
|
nullptr);
|
|
// Clean up old stores.
|
|
for (Operation *writeOp : group)
|
|
rewriter.eraseOp(writeOp);
|
|
}
|
|
}
|