Store (cache) the Argument number (index in the argument list) inside the BlockArgumentImpl (NFC)

[mehdi_amini]

Summary

This avoids linear search in BlockArgument::getArgNumber().

[rriddle]

mlir/include/mlir/IR/Value.h

287

The comment here is off.

288

Can you make this private? Block is already a friend.

mlir/lib/IR/Block.cpp

161

+ index + 1? The argument at index should already have its index set. Also, llvm::drop_begin instead?

181

llvm::drop_begin ?

194–195

Using eraseArgument here is going to result in updating the indices multiple times. Also can you merge indice updates into this loop?

[mehdi_amini]

mlir/lib/IR/Block.cpp

194–195

Using eraseArgument here is going to result in updating the indices multiple times

Good point!

Also can you merge indice updates into this loop?

I'm not sure how though: the loop is going backward, you can't know on the first element you're deleting what its final index will be?
(at least not easily)

[mehdi_amini]

Address comments

[rriddle]

mlir/lib/IR/Block.cpp

160

nit: ++index; (You could also merge it into the drop_begin call)

194–195

I'm not sure how though: the loop is going backward, you can't know on the first element you're deleting what its final index will be?

(at least not easily)
Ah right, missed that this was in reverse.

202

You could use eraseIndices.find_first() to get the first removed argument.

[mehdi_amini]

mlir/lib/IR/Block.cpp

202

Updated, but I'm not totally sure it's better for the common case: the code in find_first is not trivial compared to the vector traversal here.

[mehdi_amini]

Address comments

mlir/include/mlir/IR/Value.h

 
 //===----------------------------------------------------------------------===//
 // BlockArgument
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// The internal implementation of a BlockArgument.
 class BlockArgumentImpl : public IRObjectWithUseList<OpOperand> {
-  BlockArgumentImpl(Type type, Block *owner) : type(type), owner(owner) {}
+  BlockArgumentImpl(Type type, Block *owner, int64_t index)
+      : type(type), owner(owner), index(index) {}
 
   /// The type of this argument.
   Type type;
 
   /// The owner of this argument.
   Block *owner;
 
+  /// The position in the argument list.
+  int64_t index;
+
   /// Allow access to owner and constructor.
   friend BlockArgument;
 };
 } // end namespace detail
 
 /// Block arguments are values.
 class BlockArgument : public Value {
 public:
   using Value::Value;
 
   static bool classof(Value value) {
     return value.getKind() == Kind::BlockArgument;
   }
 
   /// Returns the block that owns this argument.
   Block *getOwner() const { return getImpl()->owner; }
 
   /// Return the type of this value.
   Type getType() const { return getImpl()->type; }
 
   /// Set the type of this value.
   void setType(Type newType) { getImpl()->type = newType; }
 
   /// Returns the number of this argument.
-  unsigned getArgNumber() const;
+  unsigned getArgNumber() const { return getImpl()->index; }
 
 private:
   /// Allocate a new argument with the given type and owner.
-  static BlockArgument create(Type type, Block *owner) {
-    return new detail::BlockArgumentImpl(type, owner);
+  static BlockArgument create(Type type, Block *owner, int64_t index) {
+    return new detail::BlockArgumentImpl(type, owner, index);
   }
 
   /// Destroy and deallocate this argument.
   void destroy() { delete getImpl(); }
 
   /// Get a raw pointer to the internal implementation.
   detail::BlockArgumentImpl *getImpl() const {
     return reinterpret_cast<detail::BlockArgumentImpl *>(
         ownerAndKind.getPointer());
   }
 
-  /// Allow access to `create` and `destroy`.
+  /// Cache the position in the block argument list.
+  void setArgNumber(int64_t index) { getImpl()->index = index; }
+
+  /// Allow access to `create`, `destroy` and `setArgNumber`.
   friend Block;
 
   /// Allow access to 'getImpl'.
   friend Value;
 };
 
 //===----------------------------------------------------------------------===//
 // OpResult

mlir/lib/IR/Block.cpp

 //===- Block.cpp - MLIR Block Class ---------------------------------------===//
 //
 // 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/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/ADT/BitVector.h"
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
-// BlockArgument
-//===----------------------------------------------------------------------===//
-
-/// Returns the number of this argument.
-unsigned BlockArgument::getArgNumber() const {
-  // Arguments are not stored in place, so we have to find it within the list.
-  auto argList = getOwner()->getArguments();
-  return std::distance(argList.begin(), llvm::find(argList, *this));
-}
-
-//===----------------------------------------------------------------------===//
 // Block
 //===----------------------------------------------------------------------===//
 
 Block::~Block() {
   assert(!verifyOpOrder() && "Expected valid operation ordering.");
   clear();
   for (BlockArgument arg : arguments)
     arg.destroy();

 //===----------------------------------------------------------------------===//
 
 /// Return a range containing the types of the arguments for this block.
 auto Block::getArgumentTypes() -> ValueTypeRange<BlockArgListType> {
   return ValueTypeRange<BlockArgListType>(getArguments());
 }
 
 BlockArgument Block::addArgument(Type type) {
-  BlockArgument arg = BlockArgument::create(type, this);
+  BlockArgument arg = BlockArgument::create(type, this, arguments.size());
   arguments.push_back(arg);
   return arg;
 }
 
 /// Add one argument to the argument list for each type specified in the list.
 auto Block::addArguments(TypeRange types) -> iterator_range<args_iterator> {
   size_t initialSize = arguments.size();
   arguments.reserve(initialSize + types.size());
   for (auto type : types)
     addArgument(type);
   return {arguments.data() + initialSize, arguments.data() + arguments.size()};
 }
 
 BlockArgument Block::insertArgument(unsigned index, Type type) {
-  auto arg = BlockArgument::create(type, this);
+  auto arg = BlockArgument::create(type, this, index);
   assert(index <= arguments.size());
   arguments.insert(arguments.begin() + index, arg);
+  // Update the cached position for all the arguments after the newly inserted
+  // one.
+  ++index;
+  for (BlockArgument arg : llvm::drop_begin(arguments, index))
+    arg.setArgNumber(index++);
   return arg;
 }
 
+/// Insert one value to the given position of the argument list. The existing
+/// arguments are shifted. The block is expected not to have predecessors.
+BlockArgument Block::insertArgument(args_iterator it, Type type) {
+  assert(llvm::empty(getPredecessors()) &&
+         "cannot insert arguments to blocks with predecessors");
+  return insertArgument(it->getArgNumber(), type);
+}
+
 void Block::eraseArgument(unsigned index) {
   assert(index < arguments.size());
   arguments[index].destroy();
   arguments.erase(arguments.begin() + index);
+  for (BlockArgument arg : llvm::drop_begin(arguments, index))
+    arg.setArgNumber(index++);
 }
 
 void Block::eraseArguments(ArrayRef<unsigned> argIndices) {
   llvm::BitVector eraseIndices(getNumArguments());
   for (unsigned i : argIndices)
     eraseIndices.set(i);
   eraseArguments(eraseIndices);
 }
 
 void Block::eraseArguments(llvm::BitVector eraseIndices) {
   // We do this in reverse so that we erase later indices before earlier
   // indices, to avoid shifting the later indices.
   unsigned originalNumArgs = getNumArguments();
-  for (unsigned i = 0; i < originalNumArgs; ++i)
-    if (eraseIndices.test(originalNumArgs - i - 1))
-      eraseArgument(originalNumArgs - i - 1);
+  for (unsigned i = 0; i < originalNumArgs; ++i) {
+    int64_t currentPos = originalNumArgs - i - 1;
+    if (eraseIndices.test(currentPos)) {
+      arguments[currentPos].destroy();
+      arguments.erase(arguments.begin() + currentPos);
     }
-
-/// Insert one value to the given position of the argument list. The existing
-/// arguments are shifted. The block is expected not to have predecessors.
-BlockArgument Block::insertArgument(args_iterator it, Type type) {
-  assert(llvm::empty(getPredecessors()) &&
-         "cannot insert arguments to blocks with predecessors");
-
-  // Use the args_iterator (on the BlockArgListType) to compute the insertion
-  // iterator in the underlying argument storage.
-  size_t distance = std::distance(args_begin(), it);
-  auto arg = BlockArgument::create(type, this);
-  arguments.insert(std::next(arguments.begin(), distance), arg);
-  return arg;
+  }
+  // Update the cached position for the arguments after the first erased one.
+  int64_t index = 0;
+  for (BlockArgument arg :
+       llvm::drop_begin(arguments, eraseIndices.find_first()))
+    arg.setArgNumber(index++);
 }
 
 //===----------------------------------------------------------------------===//
 // Terminator management
 //===----------------------------------------------------------------------===//
 
 /// Get the terminator operation of this block. This function asserts that
 /// the block has a valid terminator operation.