[MLIR] Change Operation::create() methods to use Value/Type/Block ranges.

[jurahul]

Summary

• Change Operation::create() methods to use TypeRange for result types, ValueRange for operands and SuccessorRange for successors.
• Extend SuccessorRange to be able to represent a range of blocks from BlockOperands or from an array of Block pointers.

[rriddle]

Thanks!

mlir/include/mlir/IR/BlockSupport.h

58

I would not do this for SuccessorRange. It is intended to have the same limited scope that OperandRange and ResultRange have. If we want something general I would suggest adding a BlockRange instead.

[jurahul]

mlir/include/mlir/IR/BlockSupport.h

58

Thanks. I looked into how ResultRange, OperandRange and ValueRange are layered. So I guess your suggestion is to keep SuccessorRange as is, and add a new BlockRange which can work with either a SuccessorRange or an ArrayRef<Block *>, right? And then the Operation::create() functions would use a BlockRange argument for successors.

[rriddle]

mlir/include/mlir/IR/BlockSupport.h

58

Yes, exactly.

[jurahul]

Review feedback

[bondhugula]

mlir/include/mlir/IR/BlockSupport.h

82

Missing doc comment.

[jurahul]

Add comment for BlockRange

[rriddle]

Thanks!

mlir/include/mlir/IR/BlockSupport.h

 
 //===----------------------------------------------------------------------===//
 // Successors
 //===----------------------------------------------------------------------===//
 
 /// This class implements the successor iterators for Block.
 class SuccessorRange final
     : public llvm::detail::indexed_accessor_range_base<
           SuccessorRange, BlockOperand *, Block *, Block *, Block *> {
 public:
   using RangeBaseT::RangeBaseT;
   SuccessorRange(Block *block);
   SuccessorRange(Operation *term);
 
 private:
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static BlockOperand *offset_base(BlockOperand *object, ptrdiff_t index) {
     return object + index;
   }
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static Block *dereference_iterator(BlockOperand *object, ptrdiff_t index) {
     return object[index].get();
   }
 
   /// Allow access to `offset_base` and `dereference_iterator`.
   friend RangeBaseT;
 };
 
 //===----------------------------------------------------------------------===//
+// BlockRange
+//===----------------------------------------------------------------------===//
+
+/// This class provides an abstraction over the different types of ranges over
+/// Blocks. In many cases, this prevents the need to explicitly materialize a
+/// SmallVector/std::vector. This class should be used in places that are not
+/// suitable for a more derived type (e.g. ArrayRef) or a template range
+/// parameter.
+class BlockRange final
+    : public llvm::detail::indexed_accessor_range_base<
+          BlockRange, llvm::PointerUnion<BlockOperand *, Block *const *>,
+          Block *, Block *, Block *> {
+public:
+  using RangeBaseT::RangeBaseT;
+  BlockRange(ArrayRef<Block *> blocks = llvm::None);
+  BlockRange(SuccessorRange successors);
+  template <typename Arg,
+            typename = typename std::enable_if_t<
+                std::is_constructible<ArrayRef<Block *>, Arg>::value>>
+  BlockRange(Arg &&arg)
+      : BlockRange(ArrayRef<Block *>(std::forward<Arg>(arg))) {}
+  BlockRange(std::initializer_list<Block *> blocks)
+      : BlockRange(ArrayRef<Block *>(blocks)) {}
+
+private:
+  /// The owner of the range is either:
+  /// * A pointer to the first element of an array of block operands.
+  /// * A pointer to the first element of an array of Block *.
+  using OwnerT = llvm::PointerUnion<BlockOperand *, Block *const *>;
+
+  /// See `llvm::detail::indexed_accessor_range_base` for details.
+  static OwnerT offset_base(OwnerT object, ptrdiff_t index);
+
+  /// See `llvm::detail::indexed_accessor_range_base` for details.
+  static Block *dereference_iterator(OwnerT object, ptrdiff_t index);
+
+  /// Allow access to `offset_base` and `dereference_iterator`.
+  friend RangeBaseT;
+};
+
+//===----------------------------------------------------------------------===//
 // Operation Iterators
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// A utility iterator that filters out operations that are not 'OpT'.
 template <typename OpT, typename IteratorT>
 class op_filter_iterator
     : public llvm::filter_iterator<IteratorT, bool (*)(Operation &)> {

mlir/include/mlir/IR/Operation.h

 class Operation final
     : public llvm::ilist_node_with_parent<Operation, Block>,
       private llvm::TrailingObjects<Operation, detail::InLineOpResult,
                                     detail::TrailingOpResult, BlockOperand,
                                     Region, detail::OperandStorage> {
 public:
   /// Create a new Operation with the specific fields.
   static Operation *create(Location location, OperationName name,
-                           ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
+                           TypeRange resultTypes, ValueRange operands,
                            ArrayRef<NamedAttribute> attributes,
-                           ArrayRef<Block *> successors, unsigned numRegions);
+                           BlockRange successors, unsigned numRegions);
 
   /// Overload of create that takes an existing MutableDictionaryAttr to avoid
   /// unnecessarily uniquing a list of attributes.
   static Operation *create(Location location, OperationName name,
-                           ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
+                           TypeRange resultTypes, ValueRange operands,
                            MutableDictionaryAttr attributes,
-                           ArrayRef<Block *> successors, unsigned numRegions);
+                           BlockRange successors, unsigned numRegions);
 
   /// Create a new Operation from the fields stored in `state`.
   static Operation *create(const OperationState &state);
 
   /// Create a new Operation with the specific fields.
   static Operation *create(Location location, OperationName name,
-                           ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
+                           TypeRange resultTypes, ValueRange operands,
                            MutableDictionaryAttr attributes,
-                           ArrayRef<Block *> successors = {},
+                           BlockRange successors = {},
                            RegionRange regions = {});
 
   /// The name of an operation is the key identifier for it.
   OperationName getName() { return name; }
 
   /// If this operation has a registered operation description, return it.
   /// Otherwise return null.
   const AbstractOperation *getAbstractOperation() {
private:
   /// Update the order index of this operation of this operation if necessary,
   /// potentially recomputing the order of the parent block.
   void updateOrderIfNecessary();
 
   /// Returns true if this operation has a valid order.
   bool hasValidOrder() { return orderIndex != kInvalidOrderIdx; }
 
 private:
-  Operation(Location location, OperationName name, ArrayRef<Type> resultTypes,
+  Operation(Location location, OperationName name, TypeRange resultTypes,
             unsigned numSuccessors, unsigned numRegions,
             const MutableDictionaryAttr &attributes, bool hasOperandStorage);
 
   // Operations are deleted through the destroy() member because they are
   // allocated with malloc.
   ~Operation();
 
   /// Returns the operand storage object.

mlir/include/mlir/IR/OperationSupport.h

 #include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/TrailingObjects.h"
 #include <memory>
 
 namespace mlir {
 class Block;
+class BlockRange;
 class Dialect;
 class Operation;
 struct OperationState;
 class OpAsmParser;
 class OpAsmParserResult;
 class OpAsmPrinter;
 class OperandRange;
 class OpFoldResult;
 class ParseResult;
 class Pattern;
 class Region;
 class ResultRange;
 class RewritePattern;
-class SuccessorRange;
 class Type;
 class Value;
 class ValueRange;
 template <typename ValueRangeT> class ValueTypeRange;
 
 class OwningRewritePatternList;
 
 //===----------------------------------------------------------------------===//
void addAttribute(Identifier name, Attribute attr) {
     attributes.append(name, attr);
   }
 
   /// Add an array of named attributes.
   void addAttributes(ArrayRef<NamedAttribute> newAttributes) {
     attributes.append(newAttributes);
   }
 
-  /// Add an array of successors.
-  void addSuccessors(ArrayRef<Block *> newSuccessors) {
-    successors.append(newSuccessors.begin(), newSuccessors.end());
-  }
   void addSuccessors(Block *successor) { successors.push_back(successor); }
-  void addSuccessors(SuccessorRange newSuccessors);
+  void addSuccessors(BlockRange newSuccessors);
 
   /// Create a region that should be attached to the operation.  These regions
   /// can be filled in immediately without waiting for Operation to be
   /// created.  When it is, the region bodies will be transferred.
   Region *addRegion();
 
   /// Take a region that should be attached to the Operation.  The body of the
   /// region will be transferred when the Operation is constructed.  If the

mlir/lib/IR/Block.cpp

 }
 
 /// Get the successor number in the predecessor terminator.
 unsigned PredecessorIterator::getSuccessorIndex() const {
   return I->getOperandNumber();
 }
 
 //===----------------------------------------------------------------------===//
-// Successors
+// SuccessorRange
 //===----------------------------------------------------------------------===//
 
 SuccessorRange::SuccessorRange(Block *block) : SuccessorRange(nullptr, 0) {
   if (Operation *term = block->getTerminator())
     if ((count = term->getNumSuccessors()))
       base = term->getBlockOperands().data();
 }
 
 SuccessorRange::SuccessorRange(Operation *term) : SuccessorRange(nullptr, 0) {
   if ((count = term->getNumSuccessors()))
     base = term->getBlockOperands().data();
 }
+
+//===----------------------------------------------------------------------===//
+// BlockRange
+//===----------------------------------------------------------------------===//
+
+BlockRange::BlockRange(ArrayRef<Block *> blocks) : BlockRange(nullptr, 0) {
+  if ((count = blocks.size()))
+    base = blocks.data();
+}
+
+BlockRange::BlockRange(SuccessorRange successors)
+    : BlockRange(successors.begin().getBase(), successors.size()) {}
+
+/// See `llvm::detail::indexed_accessor_range_base` for details.
+BlockRange::OwnerT BlockRange::offset_base(OwnerT object, ptrdiff_t index) {
+  if (auto *operand = object.dyn_cast<BlockOperand *>())
+    return {operand + index};
+  return {object.dyn_cast<Block *const *>() + index};
+}
+
+/// See `llvm::detail::indexed_accessor_range_base` for details.
+Block *BlockRange::dereference_iterator(OwnerT object, ptrdiff_t index) {
+  if (const auto *operand = object.dyn_cast<BlockOperand *>())
+    return operand[index].get();
+  return object.dyn_cast<Block *const *>()[index];
+}

mlir/lib/IR/Operation.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // Operation
 //===----------------------------------------------------------------------===//
 
 /// Create a new Operation with the specific fields.
 Operation *Operation::create(Location location, OperationName name,
-                             ArrayRef<Type> resultTypes,
-                             ArrayRef<Value> operands,
+                             TypeRange resultTypes, ValueRange operands,
                              ArrayRef<NamedAttribute> attributes,
-                             ArrayRef<Block *> successors,
-                             unsigned numRegions) {
+                             BlockRange successors, unsigned numRegions) {
   return create(location, name, resultTypes, operands,
                 MutableDictionaryAttr(attributes), successors, numRegions);
 }
 
 /// Create a new Operation from operation state.
 Operation *Operation::create(const OperationState &state) {
-  return Operation::create(state.location, state.name, state.types,
-                           state.operands, state.attributes, state.successors,
-                           state.regions);
+  return create(state.location, state.name, state.types, state.operands,
+                state.attributes, state.successors, state.regions);
 }
 
 /// Create a new Operation with the specific fields.
 Operation *Operation::create(Location location, OperationName name,
-                             ArrayRef<Type> resultTypes,
-                             ArrayRef<Value> operands,
+                             TypeRange resultTypes, ValueRange operands,
                              MutableDictionaryAttr attributes,
-                             ArrayRef<Block *> successors,
-                             RegionRange regions) {
+                             BlockRange successors, RegionRange regions) {
   unsigned numRegions = regions.size();
   Operation *op = create(location, name, resultTypes, operands, attributes,
                          successors, numRegions);
   for (unsigned i = 0; i < numRegions; ++i)
     if (regions[i])
       op->getRegion(i).takeBody(*regions[i]);
   return op;
 }
 
 /// Overload of create that takes an existing MutableDictionaryAttr to avoid
 /// unnecessarily uniquing a list of attributes.
 Operation *Operation::create(Location location, OperationName name,
-                             ArrayRef<Type> resultTypes,
-                             ArrayRef<Value> operands,
+                             TypeRange resultTypes, ValueRange operands,
                              MutableDictionaryAttr attributes,
-                             ArrayRef<Block *> successors,
-                             unsigned numRegions) {
+                             BlockRange successors, unsigned numRegions) {
   // We only need to allocate additional memory for a subset of results.
   unsigned numTrailingResults = OpResult::getNumTrailing(resultTypes.size());
   unsigned numInlineResults = OpResult::getNumInline(resultTypes.size());
   unsigned numSuccessors = successors.size();
   unsigned numOperands = operands.size();
 
   // If the operation is known to have no operands, don't allocate an operand
   // storage.
Operation *Operation::create(Location location, OperationName name,
   auto blockOperands = op->getBlockOperands();
   for (unsigned i = 0; i != numSuccessors; ++i)
     new (&blockOperands[i]) BlockOperand(op, successors[i]);
 
   return op;
 }
 
 Operation::Operation(Location location, OperationName name,
-                     ArrayRef<Type> resultTypes, unsigned numSuccessors,
+                     TypeRange resultTypes, unsigned numSuccessors,
                      unsigned numRegions,
                      const MutableDictionaryAttr &attributes,
                      bool hasOperandStorage)
     : location(location), numSuccs(numSuccessors), numRegions(numRegions),
       hasOperandStorage(hasOperandStorage), hasSingleResult(false), name(name),
       attrs(attributes) {
   if (!resultTypes.empty()) {
     // If there is a single result it is stored in-place, otherwise use a tuple.
for (auto opValue : getOperands())
     operands.push_back(mapper.lookupOrDefault(opValue));
 
   // Remap the successors.
   successors.reserve(getNumSuccessors());
   for (Block *successor : getSuccessors())
     successors.push_back(mapper.lookupOrDefault(successor));
 
   // Create the new operation.
-  auto *newOp = Operation::create(getLoc(), getName(), getResultTypes(),
-                                  operands, attrs, successors, getNumRegions());
+  auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs,
+                       successors, getNumRegions());
 
   // Remember the mapping of any results.
   for (unsigned i = 0, e = getNumResults(); i != e; ++i)
     mapper.map(getResult(i), newOp->getResult(i));
 
   return newOp;
 }
 

mlir/lib/IR/OperationSupport.cpp

OperationState::OperationState(Location location, StringRef name,
   for (std::unique_ptr<Region> &r : regions)
     this->regions.push_back(std::move(r));
 }
 
 void OperationState::addOperands(ValueRange newOperands) {
   operands.append(newOperands.begin(), newOperands.end());
 }
 
-void OperationState::addSuccessors(SuccessorRange newSuccessors) {
+void OperationState::addSuccessors(BlockRange newSuccessors) {
   successors.append(newSuccessors.begin(), newSuccessors.end());
 }
 
 Region *OperationState::addRegion() {
   regions.emplace_back(new Region);
   return regions.back().get();
 }