[MLIR] Add argument related API to Region

[jurahul]

Summary

• Arguments of the first block of a region are considered region arguments.
• Add API on Region class to deal with these arguments directly instead of using the front() block.
• Changed several instances of existing code that can use this API
• Fixes https://bugs.llvm.org/show_bug.cgi?id=46535

[bondhugula]

Please check the build failures though.

[mehdi_amini]

I'd like to hear from @stephenneuendorffer about how this fits with D80358 ; we may want to assert on the type of Region, or use some sort of interface I'm not sure.

[jurahul]

Looking at https://reviews.llvm.org/D80358, I don't see the definition of region arguments being entry block arguments changing.

[jurahul]

And build failure are clang_tidy warnings for args_begin() etc functions. @mehdi_amini , how are these warnings currently suppressed in say Block.h? Or do we just ignore them? Also, based on the fact that definition of region entry block and its arguments being region arguments is not changing, can you unblock the change?

[mehdi_amini]

In D83599#2148908, @jurahul wrote:

Looking at https://reviews.llvm.org/D80358, I don't see the definition of region arguments being entry block arguments changing.

Yeah I sync'd with Stephen today and there is no concern.

[mehdi_amini]

(you can ignore the clang-tidy warnings here)

[stephenneuendorffer]

In D83599#2149078, @mehdi_amini wrote:

In D83599#2148908, @jurahul wrote:

Looking at https://reviews.llvm.org/D80358, I don't see the definition of region arguments being entry block arguments changing.

Yeah I sync'd with Stephen today and there is no concern.

Yes, I agree. I think it makes sense for Regions to have the same arguments as their blocks. However, the relationship between operation arguments and region arguments is something that is determined by the semantics of each individual operation. With FuncOp, the function arguments and region arguments have the same type and arity, but in general, they may not.

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

unsigned getNumWorkgroupAttributions() {
       return getAttrOfType<IntegerAttr>(getNumWorkgroupAttributionsAttrName())
           .getInt();
     }
 
     /// 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());
+          std::next(getBody().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);
 
     /// Returns the number of buffers located in the private memory.
     unsigned getNumPrivateAttributions() {
-      return getBody().front().getNumArguments() - getType().getNumInputs() -
+      return getBody().getNumArguments() - 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(),
+          std::next(getBody().args_begin(),
                     getType().getNumInputs() + getNumWorkgroupAttributions());
-      return {begin, getBody().front().args_end()};
+      return {begin, getBody().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.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def GenericAtomicRMWOp : Std_Op<"generic_atomic_rmw", [
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "Value memref, ValueRange ivs">
   ];
 
   let extraClassDeclaration = [{
     // The value stored in memref[ivs].
     Value getCurrentValue() {
-      return body().front().getArgument(0);
+      return body().getArgument(0);
     }
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
   }];
 }
 
 def AtomicYieldOp : Std_Op<"atomic_yield", [

mlir/include/mlir/IR/FunctionSupport.h

unsigned getNumArguments() {
     return static_cast<ConcreteType *>(this)->getNumFuncArguments();
   }
 
   unsigned getNumResults() {
     return static_cast<ConcreteType *>(this)->getNumFuncResults();
   }
 
   /// Gets argument.
-  BlockArgument getArgument(unsigned idx) {
-    return getBlocks().front().getArgument(idx);
-  }
+  BlockArgument getArgument(unsigned idx) { return getBody().getArgument(idx); }
 
   /// Support argument iteration.
-  using args_iterator = Block::args_iterator;
-  args_iterator args_begin() { return front().args_begin(); }
-  args_iterator args_end() { return front().args_end(); }
-  Block::BlockArgListType getArguments() { return front().getArguments(); }
+  using args_iterator = Region::args_iterator;
+  args_iterator args_begin() { return getBody().args_begin(); }
+  args_iterator args_end() { return getBody().args_end(); }
+  Block::BlockArgListType getArguments() { return getBody().getArguments(); }
 
   //===--------------------------------------------------------------------===//
   // Argument Attributes
   //===--------------------------------------------------------------------===//
 
   /// FunctionLike operations allow for attaching attributes to each of the
   /// respective function arguments. These argument attributes are stored as
   /// DictionaryAttrs in the main operation attribute dictionary. The name of

mlir/include/mlir/IR/Region.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_REGION_H
 #define MLIR_IR_REGION_H
 
 #include "mlir/IR/Block.h"
 
 namespace mlir {
+class TypeRange;
+template <typename ValueRangeT>
+class ValueTypeRange;
 class BlockAndValueMapping;
 
 /// This class contains a list of basic blocks and a link to the parent
 /// operation it is attached to.
 class Region {
 public:
   Region() = default;
   explicit Region(Operation *container);
public:
   Block &front() { return blocks.front(); }
 
   /// getSublistAccess() - Returns pointer to member of region.
   static BlockListType Region::*getSublistAccess(Block *) {
     return &Region::blocks;
   }
 
   //===--------------------------------------------------------------------===//
+  // Argument Handling
+  //===--------------------------------------------------------------------===//
+
+  // This is the list of arguments to the block.
+  using BlockArgListType = MutableArrayRef<BlockArgument>;
+  BlockArgListType getArguments() {
+    return empty() ? BlockArgListType() : front().getArguments();
+  }
+  using args_iterator = BlockArgListType::iterator;
+  using reverse_args_iterator = BlockArgListType::reverse_iterator;
+  args_iterator args_begin() { return getArguments().begin(); }
+  args_iterator args_end() { return getArguments().end(); }
+  reverse_args_iterator args_rbegin() { return getArguments().rbegin(); }
+  reverse_args_iterator args_rend() { return getArguments().rend(); }
+
+  bool args_empty() { return getArguments().empty(); }
+
+  /// Add one value to the argument list.
+  BlockArgument addArgument(Type type) { return front().addArgument(type); }
+
+  /// Insert one value to the position in the argument list indicated by the
+  /// given iterator. The existing arguments are shifted. The block is expected
+  /// not to have predecessors.
+  BlockArgument insertArgument(args_iterator it, Type type) {
+    return front().insertArgument(it, type);
+  }
+
+  /// Add one argument to the argument list for each type specified in the list.
+  iterator_range<args_iterator> addArguments(TypeRange types);
+
+  /// Add one value to the argument list at the specified position.
+  BlockArgument insertArgument(unsigned index, Type type) {
+    return front().insertArgument(index, type);
+  }
+
+  /// Erase the argument at 'index' and remove it from the argument list.
+  void eraseArgument(unsigned index) { front().eraseArgument(index); }
+
+  unsigned getNumArguments() { return getArguments().size(); }
+  BlockArgument getArgument(unsigned i) { return getArguments()[i]; }
+
+  //===--------------------------------------------------------------------===//
   // Operation list utilities
   //===--------------------------------------------------------------------===//
 
   /// This class provides iteration over the held operations of blocks directly
   /// within a region.
   class OpIterator final
       : public llvm::iterator_facade_base<OpIterator, std::forward_iterator_tag,
                                           Operation> {

mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp

for (auto config : llvm::zip(mapping, parallelOp.getInductionVars(),
     if (!annotation)
       return parallelOp.emitOpError()
              << "expected mapping attribute for lowering to GPU";
     Value newIndex;
     gpu::Processor processor = gpu::getProcessor(annotation);
 
     if (isMappedToProcessor(processor)) {
       // Use the corresponding thread/grid index as replacement for the loop iv.
-      Value operand = launchOp.body().front().getArgument(
-          getLaunchOpArgumentNum(processor));
+      Value operand =
+          launchOp.body().getArgument(getLaunchOpArgumentNum(processor));
       // Take the indexmap and add the lower bound and step computations in.
       // This computes operand * step + lowerBound.
       // Use an affine map here so that it composes nicely with the provided
       // annotation.
       AffineMap lowerAndStep = AffineMap::get(
           1, 2,
           rewriter.getAffineDimExpr(0) * rewriter.getAffineSymbolExpr(0) +
               rewriter.getAffineSymbolExpr(1));

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

template <typename T> static LogicalResult verifyIndexOp(T op) {
   return success();
 }
 
 static LogicalResult verifyAllReduce(gpu::AllReduceOp allReduce) {
   if (allReduce.body().empty() != allReduce.op().hasValue())
     return allReduce.emitError(
         "expected either an op attribute or a non-empty body");
   if (!allReduce.body().empty()) {
-    if (allReduce.body().front().getNumArguments() != 2)
+    if (allReduce.body().getNumArguments() != 2)
       return allReduce.emitError("expected two region arguments");
-    for (auto argument : allReduce.body().front().getArguments()) {
+    for (auto argument : allReduce.body().getArguments()) {
       if (argument.getType() != allReduce.getType())
         return allReduce.emitError("incorrect region argument type");
     }
     unsigned yieldCount = 0;
     for (Block &block : allReduce.body()) {
       if (auto yield = dyn_cast<gpu::YieldOp>(block.getTerminator())) {
         if (yield.getNumOperands() != 1)
           return allReduce.emitError("expected one gpu.yield operand");
void LaunchOp::build(OpBuilder &builder, OperationState &result,
   Block *body = new Block();
   body->addArguments(
       std::vector<Type>(kNumConfigRegionAttributes, builder.getIndexType()));
   kernelRegion->push_back(body);
 }
 
 KernelDim3 LaunchOp::getBlockIds() {
   assert(!body().empty() && "LaunchOp body must not be empty.");
-  auto args = body().front().getArguments();
+  auto args = body().getArguments();
   return KernelDim3{args[0], args[1], args[2]};
 }
 
 KernelDim3 LaunchOp::getThreadIds() {
   assert(!body().empty() && "LaunchOp body must not be empty.");
-  auto args = body().front().getArguments();
+  auto args = body().getArguments();
   return KernelDim3{args[3], args[4], args[5]};
 }
 
 KernelDim3 LaunchOp::getGridSize() {
   assert(!body().empty() && "LaunchOp body must not be empty.");
-  auto args = body().front().getArguments();
+  auto args = body().getArguments();
   return KernelDim3{args[6], args[7], args[8]};
 }
 
 KernelDim3 LaunchOp::getBlockSize() {
   assert(!body().empty() && "LaunchOp body must not be empty.");
-  auto args = body().getBlocks().front().getArguments();
+  auto args = body().getArguments();
   return KernelDim3{args[9], args[10], args[11]};
 }
 
 KernelDim3 LaunchOp::getGridSizeOperandValues() {
   return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
 }
 
 KernelDim3 LaunchOp::getBlockSizeOperandValues() {
   return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
 }
 
 static LogicalResult verify(LaunchOp op) {
   // Kernel launch takes kNumConfigOperands leading operands for grid/block
   // sizes and transforms them into kNumConfigRegionAttributes region arguments
   // for block/thread identifiers and grid/block sizes.
   if (!op.body().empty()) {
-    Block &entryBlock = op.body().front();
-    if (entryBlock.getNumArguments() !=
+    if (op.body().getNumArguments() !=
         LaunchOp::kNumConfigOperands + op.getNumOperands())
       return op.emitOpError("unexpected number of region arguments");
   }
 
   // Block terminators without successors are expected to exit the kernel region
   // and must be `gpu.terminator`.
   for (Block &block : op.body()) {
     if (block.empty())

 //===----------------------------------------------------------------------===//
 
 /// 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);
+  return getBody().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);
+  return getBody().addArgument(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:
   /// accumulated value of the same type.
   AccumulatorFactory getFactory(Region &body) {
     return AccumulatorFactory([&](Value lhs, Value rhs) {
       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);
+      mapping.map(body.getArgument(0), lhs);
+      mapping.map(body.getArgument(1), rhs);
       rewriter.cloneRegionBefore(body, *split->getParent(),
                                  split->getIterator(), mapping);
 
       // Add branch before inserted body, into body.
       block = block->getNextNode();
       create<BranchOp>(block, ValueRange());
 
       // Replace all gpu.yield ops with branch out of body.

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

static LogicalResult verifyRegionAttribute(Location loc, Type valueType,
   return success();
 }
 
 LogicalResult SPIRVDialect::verifyRegionArgAttribute(Operation *op,
                                                      unsigned regionIndex,
                                                      unsigned argIndex,
                                                      NamedAttribute attribute) {
   return verifyRegionAttribute(
-      op->getLoc(),
-      op->getRegion(regionIndex).front().getArgument(argIndex).getType(),
+      op->getLoc(), op->getRegion(regionIndex).getArgument(argIndex).getType(),
       attribute);
 }
 
 LogicalResult SPIRVDialect::verifyRegionResultAttribute(
     Operation *op, unsigned /*regionIndex*/, unsigned /*resultIndex*/,
     NamedAttribute attribute) {
   return op->emitError("cannot attach SPIR-V attributes to region result");
 }

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

void GenericAtomicRMWOp::build(OpBuilder &builder, OperationState &result,
   result.addOperands(ivs);
 
   if (auto memrefType = memref.getType().dyn_cast<MemRefType>()) {
     Type elementType = memrefType.getElementType();
     result.addTypes(elementType);
 
     Region *bodyRegion = result.addRegion();
     bodyRegion->push_back(new Block());
-    bodyRegion->front().addArgument(elementType);
+    bodyRegion->addArgument(elementType);
   }
 }
 
 static LogicalResult verify(GenericAtomicRMWOp op) {
-  auto &block = op.body().front();
-  if (block.getNumArguments() != 1)
+  auto &body = op.body();
+  if (body.getNumArguments() != 1)
     return op.emitOpError("expected single number of entry block arguments");
 
-  if (op.getResult().getType() != block.getArgument(0).getType())
+  if (op.getResult().getType() != body.getArgument(0).getType())
     return op.emitOpError(
         "expected block argument of the same type result type");
 
   bool hasSideEffects =
-      op.body()
-          .walk([&](Operation *nestedOp) {
+      body.walk([&](Operation *nestedOp) {
             if (MemoryEffectOpInterface::hasNoEffect(nestedOp))
               return WalkResult::advance();
             nestedOp->emitError("body of 'generic_atomic_rmw' should contain "
                                 "only operations with no side effects");
             return WalkResult::interrupt();
           })
           .wasInterrupted();
   return hasSideEffects ? failure() : success();

mlir/lib/IR/AsmPrinter.cpp

 
 unsigned SSANameState::getBlockID(Block *block) {
   auto it = blockIDs.find(block);
   return it != blockIDs.end() ? it->second : NameSentinel;
 }
 
 void SSANameState::shadowRegionArgs(Region &region, ValueRange namesToUse) {
   assert(!region.empty() && "cannot shadow arguments of an empty region");
-  assert(region.front().getNumArguments() == namesToUse.size() &&
+  assert(region.getNumArguments() == namesToUse.size() &&
          "incorrect number of names passed in");
   assert(region.getParentOp()->isKnownIsolatedFromAbove() &&
          "only KnownIsolatedFromAbove ops can shadow names");
 
   SmallVector<char, 16> nameStr;
   for (unsigned i = 0, e = namesToUse.size(); i != e; ++i) {
     auto nameToUse = namesToUse[i];
     if (nameToUse == nullptr)
       continue;
-    auto nameToReplace = region.front().getArgument(i);
+    auto nameToReplace = region.getArgument(i);
 
     nameStr.clear();
     llvm::raw_svector_ostream nameStream(nameStr);
     printValueID(nameToUse, /*printResultNo=*/true, nameStream);
 
     // Entry block arguments should already have a pretty "arg" name.
     assert(valueIDs[nameToReplace] == NameSentinel);
 

mlir/lib/IR/FunctionImplementation.cpp

void mlir::impl::printFunctionSignature(OpAsmPrinter &p, Operation *op,
   bool isExternal = body.empty();
 
   p << '(';
   for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {
     if (i > 0)
       p << ", ";
 
     if (!isExternal) {
-      p.printOperand(body.front().getArgument(i));
+      p.printOperand(body.getArgument(i));
       p << ": ";
     }
 
     p.printType(argTypes[i]);
     p.printOptionalAttrDict(::mlir::impl::getArgAttrs(op, i));
   }
 
   if (isVariadic) {

mlir/lib/IR/Operation.cpp

LogicalResult OpTrait::impl::verifyResultSizeAttr(Operation *op,
   return verifyValueSizeAttr(op, attrName, /*isOperand=*/false);
 }
 
 LogicalResult OpTrait::impl::verifyNoRegionArguments(Operation *op) {
   for (Region &region : op->getRegions()) {
     if (region.empty())
       continue;
 
-    if (region.front().getNumArguments() != 0) {
+    if (region.getNumArguments() != 0) {
       if (op->getNumRegions() > 1)
         return op->emitOpError("region #")
                << region.getRegionNumber() << " should have no arguments";
       else
         return op->emitOpError("region should have no arguments");
     }
   }
   return success();

mlir/lib/IR/Region.cpp

 
 /// Return a location for this region. This is the location attached to the
 /// parent container. The region must have a valid parent container.
 Location Region::getLoc() {
   assert(container && "region is not attached to a container");
   return container->getLoc();
 }
 
+/// Add one argument to the argument list for each type specified in the list.
+iterator_range<Region::args_iterator> Region::addArguments(TypeRange types) {
+  return front().addArguments(types);
+}
+
 Region *Region::getParentRegion() {
   assert(container && "region is not attached to a container");
   return container->getParentRegion();
 }
 
 Operation *Region::getParentOp() { return container; }
 
 bool Region::isProperAncestor(Region *other) {

mlir/lib/Transforms/SCCP.cpp

 
 /// This class contains various state used when computing the lattice of a
 /// callable operation.
 class CallableLatticeState {
 public:
   /// Build a lattice state with a given callable region, and a specified number
   /// of results to be initialized to the default lattice value (Unknown).
   CallableLatticeState(Region *callableRegion, unsigned numResults)
-      : callableArguments(callableRegion->front().getArguments()),
+      : callableArguments(callableRegion->getArguments()),
         resultLatticeValues(numResults) {}
 
   /// Returns the arguments to the callable region.
   Block::BlockArgListType getCallableArguments() const {
     return callableArguments;
   }
 
   /// Returns the lattice value for the results of the callable region.
for (auto callable : symbolTableBlock->getOps<CallableOpInterface>()) {
       if (!symbol)
         continue;
       callableLatticeState.try_emplace(callable, callableRegion,
                                        callable.getCallableResults().size());
 
       // If not all of the uses of this symbol are visible, we can't track the
       // state of the arguments.
       if (symbol.isPublic() || (!allUsesVisible && symbol.isNested()))
-        markAllOverdefined(callableRegion->front().getArguments());
+        markAllOverdefined(callableRegion->getArguments());
     }
     if (callableLatticeState.empty())
       return;
 
     // After computing the valid callables, walk any symbol uses to check
     // for non-call references. We won't be able to track the lattice state
     // for arguments to these callables, as we can't guarantee that we can see
     // all of its calls.

mlir/test/lib/Dialect/Test/TestPatterns.cpp

struct TestUndoBlockArgReplace : public ConversionPattern {
   TestUndoBlockArgReplace(MLIRContext *ctx)
       : ConversionPattern("test.undo_block_arg_replace", /*benefit=*/1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto illegalOp =
         rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type());
-    rewriter.replaceUsesOfBlockArgument(op->getRegion(0).front().getArgument(0),
+    rewriter.replaceUsesOfBlockArgument(op->getRegion(0).getArgument(0),
                                         illegalOp);
     rewriter.updateRootInPlace(op, [] {});
     return success();
   }
 };
 
 /// A rewrite pattern that tests the undo mechanism when erasing a block.
 struct TestUndoBlockErase : public ConversionPattern {