Fix Block::eraseArgument when block arg is also a successor operand.

[silvas]

Summary

This could trigger an assertion due to the block argument being used by
this block's own successor operands.

[mehdi_amini]

Is this something we could have a test for?

[silvas]

I couldn't think of a great way to test this. The one place I can think of that could potentially benefit from this uses updatePredTerms=false because it erases successor operands manually for other reasons:
https://github.com/llvm/llvm-project/blob/e8358455a2b662bec59bc3971c18346800d7cb00/mlir/lib/Transforms/Utils/RegionUtils.cpp#L278

Maybe the current behavior is a "feature" and not a bug? That is, this function has a contract (currently implicit) saying that the BlockArgument being erased must have no uses? Thinking about it more, doing dropAllUses() is a more general solution and changes the contract of eraseArgument.

So I think it boils down to the contract of eraseArgument. Do we want:

1. Precondition: the BlockArgument has no uses (current behavior at HEAD)
2. Precondition: the BlockArgument has no uses besides successor operands of a predecessor (behavior in this patch)
3. Precondition: the BlockArgument can have uses, and we will dropAllUses() as part of this call.

When put that way, I think this current patch is a weird middle-ground.

Would like to hear Mehdi and River's feedback on this question. I'll also update the eraseArgument comment with the result of the decision :)

[rriddle]

In D74583#1876727, @silvas wrote:

I couldn't think of a great way to test this. The one place I can think of that could potentially benefit from this uses updatePredTerms=false because it erases successor operands manually for other reasons:
https://github.com/llvm/llvm-project/blob/e8358455a2b662bec59bc3971c18346800d7cb00/mlir/lib/Transforms/Utils/RegionUtils.cpp#L278

Maybe the current behavior is a "feature" and not a bug? That is, this function has a contract (currently implicit) saying that the BlockArgument being erased must have no uses? Thinking about it more, doing dropAllUses() is a more general solution and changes the contract of eraseArgument.

So I think it boils down to the contract of eraseArgument. Do we want:

1. Precondition: the BlockArgument has no uses (current behavior at HEAD)
2. Precondition: the BlockArgument has no uses besides successor operands of a predecessor (behavior in this patch)
3. Precondition: the BlockArgument can have uses, and we will dropAllUses() as part of this call.

When put that way, I think this current patch is a weird middle-ground.

Would like to hear Mehdi and River's feedback on this question. I'll also update the eraseArgument comment with the result of the decision :)

I think we can do either 1 or 2. When I originally added the ability to update the predecessor terminators, it was to remove a bit of redundant work that the user would generally end up doing anyways. I'm generally of the opinion that we can treat uses related to block arguments specially, because we can completely reason about their meaning/context/semantics. This is also why 3 feels a bit naughty to me, it feels like it's too easy for the user to hurt themselves that way. One could say that the same logic applies to case 2, but I see it slightly differently as we can do that transformation automatically and know it is correct(i.e. the argument is just a pass through).

mlir/lib/IR/Block.cpp

BlockArgument Block::insertArgument(unsigned index, Type type) {
   assert(index <= arguments.size());
   arguments.insert(arguments.begin() + index, arg);
   return arg;
 }
 
 void Block::eraseArgument(unsigned index, bool updatePredTerms) {
   assert(index < arguments.size());
 
-  // Delete the argument.
-  arguments[index].destroy();
-  arguments.erase(arguments.begin() + index);
-
-  // If we aren't updating predecessors, there is nothing left to do.
-  if (!updatePredTerms)
-    return;
-
+  // If requested, update predecessors. We do this first since this block might
+  // be a predecessor of itself and use this block argument as a successor
+  // operand.
+  if (updatePredTerms) {
     // Erase this argument from each of the predecessor's terminator.
     for (auto predIt = pred_begin(), predE = pred_end(); predIt != predE;
          ++predIt) {
       auto *predTerminator = (*predIt)->getTerminator();
       predTerminator->eraseSuccessorOperand(predIt.getSuccessorIndex(), index);
     }
   }
 
+  // Delete the argument.
+  arguments[index].destroy();
+  arguments.erase(arguments.begin() + index);
+}
+
 /// 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.