[Mlir-commits] [mlir] 5b569ed - [mlir] Add `Block::eraseArguments` that erases a subrange

[Jeff Niu]

Author: Jeff Niu
Date: 2022-08-29T15:34:21-07:00
New Revision: 5b569ed2cdfb42fcca90f637fc586d7c36834d43

URL: https://github.com/llvm/llvm-project/commit/5b569ed2cdfb42fcca90f637fc586d7c36834d43
DIFF: https://github.com/llvm/llvm-project/commit/5b569ed2cdfb42fcca90f637fc586d7c36834d43.diff

LOG: [mlir] Add `Block::eraseArguments` that erases a subrange

This patch adds a an `eraseArguments` function that erases a subrange of
a block's arguments. This can be used inplace of the terrible pattern

```
block->eraseArguments(llvm::to_vector(llvm::seq(...)));
```

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D132890

Added: 
    

Modified: 
    mlir/include/mlir/IR/Block.h
    mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
    mlir/lib/Dialect/Affine/Utils/Utils.cpp
    mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
    mlir/lib/IR/Block.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index ee4ebcfe19730..1e4e51cc8190f 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -105,6 +105,8 @@ class Block : public IRObjectWithUseList<BlockOperand>,
 
   /// Erase the argument at 'index' and remove it from the argument list.
   void eraseArgument(unsigned index);
+  /// Erases 'num' arguments from the index 'start'.
+  void eraseArguments(unsigned start, unsigned num);
   /// Erases the arguments listed in `argIndices` and removes them from the
   /// argument list.
   /// `argIndices` is allowed to have duplicates and can be in any order.

diff  --git a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
index 46e751e8f0ff1..3b261ffbb2de9 100644
--- a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
+++ b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
@@ -35,8 +35,7 @@ static void inlineIfCase(Region &srcRegion, Region &dstRegion,
   rewriter.create<scf::YieldOp>(yield.getLoc(), yield.getInputs());
   rewriter.eraseOp(yield);
 
-  headBlock->eraseArguments(
-      llvm::to_vector<4>(llvm::seq<unsigned>(0, headBlock->getNumArguments())));
+  headBlock->eraseArguments(0, headBlock->getNumArguments());
 }
 
 static void inlineWhileCase(Region &srcRegion, Region &dstRegion,

diff  --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 66a0e3640aba6..703ff5b5d744b 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -398,8 +398,7 @@ mlir::affineParallelize(AffineForOp forOp,
   // "main" induction variable whenc coming from a non-parallel for.
   unsigned numIVs = 1;
   yieldOp->setOperands(reducedValues);
-  newPloop.getBody()->eraseArguments(
-      llvm::to_vector<4>(llvm::seq<unsigned>(numIVs, numReductions + numIVs)));
+  newPloop.getBody()->eraseArguments(numIVs, numReductions);
 
   forOp.erase();
   return success();

diff  --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
index 355049de9d02b..942430588b771 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp
@@ -162,8 +162,7 @@ mlir::scf::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes,
       thenBlock.getArgument(ivs.index())
           .replaceAllUsesExcept(newIndex, newIndex);
     }
-    thenBlock.eraseArguments(llvm::to_vector<4>(
-        llvm::seq((unsigned)0, thenBlock.getNumArguments())));
+    thenBlock.eraseArguments(0, thenBlock.getNumArguments());
   } else {
     innerLoop.getRegion().takeBody(op.getRegion());
     b.setInsertionPointToStart(innerLoop.getBody());

diff  --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 18a79d25ff5bf..cc84b0d4884f8 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -186,6 +186,15 @@ void Block::eraseArgument(unsigned index) {
     arg.setArgNumber(index++);
 }
 
+void Block::eraseArguments(unsigned start, unsigned num) {
+  assert(start + num <= arguments.size());
+  for (unsigned i = 0; i < num; ++i)
+    arguments[start + i].destroy();
+  arguments.erase(arguments.begin() + start, arguments.begin() + start + num);
+  for (BlockArgument arg : llvm::drop_begin(arguments, start))
+    arg.setArgNumber(start++);
+}
+
 void Block::eraseArguments(ArrayRef<unsigned> argIndices) {
   BitVector eraseIndices(getNumArguments());
   for (unsigned i : argIndices)