diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h --- a/mlir/include/mlir/IR/Block.h +++ b/mlir/include/mlir/IR/Block.h @@ -105,6 +105,8 @@ /// 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 --- a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp +++ b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp @@ -35,8 +35,7 @@ rewriter.create(yield.getLoc(), yield.getInputs()); rewriter.eraseOp(yield); - headBlock->eraseArguments( - llvm::to_vector<4>(llvm::seq(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 --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -398,8 +398,7 @@ // "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(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 --- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp @@ -162,8 +162,7 @@ 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 --- a/mlir/lib/IR/Block.cpp +++ b/mlir/lib/IR/Block.cpp @@ -186,6 +186,15 @@ 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 argIndices) { BitVector eraseIndices(getNumArguments()); for (unsigned i : argIndices)