From 5b569ed2cdfb42fcca90f637fc586d7c36834d43 Mon Sep 17 00:00:00 2001 From: Jeff Niu Date: Mon, 29 Aug 2022 14:32:14 -0700 Subject: [PATCH] [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 --- mlir/include/mlir/IR/Block.h | 2 ++ mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp | 3 +-- mlir/lib/Dialect/Affine/Utils/Utils.cpp | 3 +-- mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp | 3 +-- mlir/lib/IR/Block.cpp | 9 +++++++++ 5 files changed, 14 insertions(+), 6 deletions(-) diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h index ee4ebcf..1e4e51c 100644 --- a/mlir/include/mlir/IR/Block.h +++ b/mlir/include/mlir/IR/Block.h @@ -105,6 +105,8 @@ public: /// 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 46e751e..3b261ff 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(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 index 66a0e36..703ff5b 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(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 355049d..9424305 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 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 18a79d25..cc84b0d 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 argIndices) { BitVector eraseIndices(getNumArguments()); for (unsigned i : argIndices) -- 2.7.4