[Mlir-commits] [mlir] [MLIR] Preserve operand segment sizes when removing dead inputs (PR #223154)

[Ivan Butygin]

================
@@ -244,6 +244,35 @@ void RewriterBase::eraseBlock(Block *block) {
   block->erase();
 }
 
+void RewriterBase::eraseOperands(Operation *op, const BitVector &eraseIndices) {
+  assert(eraseIndices.size() == op->getNumOperands() &&
+         "expected one bit per operand");
+  if (eraseIndices.none())
+    return;
+
+  modifyOpInPlace(op, [&]() {
+    op->eraseOperands(eraseIndices);
+    if (!op->hasTrait<OpTrait::AttrSizedOperandSegments>())
+      return;
+
+    auto attrName = StringAttr::get(
+        op->getContext(),
+        OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr());
+    auto sizes = op->getAttrOfType<DenseI32ArrayAttr>(attrName);
+    assert(sizes && "expected operand segment sizes attribute");
+    SmallVector<int32_t> newSizes(sizes.asArrayRef());
+    unsigned offset = 0;
+    for (int32_t &size : newSizes) {
+      // The attribute and bit vector still use the original operand indices.
+      unsigned end = offset + size;
+      for (; offset < end; ++offset)
+        size -= eraseIndices.test(offset);
+    }
+    op->setInherentAttr(attrName,
+                        DenseI32ArrayAttr::get(op->getContext(), newSizes));
----------------
Hardcode84 wrote:

I can take a look into it later.

https://github.com/llvm/llvm-project/pull/223154