[Mlir-commits] [mlir] ad39816 - [mlir][gpu] Refactor functions for workgroup and private buffer attributions.
Wen-Heng Chung
llvmlistbot at llvm.org
Wed May 20 14:20:38 PDT 2020
Author: Wen-Heng (Jack) Chung
Date: 2020-05-20T16:20:27-05:00
New Revision: ad398164bac0e4a63588005e7061537518557eab
URL: https://github.com/llvm/llvm-project/commit/ad398164bac0e4a63588005e7061537518557eab
DIFF: https://github.com/llvm/llvm-project/commit/ad398164bac0e4a63588005e7061537518557eab.diff
LOG: [mlir][gpu] Refactor functions for workgroup and private buffer attributions.
Summary:
Consolidate interfaces adding workgroup and private buffer attributions in GPU
dialect.
Note all private buffer attributions must follow workgroup buffer attributions.
Reviewers: herhut
Subscribers: mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, csigg, arpith-jacob, mgester, lucyrfox, liufengdb, stephenneuendorffer, Joonsoo, grosul1, frgossen, Kayjukh, llvm-commits
Tags: #llvm, #mlir
Differential Revision: https://reviews.llvm.org/D79508
Added:
Modified:
mlir/include/mlir/Dialect/GPU/GPUOps.td
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index 524aa2041f49..3d623a64f425 100644
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -161,10 +161,6 @@ def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
];
let extraClassDeclaration = [{
- /// Adds a workgroup attribution of the MemRef type with the given shape and
- /// element type.
- Value addWorkgroupAttribution(ArrayRef<int64_t> shape, Type elementType);
-
/// Returns `true` if the GPU function defined by this Op is a kernel, i.e.
/// it is intended to be launched from host.
bool isKernel() {
@@ -198,25 +194,31 @@ def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
return {begin, end};
}
- // Adds a new block argument that corresponds to buffers located in
- // workgroup memory.
- BlockArgument addWorkgroupAttribution(Type type) {
- auto attrName = getNumWorkgroupAttributionsAttrName();
- auto attr = getAttrOfType<IntegerAttr>(attrName);
- setAttr(attrName, IntegerAttr::get(attr.getType(), attr.getValue() + 1));
- return getBody().front().insertArgument(
- getType().getNumInputs() + attr.getInt(), type);
- }
+ /// Adds a new block argument that corresponds to buffers located in
+ /// workgroup memory.
+ BlockArgument addWorkgroupAttribution(Type type);
+ /// Returns the number of buffers located in the private memory.
+ unsigned getNumPrivateAttributions() {
+ return getOperation()->getNumOperands() - getType().getNumInputs() -
+ getNumWorkgroupAttributions();
+ }
+
/// Returns a list of block arguments that correspond to buffers located in
/// the private memory.
ArrayRef<BlockArgument> getPrivateAttributions() {
+ // Buffers on the private memory always come after buffers on the workgroup
+ // memory.
auto begin =
std::next(getBody().front().args_begin(),
getType().getNumInputs() + getNumWorkgroupAttributions());
return {begin, getBody().front().args_end()};
}
+ /// Adds a new block argument that corresponds to buffers located in
+ /// private memory.
+ BlockArgument addPrivateAttribution(Type type);
+
/// Returns the name of the attribute containing the number of buffers
/// located in the workgroup memory.
static StringRef getNumWorkgroupAttributionsAttrName() {
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 71cec77a0a4a..9a5f6fcefbca 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -457,22 +457,22 @@ static LogicalResult verify(LaunchFuncOp op) {
// GPUFuncOp
//===----------------------------------------------------------------------===//
-/// Adds a workgroup attribution to "op" of the MemRef type with the given shape
-/// and element type.
-Value GPUFuncOp::addWorkgroupAttribution(ArrayRef<int64_t> shape,
- Type elementType) {
- unsigned pos = getNumFuncArguments() + getNumWorkgroupAttributions();
- Block &bodyBlock = body().front();
- Value attribution = bodyBlock.insertArgument(
- std::next(bodyBlock.args_begin(), pos),
- MemRefType::get(shape, elementType, /*affineMapComposition=*/{},
- GPUDialect::getWorkgroupAddressSpace()));
- auto numWorkgroupBuffersAttr =
- getAttrOfType<IntegerAttr>(getNumWorkgroupAttributionsAttrName());
- setAttr(getNumWorkgroupAttributionsAttrName(),
- IntegerAttr::get(numWorkgroupBuffersAttr.getType(),
- numWorkgroupBuffersAttr.getValue() + 1));
- return attribution;
+/// Adds a new block argument that corresponds to buffers located in
+/// workgroup memory.
+BlockArgument GPUFuncOp::addWorkgroupAttribution(Type type) {
+ auto attrName = getNumWorkgroupAttributionsAttrName();
+ auto attr = getAttrOfType<IntegerAttr>(attrName);
+ setAttr(attrName, IntegerAttr::get(attr.getType(), attr.getValue() + 1));
+ return getBody().front().insertArgument(
+ getType().getNumInputs() + attr.getInt(), type);
+}
+
+/// Adds a new block argument that corresponds to buffers located in
+/// private memory.
+BlockArgument GPUFuncOp::addPrivateAttribution(Type type) {
+ // Buffers on the private memory always come after buffers on the workgroup
+ // memory.
+ return getBody().front().addArgument(type);
}
void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
@@ -675,10 +675,13 @@ static LogicalResult verifyAttributions(Operation *op,
LogicalResult GPUFuncOp::verifyBody() {
unsigned numFuncArguments = getNumArguments();
unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
+ unsigned numPrivateAttributions = getNumPrivateAttributions();
unsigned numBlockArguments = front().getNumArguments();
- if (numBlockArguments < numFuncArguments + numWorkgroupAttributions)
+ if (numBlockArguments <
+ numFuncArguments + numWorkgroupAttributions + numPrivateAttributions)
return emitOpError() << "expected at least "
- << numFuncArguments + numWorkgroupAttributions
+ << numFuncArguments + numWorkgroupAttributions +
+ numPrivateAttributions
<< " arguments to body region";
ArrayRef<Type> funcArgTypes = getType().getInputs();
diff --git a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
index 57c497f613cb..6f6f1c27241c 100644
--- a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
@@ -151,7 +151,8 @@ struct GpuAllReduceRewriter {
/// Adds type to funcOp's workgroup attributions.
Value createWorkgroupBuffer() {
- int workgroupMemoryAddressSpace = 3;
+ int workgroupMemoryAddressSpace =
+ gpu::GPUDialect::getWorkgroupAddressSpace();
auto bufferType =
MemRefType::get({kSubgroupSize}, valueType, ArrayRef<AffineMap>{},
workgroupMemoryAddressSpace);
diff --git a/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp b/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
index 3e3ea155d960..d99105e04684 100644
--- a/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
@@ -162,8 +162,12 @@ void mlir::promoteToWorkgroupMemory(GPUFuncOp op, unsigned arg) {
auto type = value.getType().dyn_cast<MemRefType>();
assert(type && type.hasStaticShape() && "can only promote memrefs");
- Value attribution =
- op.addWorkgroupAttribution(type.getShape(), type.getElementType());
+ // Get the type of the buffer in the workgroup memory.
+ int workgroupMemoryAddressSpace = gpu::GPUDialect::getWorkgroupAddressSpace();
+ auto bufferType = MemRefType::get(type.getShape(), type.getElementType(), {},
+ workgroupMemoryAddressSpace);
+
+ Value attribution = op.addWorkgroupAttribution(bufferType);
// Replace the uses first since only the original uses are currently present.
// Then insert the copies.
More information about the Mlir-commits
mailing list