[mlir][gpu] Refactor functions for workgroup and private buffer attributions.

[whchung]

Summary

Consolidate interfaces adding workgroup and private buffer attributions in GPU
dialect.

Note all private buffer attributions must follow workgroup buffer attributions.

[herhut]

Thanks for adding this!

mlir/include/mlir/Dialect/GPU/GPUOps.td

215

It is not clear to me why we have two versions of this function. While you have not caused this, could you clean it up and unify them? One could implement addWorkgroupAttribution(ArrayRef<int64_t> shape, Type elementType) by means of BlockArgument addWorkgroupAttribution(Type type).

I see only one use of the former form in MemoryPromotion. Changing that use to the type based one and removing the shape/elementType versions would be best.

Also, both of these could be implemented in the cpp file and only be declared here.

220

Use getNumWorkgroupAttributions here?

[whchung]

Address review comments.

[whchung]

mlir/include/mlir/Dialect/GPU/GPUOps.td

215

@herhut I was somewhat puzzled by the two versions of interfaces too. I've revised the patch so we we consolidate the interfaces. Could you help conduct another round of review? Thanks.

[herhut]

Thanks for the cleanup.

mlir/include/mlir/Dialect/GPU/GPUOps.td

195

These getters that give access to a range should remain in the header to enable inlining (and maybe avoid materializing the ArrayRef).

222

I don't think this is needed. The contract could simple be that all trailing operands after function arguments and workgroup attributions are private. Storing an extra count does not add benefit. You would then need to verify that arguments + all attributions = operand count.

So let's just drop the extra attribute.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

478

This would just be {begin, getBody().front().args_end()}

[herhut]

@whchung Are you still interested in moving this forward?

[whchung]

Address code review comments.

[whchung]

In D79508#2043577, @herhut wrote:

@whchung Are you still interested in moving this forward?

@herhut I just revised the patch to address your review comments. Getting rid of the counter for private allocations is less than ideal in my applications but I've found ways to get around them. Could you give this patch another look? Thanks.

[herhut]

In D79508#2044128, @whchung wrote:

In D79508#2043577, @herhut wrote:

@whchung Are you still interested in moving this forward?

@herhut I just revised the patch to address your review comments. Getting rid of the counter for private allocations is less than ideal in my applications but I've found ways to get around them. Could you give this patch another look? Thanks.

I think you misunderstood my comment. I liked having the getters and interface for private attributions. My point was that you do not need to store a count as an attribute for the private attributions but instead could always compute the number of private attributes via getNumOperands() - getType().getNumInputs() - getNumWorkgroupAttributions(). Sorry for not being clear.

Using that approach should also cover your use cases?

[whchung]

In D79508#2044267, @herhut wrote:

In D79508#2044128, @whchung wrote:

In D79508#2043577, @herhut wrote:

@whchung Are you still interested in moving this forward?

@herhut I just revised the patch to address your review comments. Getting rid of the counter for private allocations is less than ideal in my applications but I've found ways to get around them. Could you give this patch another look? Thanks.

I think you misunderstood my comment. I liked having the getters and interface for private attributions. My point was that you do not need to store a count as an attribute for the private attributions but instead could always compute the number of private attributes via getNumOperands() - getType().getNumInputs() - getNumWorkgroupAttributions(). Sorry for not being clear.

Using that approach should also cover your use cases?

@herhut Yes, that's exactly how I compute the number of private attributes right now in my application.

BTW it seems nowadays pre-merge checks would always fail on Windows due to logic in SPIR-V dialect.

[2020-05-19T14:57:38.942Z] tools\mlir\include\mlir/Dialect/SPIRV/SPIRVSerialization.inc(4334): fatal error C1061: compiler limit: blocks nested too deeply

[mravishankar]

BTW it seems nowadays pre-merge checks would always fail on Windows due to logic in SPIR-V dialect.

[2020-05-19T14:57:38.942Z] tools\mlir\include\mlir/Dialect/SPIRV/SPIRVSerialization.inc(4334): fatal error C1061: compiler limit: blocks nested too deeply

This should be fixed by patch https://github.com/llvm/llvm-project/commit/d5b1643c74eeae327d85c75fe79fd98edb1014f9

[herhut]

In D79508#2044398, @whchung wrote:

In D79508#2044267, @herhut wrote:

In D79508#2044128, @whchung wrote:

In D79508#2043577, @herhut wrote:

@whchung Are you still interested in moving this forward?

@herhut I just revised the patch to address your review comments. Getting rid of the counter for private allocations is less than ideal in my applications but I've found ways to get around them. Could you give this patch another look? Thanks.

I think you misunderstood my comment. I liked having the getters and interface for private attributions. My point was that you do not need to store a count as an attribute for the private attributions but instead could always compute the number of private attributes via getNumOperands() - getType().getNumInputs() - getNumWorkgroupAttributions(). Sorry for not being clear.

Using that approach should also cover your use cases?

@herhut Yes, that's exactly how I compute the number of private attributes right now in my application.

Why not add a getNumPrivateAttributions to the GPUFunc for this? Then the interface is exactly as before.

With that change, I think this is good to land.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

478

I would prefer if new private attributions would also be added at the end of the existing private attributes, rather than at the front. You could just use addArgument for this.

[whchung]

Reinstate getNumPrivateAttributions() in GPUFunc.

Use it in verifier logic, and addPrivateAttribution() logic.

[whchung]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

478

@herhut addArgument might not guard the case addPrivateAttribution is used prior to addWorkgroupAttribution. Now I have revised the patch to re-introduce getNumPrivateAttribution we can leverage it here.

[herhut]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

478

I don't understand this. Private attributions are always at the end. So using getBody().front().addArgument(...) would insert them at the end, where they belong. If you then use addWorkgroupAttribution, it will insert in front of the private one, as getNumWorkgroupAttributions would return 0, so it inserts directly after the function arguments.

[whchung]

Simplify logic.

[whchung]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

478

Yes you are right. I've revised the patch with simplified logic.

[herhut]

Thanks!

[rriddle]

mlir/include/mlir/Dialect/GPU/GPUOps.td

203

This doesn't make sense to me, when does GPUFuncOp ever have operands? Did you mean to use numArguments here? This looks like a lack of testing coverage if so.

[whchung]

mlir/include/mlir/Dialect/GPU/GPUOps.td

203

Thanks for catching this. Indeed this is a bug not covered by any existing tests. I'll submit another patch addressing this with a test.

[whchung]

mlir/include/mlir/Dialect/GPU/GPUOps.td

203

@rriddle I submitted D80766 to fix this.

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, StringRef name, "
               "FunctionType type, ArrayRef<Type> workgroupAttributions = {}, "
               "ArrayRef<Type> privateAttributions = {}, "
               "ArrayRef<NamedAttribute> attrs = {}">
   ];
 
   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() {
       return getAttrOfType<UnitAttr>(GPUDialect::getKernelFuncAttrName()) !=
              nullptr;
     }
 
     /// Change the type of this function in place. This is an extremely
let extraClassDeclaration = [{
 
     /// Returns a list of block arguments that correspond to buffers located in
     /// the workgroup memory
     ArrayRef<BlockArgument> getWorkgroupAttributions() {
       auto begin =
           std::next(getBody().front().args_begin(), getType().getNumInputs());
       auto end = std::next(begin, getNumWorkgroupAttributions());
       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 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() {
       return "workgroup_attributions";
     }
 
     // FunctionLike trait needs access to the functions below.
     friend class OpTrait::FunctionLike<GPUFuncOp>;
 
     /// Hooks for the input/output type enumeration in FunctionLike .
     unsigned getNumFuncArguments() { return getType().getNumInputs(); }
     unsigned getNumFuncResults() { return getType().getNumResults(); }
 
     /// Returns the keywords used in the custom syntax for this Op.
     static StringRef getWorkgroupKeyword() { return "workgroup"; }

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

static LogicalResult verify(LaunchFuncOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // 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.
+  auto workgroupAttrCount = getNumWorkgroupAttributions();
+
+  return getBody().front().insertArgument(
+      getType().getNumInputs() + workgroupAttrCount, type);
 }
 
 void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
                       StringRef name, FunctionType type,
                       ArrayRef<Type> workgroupAttributions,
                       ArrayRef<Type> privateAttributions,
                       ArrayRef<NamedAttribute> attrs) {
   result.addAttribute(SymbolTable::getSymbolAttrName(),

mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp

private:
   // Creates dimension op of type T, with the result casted to int32.
   template <typename T> Value getDimOp(StringRef dimension) {
     Value dim = create<T>(indexType, rewriter.getStringAttr(dimension));
     return create<IndexCastOp>(int32Type, dim);
   }
 
   /// 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);
     return funcOp.addWorkgroupAttribution(bufferType);
   }
 
   /// Returns an accumulator factory using either the op attribute or the body
   /// region.

mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp

 
 /// Promotes a function argument to workgroup memory in the given function. The
 /// copies will be inserted in the beginning and in the end of the function.
 void mlir::promoteToWorkgroupMemory(GPUFuncOp op, unsigned arg) {
   Value value = op.getArgument(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.
   value.replaceAllUsesWith(attribution);
   insertCopies(op.getBody(), op.getLoc(), value, attribution);
 }