[Mlir-commits] [mlir] [MLIR] Preserve operand segment sizes when removing dead inputs (PR #223154)
Ivan Butygin
llvmlistbot at llvm.org
Sun Sep 13 03:03:51 PDT 2026
https://github.com/Hardcode84 updated https://github.com/llvm/llvm-project/pull/223154
>From 744f04be0a2092dd6b9817eb7aef8586c2b1def3 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 12 Sep 2026 16:11:23 +0200
Subject: [PATCH 1/7] [MLIR] Preserve operand segment sizes when removing dead
region inputs
---
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 32 ++++-
.../region-branch-operand-segments.mlir | 125 ++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 42 ++++++
mlir/test/lib/Dialect/Test/TestOps.td | 26 ++++
4 files changed, 224 insertions(+), 1 deletion(-)
create mode 100644 mlir/test/Transforms/region-branch-operand-segments.mlir
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 475969a7c0d09..c7b90df464a98 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -919,11 +919,41 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
}
}
+ StringAttr attrName;
// Erase operands.
for (auto &pair : operandsToRemove) {
Operation *op = pair.first;
BitVector &operands = pair.second;
- rewriter.modifyOpInPlace(op, [&]() { op->eraseOperands(operands); });
+ rewriter.modifyOpInPlace(op, [&]() {
+ if (!op->hasTrait<OpTrait::AttrSizedOperandSegments>()) {
+ op->eraseOperands(operands);
+ return;
+ }
+ if (!attrName)
+ attrName = rewriter.getStringAttr(OpTrait::AttrSizedOperandSegments<
+ void>::getOperandSegmentSizeAttr());
+ auto sizes = op->getAttrOfType<DenseI32ArrayAttr>(attrName);
+ assert(sizes && "expected operand segment sizes attribute");
+ unsigned offset = op->getNumOperands();
+ // Erase from the last segment first to preserve the original indices.
+ for (unsigned segment :
+ llvm::reverse(llvm::seq<unsigned>(sizes.size()))) {
+ unsigned size = sizes.asArrayRef()[segment];
+ offset -= size;
+ SmallVector<Value> retained;
+ for (unsigned i : llvm::seq(size))
+ if (!operands.test(offset + i))
+ retained.push_back(op->getOperand(offset + i));
+ if (retained.size() == size)
+ continue;
+ // Use getAttr to keep changes to later segments.
+ MutableOperandRange range(
+ op, offset, size,
+ MutableOperandRange::OperandSegment(
+ segment, {attrName, op->getAttr(attrName)}));
+ range.assign(retained);
+ }
+ });
}
// Erase block arguments.
diff --git a/mlir/test/Transforms/region-branch-operand-segments.mlir b/mlir/test/Transforms/region-branch-operand-segments.mlir
new file mode 100644
index 0000000000000..e8a74035665d7
--- /dev/null
+++ b/mlir/test/Transforms/region-branch-operand-segments.mlir
@@ -0,0 +1,125 @@
+// RUN: mlir-opt %s --canonicalize --split-input-file | FileCheck %s
+
+// Entry and backedge values differ, so live block arguments cannot fold.
+// Check operand order, block arguments, and segment sizes on the parent and
+// on both terminators. The control operands must remain in place.
+
+// CHECK-LABEL: func.func @remove_from_both_segments(
+// CHECK-SAME: %[[CONDITION:[^ ,)]+]]: i1, %[[STOP:[^ ,)]+]]: i1, %[[A:[^ ,)]+]]: i32, %[[B:[^ ,)]+]]: i32, %[[C:[^ ,)]+]]: i32, %[[D:[^ ,)]+]]: i32, %[[E:[^ ,)]+]]: i32, %[[F:[^ ,)]+]]: i32) {
+// CHECK: "test.segmented_region_branch"(%[[CONDITION]], %[[A]], %[[C]], %[[E]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 2, 1, 1>}> ({
+// CHECK: ^bb0(%[[R0_0:[^ ,)]+]]: i32, %[[R0_2:[^ ,)]+]]: i32):
+// CHECK: call @use2(%[[R0_0]], %[[R0_2]])
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[D]], %[[F]], %[[B]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 2, 1, 1>}>
+// CHECK: ^bb0(%[[R1_1:[^ ,)]+]]: i32):
+// CHECK: call @use1(%[[R1_1]])
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[D]], %[[F]], %[[B]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 2, 1, 1>}>
+// CHECK: })
+func.func @remove_from_both_segments(%condition: i1, %stop: i1, %a: i32, %b: i32, %c: i32, %d: i32, %e: i32, %f: i32) {
+ "test.segmented_region_branch"(%condition, %a, %b, %c, %d, %e, %f, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> ({
+ ^bb0(%r0_0: i32, %r0_1: i32, %r0_2: i32):
+ func.call @use2(%r0_0, %r0_2) : (i32, i32) -> ()
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ }, {
+ ^bb0(%r1_0: i32, %r1_1: i32, %r1_2: i32):
+ func.call @use1(%r1_1) : (i32) -> ()
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ }) : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ return
+}
+func.func private @use1(i32)
+func.func private @use2(i32, i32)
+
+// -----
+
+// CHECK-LABEL: func.func @remove_entire_segment(
+// CHECK-SAME: %[[CONDITION:[^ ,)]+]]: i1, %[[STOP:[^ ,)]+]]: i1, %[[A:[^ ,)]+]]: i32, %[[B:[^ ,)]+]]: i32, %[[C:[^ ,)]+]]: i32, %[[D:[^ ,)]+]]: i32, %[[E:[^ ,)]+]]: i32, %[[F:[^ ,)]+]]: i32) {
+// CHECK: "test.segmented_region_branch"(%[[CONDITION]], %[[D]], %[[F]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 0, 2, 1>}> ({
+// CHECK-NOT: ^bb
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[A]], %[[C]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 0, 2, 1>}>
+// CHECK: ^bb0(%[[R1_0:[^ ,)]+]]: i32, %[[R1_2:[^ ,)]+]]: i32):
+// CHECK: call @use2(%[[R1_0]], %[[R1_2]])
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[A]], %[[C]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 0, 2, 1>}>
+// CHECK: })
+func.func @remove_entire_segment(%condition: i1, %stop: i1, %a: i32, %b: i32, %c: i32, %d: i32, %e: i32, %f: i32) {
+ "test.segmented_region_branch"(%condition, %a, %b, %c, %d, %e, %f, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> ({
+ ^bb0(%r0_0: i32, %r0_1: i32, %r0_2: i32):
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ }, {
+ ^bb0(%r1_0: i32, %r1_1: i32, %r1_2: i32):
+ func.call @use2(%r1_0, %r1_2) : (i32, i32) -> ()
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ }) : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ return
+}
+func.func private @use2(i32, i32)
+
+// -----
+
+// CHECK-LABEL: func.func @remove_all_successor_operands(
+// CHECK-SAME: %[[CONDITION:[^ ,)]+]]: i1, %[[STOP:[^ ,)]+]]: i1, %[[A:[^ ,)]+]]: i32, %[[B:[^ ,)]+]]: i32, %[[C:[^ ,)]+]]: i32, %[[D:[^ ,)]+]]: i32, %[[E:[^ ,)]+]]: i32, %[[F:[^ ,)]+]]: i32) {
+// CHECK: "test.segmented_region_branch"(%[[CONDITION]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 0, 0, 1>}> ({
+// CHECK-NOT: ^bb
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 0, 0, 1>}>
+// CHECK-NOT: ^bb
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[STOP]]) <{operandSegmentSizes = array<i32: 1, 0, 0, 1>}>
+// CHECK: })
+func.func @remove_all_successor_operands(%condition: i1, %stop: i1, %a: i32, %b: i32, %c: i32, %d: i32, %e: i32, %f: i32) {
+ "test.segmented_region_branch"(%condition, %a, %b, %c, %d, %e, %f, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> ({
+ ^bb0(%r0_0: i32, %r0_1: i32, %r0_2: i32):
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ }, {
+ ^bb0(%r1_0: i32, %r1_1: i32, %r1_2: i32):
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c, %stop) <{operandSegmentSizes = array<i32: 1, 3, 3, 1>}> : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ }) : (i1, i32, i32, i32, i32, i32, i32, i1) -> ()
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @empty_segments_and_absent_control(
+// CHECK-SAME: %[[CONDITION:[^ ,)]+]]: i1, %[[STOP:[^ ,)]+]]: i1, %[[A:[^ ,)]+]]: i32, %[[B:[^ ,)]+]]: i32, %[[C:[^ ,)]+]]: i32, %[[D:[^ ,)]+]]: i32, %[[E:[^ ,)]+]]: i32, %[[F:[^ ,)]+]]: i32) {
+// CHECK: "test.segmented_region_branch"(%[[CONDITION]], %[[E]]) <{operandSegmentSizes = array<i32: 1, 0, 1, 0>}> ({
+// CHECK-NOT: ^bb
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[B]]) <{operandSegmentSizes = array<i32: 1, 0, 1, 0>}>
+// CHECK: ^bb0(%[[R1_1:[^ ,)]+]]: i32):
+// CHECK: call @use1(%[[R1_1]])
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[B]]) <{operandSegmentSizes = array<i32: 1, 0, 1, 0>}>
+// CHECK: })
+func.func @empty_segments_and_absent_control(%condition: i1, %stop: i1, %a: i32, %b: i32, %c: i32, %d: i32, %e: i32, %f: i32) {
+ "test.segmented_region_branch"(%condition, %d, %e, %f) <{operandSegmentSizes = array<i32: 1, 0, 3, 0>}> ({
+ ^bb0:
+ "test.segmented_region_branch_terminator"(%condition, %a, %b, %c) <{operandSegmentSizes = array<i32: 1, 0, 3, 0>}> : (i1, i32, i32, i32) -> ()
+ }, {
+ ^bb0(%r1_0: i32, %r1_1: i32, %r1_2: i32):
+ func.call @use1(%r1_1) : (i32) -> ()
+ "test.segmented_region_branch_terminator"(%condition, %a, %b, %c) <{operandSegmentSizes = array<i32: 1, 0, 3, 0>}> : (i1, i32, i32, i32) -> ()
+ }) : (i1, i32, i32, i32) -> ()
+ return
+}
+func.func private @use1(i32)
+
+// -----
+
+// CHECK-LABEL: func.func @retain_all_operands(
+// CHECK-SAME: %[[CONDITION:[^ ,)]+]]: i1, %[[STOP:[^ ,)]+]]: i1, %[[A:[^ ,)]+]]: i32, %[[B:[^ ,)]+]]: i32, %[[C:[^ ,)]+]]: i32, %[[D:[^ ,)]+]]: i32, %[[E:[^ ,)]+]]: i32, %[[F:[^ ,)]+]]: i32) {
+// CHECK: "test.segmented_region_branch"(%[[CONDITION]], %[[A]], %[[B]], %[[C]], %[[D]], %[[E]], %[[F]]) <{operandSegmentSizes = array<i32: 1, 3, 3, 0>}> ({
+// CHECK: ^bb0(%[[R0_0:[^ ,)]+]]: i32, %[[R0_1:[^ ,)]+]]: i32, %[[R0_2:[^ ,)]+]]: i32):
+// CHECK: call @use3(%[[R0_0]], %[[R0_1]], %[[R0_2]])
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[D]], %[[E]], %[[F]], %[[A]], %[[B]], %[[C]]) <{operandSegmentSizes = array<i32: 1, 3, 3, 0>}>
+// CHECK: ^bb0(%[[R1_0:[^ ,)]+]]: i32, %[[R1_1:[^ ,)]+]]: i32, %[[R1_2:[^ ,)]+]]: i32):
+// CHECK: call @use3(%[[R1_0]], %[[R1_1]], %[[R1_2]])
+// CHECK: "test.segmented_region_branch_terminator"(%[[CONDITION]], %[[D]], %[[E]], %[[F]], %[[A]], %[[B]], %[[C]]) <{operandSegmentSizes = array<i32: 1, 3, 3, 0>}>
+// CHECK: })
+func.func @retain_all_operands(%condition: i1, %stop: i1, %a: i32, %b: i32, %c: i32, %d: i32, %e: i32, %f: i32) {
+ "test.segmented_region_branch"(%condition, %a, %b, %c, %d, %e, %f) <{operandSegmentSizes = array<i32: 1, 3, 3, 0>}> ({
+ ^bb0(%r0_0: i32, %r0_1: i32, %r0_2: i32):
+ func.call @use3(%r0_0, %r0_1, %r0_2) : (i32, i32, i32) -> ()
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c) <{operandSegmentSizes = array<i32: 1, 3, 3, 0>}> : (i1, i32, i32, i32, i32, i32, i32) -> ()
+ }, {
+ ^bb0(%r1_0: i32, %r1_1: i32, %r1_2: i32):
+ func.call @use3(%r1_0, %r1_1, %r1_2) : (i32, i32, i32) -> ()
+ "test.segmented_region_branch_terminator"(%condition, %d, %e, %f, %a, %b, %c) <{operandSegmentSizes = array<i32: 1, 3, 3, 0>}> : (i1, i32, i32, i32, i32, i32, i32) -> ()
+ }) : (i1, i32, i32, i32, i32, i32, i32) -> ()
+ return
+}
+func.func private @use3(i32, i32, i32)
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 4c683358e71a5..3effcfa8f040e 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1322,6 +1322,48 @@ LogicalResult TestOpWithPropertiesAndInferredType::inferReturnTypes(
return success();
}
+//===----------------------------------------------------------------------===//
+// SegmentedRegionBranchOp / SegmentedRegionBranchTerminatorOp
+//===----------------------------------------------------------------------===//
+
+void SegmentedRegionBranchOp::getCanonicalizationPatterns(
+ RewritePatternSet &patterns, MLIRContext *context) {
+ populateRegionBranchOpInterfaceCanonicalizationPatterns(patterns,
+ getOperationName());
+}
+
+void SegmentedRegionBranchOp::getSuccessorRegions(
+ RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {
+ regions.emplace_back(&getFirstRegion());
+ regions.emplace_back(&getSecondRegion());
+ if (!point.isParent())
+ regions.emplace_back(getOperation());
+}
+
+ValueRange
+SegmentedRegionBranchOp::getSuccessorInputs(RegionSuccessor successor) {
+ if (successor.isOperation())
+ return {};
+ return successor.getSuccessor()->front().getArguments();
+}
+
+OperandRange
+SegmentedRegionBranchOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return successor.getSuccessor() == &getFirstRegion() ? getFirst()
+ : getSecond();
+}
+
+MutableOperandRange
+SegmentedRegionBranchTerminatorOp::getMutableSuccessorOperands(
+ RegionSuccessor successor) {
+ if (successor.isOperation())
+ return MutableOperandRange(getOperation(), getNumOperands(), 0);
+ auto parent = cast<SegmentedRegionBranchOp>((*this)->getParentOp());
+ return successor.getSuccessor() == &parent.getFirstRegion()
+ ? getFirstMutable()
+ : getSecondMutable();
+}
+
//===----------------------------------------------------------------------===//
// LoopBlockOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 492cd44833380..96a31bf06b30e 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2797,6 +2797,32 @@ def CompareOp : TEST_Op<"compare"> {
// Test RegionBranchOpInterface
//===----------------------------------------------------------------------===//
+def SegmentedRegionBranchOp : TEST_Op<"segmented_region_branch", [
+ AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
+ let description = [{
+ Select one of two regions and pass the corresponding operand segment to
+ its block arguments. Each region can branch to either region or exit.
+ The control operands are not successor operands.
+ }];
+ let arguments = (ins I1:$condition, Variadic<AnyType>:$first,
+ Variadic<AnyType>:$second, Optional<I1>:$exitCondition);
+ let regions = (region SizedRegion<1>:$firstRegion,
+ SizedRegion<1>:$secondRegion);
+ let hasCanonicalizer = 1;
+}
+
+def SegmentedRegionBranchTerminatorOp
+ : TEST_Op<"segmented_region_branch_terminator", [
+ AttrSizedOperandSegments, Terminator,
+ HasParent<"SegmentedRegionBranchOp">,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>]> {
+ let arguments = (ins I1:$condition, Variadic<AnyType>:$first,
+ Variadic<AnyType>:$second, Optional<I1>:$exitCondition);
+}
+
def RegionIfYieldOp : TEST_Op<"region_if_yield",
[NoMemoryEffect, ReturnLike, Terminator]> {
let arguments = (ins Variadic<AnyType>:$results);
>From 615daeb263721e91aeeae9d25ed700eccc541e1f Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 12 Sep 2026 18:56:50 +0200
Subject: [PATCH 2/7] [MLIR] Extract segment-aware operand erasure into
OperationSupport
---
mlir/include/mlir/IR/OperationSupport.h | 10 ++++++
mlir/lib/IR/OperationSupport.cpp | 31 ++++++++++++++++++
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 32 +------------------
3 files changed, 42 insertions(+), 31 deletions(-)
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 54beafdeab27e..ed210192bb145 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1150,6 +1150,16 @@ struct OperationState {
MLIRContext *getContext() const { return location->getContext(); }
};
+//===----------------------------------------------------------------------===//
+// Operand mutation
+//===----------------------------------------------------------------------===//
+
+/// Erase the operands selected by `eraseIndices` and update operandSegmentSizes
+/// if the operation has AttrSizedOperandSegments. The bit vector must have one
+/// entry for each original operand. The caller must select operands that the
+/// operation permits removing and update any other dependent metadata.
+void eraseOperands(Operation *op, const BitVector &eraseIndices);
+
//===----------------------------------------------------------------------===//
// OperandStorage
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index 4ecb73398aa77..d27185c0a69b5 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -226,6 +226,37 @@ void OperationState::addRegions(
addRegion(std::move(region));
}
+//===----------------------------------------------------------------------===//
+// Operand mutation
+//===----------------------------------------------------------------------===//
+
+void mlir::eraseOperands(Operation *op, const BitVector &eraseIndices) {
+ assert(eraseIndices.size() == op->getNumOperands() &&
+ "expected one bit per operand");
+ if (eraseIndices.none())
+ return;
+
+ 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));
+}
+
//===----------------------------------------------------------------------===//
// OperandStorage
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index c7b90df464a98..d4d8cc15f65b6 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -919,41 +919,11 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
}
}
- StringAttr attrName;
// Erase operands.
for (auto &pair : operandsToRemove) {
Operation *op = pair.first;
BitVector &operands = pair.second;
- rewriter.modifyOpInPlace(op, [&]() {
- if (!op->hasTrait<OpTrait::AttrSizedOperandSegments>()) {
- op->eraseOperands(operands);
- return;
- }
- if (!attrName)
- attrName = rewriter.getStringAttr(OpTrait::AttrSizedOperandSegments<
- void>::getOperandSegmentSizeAttr());
- auto sizes = op->getAttrOfType<DenseI32ArrayAttr>(attrName);
- assert(sizes && "expected operand segment sizes attribute");
- unsigned offset = op->getNumOperands();
- // Erase from the last segment first to preserve the original indices.
- for (unsigned segment :
- llvm::reverse(llvm::seq<unsigned>(sizes.size()))) {
- unsigned size = sizes.asArrayRef()[segment];
- offset -= size;
- SmallVector<Value> retained;
- for (unsigned i : llvm::seq(size))
- if (!operands.test(offset + i))
- retained.push_back(op->getOperand(offset + i));
- if (retained.size() == size)
- continue;
- // Use getAttr to keep changes to later segments.
- MutableOperandRange range(
- op, offset, size,
- MutableOperandRange::OperandSegment(
- segment, {attrName, op->getAttr(attrName)}));
- range.assign(retained);
- }
- });
+ rewriter.modifyOpInPlace(op, [&]() { eraseOperands(op, operands); });
}
// Erase block arguments.
>From 469ebea2d30ae1cbe8e6f107065b94033c3673eb Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sun, 13 Sep 2026 10:43:11 +0200
Subject: [PATCH 3/7] [MLIR] Move segment-aware operand erasure to RewriterBase
---
mlir/include/mlir/IR/OperationSupport.h | 10 ------
mlir/include/mlir/IR/PatternMatch.h | 7 +++++
mlir/lib/IR/OperationSupport.cpp | 31 -------------------
mlir/lib/IR/PatternMatch.cpp | 29 +++++++++++++++++
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 2 +-
5 files changed, 37 insertions(+), 42 deletions(-)
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index ed210192bb145..54beafdeab27e 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1150,16 +1150,6 @@ struct OperationState {
MLIRContext *getContext() const { return location->getContext(); }
};
-//===----------------------------------------------------------------------===//
-// Operand mutation
-//===----------------------------------------------------------------------===//
-
-/// Erase the operands selected by `eraseIndices` and update operandSegmentSizes
-/// if the operation has AttrSizedOperandSegments. The bit vector must have one
-/// entry for each original operand. The caller must select operands that the
-/// operation permits removing and update any other dependent metadata.
-void eraseOperands(Operation *op, const BitVector &eraseIndices);
-
//===----------------------------------------------------------------------===//
// OperandStorage
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 531051e42b824..438766fcd24ca 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -545,6 +545,13 @@ class RewriterBase : public OpBuilder {
/// This method erases all operations in a block.
virtual void eraseBlock(Block *block);
+ /// Erase the operands selected by `eraseIndices` and update
+ /// operandSegmentSizes if the operation has AttrSizedOperandSegments. The bit
+ /// vector must have one entry for each original operand. The caller must
+ /// select operands that the operation permits removing and update any other
+ /// dependent metadata.
+ void eraseOperands(Operation *op, const BitVector &eraseIndices);
+
/// Erase the specified results of the given operation. Results cannot be
/// erased directly, so the implementation creates a new replacement
/// operation and erases the original operation. The new operation is
diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index d27185c0a69b5..4ecb73398aa77 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -226,37 +226,6 @@ void OperationState::addRegions(
addRegion(std::move(region));
}
-//===----------------------------------------------------------------------===//
-// Operand mutation
-//===----------------------------------------------------------------------===//
-
-void mlir::eraseOperands(Operation *op, const BitVector &eraseIndices) {
- assert(eraseIndices.size() == op->getNumOperands() &&
- "expected one bit per operand");
- if (eraseIndices.none())
- return;
-
- 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));
-}
-
//===----------------------------------------------------------------------===//
// OperandStorage
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 0b080d2920230..6c4e0c91fc9f3 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -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));
+ });
+}
+
Operation *RewriterBase::eraseOpResults(Operation *op,
const BitVector &eraseIndices) {
assert(op->getNumResults() == eraseIndices.size() &&
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index d4d8cc15f65b6..d505ff0772a9c 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -923,7 +923,7 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
for (auto &pair : operandsToRemove) {
Operation *op = pair.first;
BitVector &operands = pair.second;
- rewriter.modifyOpInPlace(op, [&]() { eraseOperands(op, operands); });
+ rewriter.eraseOperands(op, operands);
}
// Erase block arguments.
>From fe1ab2081f12c2124b79397338b102723a3cbbaa Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sun, 13 Sep 2026 10:51:57 +0200
Subject: [PATCH 4/7] [MLIR] Preserve return operand segments in
RemoveDeadValues
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 2 +-
.../remove-dead-values-return-segments.mlir | 77 +++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOps.td | 5 ++
3 files changed, 83 insertions(+), 1 deletion(-)
create mode 100644 mlir/test/Transforms/remove-dead-values-return-segments.mlir
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index ffaa0663b0457..a4c2aa59b4421 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -713,7 +713,7 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
<< OpWithFlags(o.op,
OpPrintingFlags().skipRegions().printGenericOpForm());
});
- o.op->eraseOperands(o.nonLive);
+ rewriter.eraseOperands(o.op, o.nonLive);
}
}
diff --git a/mlir/test/Transforms/remove-dead-values-return-segments.mlir b/mlir/test/Transforms/remove-dead-values-return-segments.mlir
new file mode 100644
index 0000000000000..46a16e528c291
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-return-segments.mlir
@@ -0,0 +1,77 @@
+// RUN: mlir-opt %s -split-input-file -remove-dead-values="canonicalize=0" | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -remove-dead-values="canonicalize=1" | FileCheck %s
+
+// Remove unused function results and update the return operand segments.
+// Check with and without canonicalization to test the cleanup pass directly.
+
+// CHECK-LABEL: func.func private @remove_last_segment() -> i32 {
+// CHECK-DAG: %[[V0:[^ ]+]] = arith.constant 11 : i32
+// CHECK: "test.segmented_return"(%[[V0]]) <{operandSegmentSizes = array<i32: 1, 0>}> : (i32) -> ()
+// CHECK-LABEL: func.func @caller()
+// CHECK: call @remove_last_segment() : () -> i32
+func.func private @remove_last_segment() -> (i32, i32) {
+ %v0 = arith.constant 11 : i32
+ %v1 = arith.constant 22 : i32
+ "test.segmented_return"(%v0, %v1)
+ <{operandSegmentSizes = array<i32: 1, 1>}> : (i32, i32) -> ()
+}
+func.func @caller() -> i32 {
+ %r:2 = func.call @remove_last_segment() : () -> (i32, i32)
+ return %r#0 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @remove_first_segment() -> i32 {
+// CHECK-DAG: %[[V1:[^ ]+]] = arith.constant 22 : i32
+// CHECK: "test.segmented_return"(%[[V1]]) <{operandSegmentSizes = array<i32: 0, 1>}> : (i32) -> ()
+// CHECK-LABEL: func.func @caller()
+// CHECK: call @remove_first_segment() : () -> i32
+func.func private @remove_first_segment() -> (i32, i32) {
+ %v0 = arith.constant 11 : i32
+ %v1 = arith.constant 22 : i32
+ "test.segmented_return"(%v0, %v1)
+ <{operandSegmentSizes = array<i32: 1, 1>}> : (i32, i32) -> ()
+}
+func.func @caller() -> i32 {
+ %r:2 = func.call @remove_first_segment() : () -> (i32, i32)
+ return %r#1 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @shrink_both_segments() -> (i32, i32) {
+// CHECK-DAG: %[[V0:[^ ]+]] = arith.constant 11 : i32
+// CHECK-DAG: %[[V3:[^ ]+]] = arith.constant 44 : i32
+// CHECK: "test.segmented_return"(%[[V0]], %[[V3]]) <{operandSegmentSizes = array<i32: 1, 1>}> : (i32, i32) -> ()
+// CHECK-LABEL: func.func @caller()
+// CHECK: call @shrink_both_segments() : () -> (i32, i32)
+func.func private @shrink_both_segments() -> (i32, i32, i32, i32) {
+ %v0 = arith.constant 11 : i32
+ %v1 = arith.constant 22 : i32
+ %v2 = arith.constant 33 : i32
+ %v3 = arith.constant 44 : i32
+ "test.segmented_return"(%v0, %v1, %v2, %v3)
+ <{operandSegmentSizes = array<i32: 2, 2>}> : (i32, i32, i32, i32) -> ()
+}
+func.func @caller() -> (i32, i32) {
+ %r:4 = func.call @shrink_both_segments() : () -> (i32, i32, i32, i32)
+ return %r#0, %r#3 : i32, i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @remove_all_results() {
+// CHECK: "test.segmented_return"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> ()
+// CHECK-LABEL: func.func @caller()
+// CHECK: call @remove_all_results() : () -> ()
+func.func private @remove_all_results() -> (i32, i32) {
+ %v0 = arith.constant 11 : i32
+ %v1 = arith.constant 22 : i32
+ "test.segmented_return"(%v0, %v1)
+ <{operandSegmentSizes = array<i32: 1, 1>}> : (i32, i32) -> ()
+}
+func.func @caller() {
+ %r:2 = func.call @remove_all_results() : () -> (i32, i32)
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 96a31bf06b30e..bbf6c4bbfe648 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2366,6 +2366,11 @@ def BlackHoleOp : TEST_Op<"blackhole">,
//===----------------------------------------------------------------------===//
def TestRegionBuilderOp : TEST_Op<"region_builder">;
+def TestSegmentedReturnOp : TEST_Op<"segmented_return", [
+ Pure, ReturnLike, Terminator, AttrSizedOperandSegments]> {
+ let arguments = (ins Variadic<AnyType>:$first, Variadic<AnyType>:$second);
+}
+
def TestReturnOp : TEST_Op<"return", [Pure, ReturnLike, Terminator]> {
let arguments = (ins Variadic<AnyType>);
let builders = [OpBuilder<(ins),
>From 01e7ce02fdeda8e1d3f69e4f1ed789f84dfc36b3 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sun, 13 Sep 2026 11:04:27 +0200
Subject: [PATCH 5/7] [MLIR][LLVM] Parse invoke operands in segment order
---
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 8 ++++----
.../LLVMIR/invoke-operand-segments.mlir | 18 ++++++++++++++++++
2 files changed, 22 insertions(+), 4 deletions(-)
create mode 100644 mlir/test/Dialect/LLVMIR/invoke-operand-segments.mlir
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 52cd73f521326..edd2d540132d0 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1799,15 +1799,15 @@ ParseResult InvokeOp::parse(OpAsmParser &parser, OperationState &result) {
parser.getBuilder(), result, argAttrs, resultAttrs,
getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
+ result.addSuccessors({normalDest, unwindDest});
+ result.addOperands(normalOperands);
+ result.addOperands(unwindOperands);
+
if (resolveOpBundleOperands(parser, opBundlesLoc, result, opBundleOperands,
opBundleOperandTypes,
getOpBundleSizesAttrName(result.name)))
return failure();
- result.addSuccessors({normalDest, unwindDest});
- result.addOperands(normalOperands);
- result.addOperands(unwindOperands);
-
int32_t numOpBundleOperands = 0;
for (const auto &operands : opBundleOperands)
numOpBundleOperands += operands.size();
diff --git a/mlir/test/Dialect/LLVMIR/invoke-operand-segments.mlir b/mlir/test/Dialect/LLVMIR/invoke-operand-segments.mlir
new file mode 100644
index 0000000000000..e7fe62bac27f6
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/invoke-operand-segments.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt %s | FileCheck %s
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// Use different types to detect operands placed in the wrong segment.
+llvm.func @target(i32)
+llvm.func @__gxx_personality_v0(...) -> i32
+
+// CHECK-LABEL: llvm.func @invoke_segments(
+// CHECK-SAME: %[[ARG:[^:]+]]: i32, %[[NORMAL:[^:]+]]: i16, %[[UNWIND:[^:]+]]: i64, %[[TAG:[^:]+]]: i8
+// CHECK: llvm.invoke @target(%[[ARG]]) to ^{{[^(]+}}(%[[NORMAL]] : i16) unwind ^{{[^(]+}}(%[[UNWIND]] : i64) ["tag"(%[[TAG]] : i8)] : (i32) -> ()
+llvm.func @invoke_segments(%arg: i32, %normal: i16, %unwind: i64, %tag: i8) attributes {personality = @__gxx_personality_v0} {
+ llvm.invoke @target(%arg) to ^normal(%normal : i16) unwind ^unwind(%unwind : i64) ["tag"(%tag : i8)] : (i32) -> ()
+^normal(%n: i16):
+ llvm.return
+^unwind(%u: i64):
+ %lp = llvm.landingpad cleanup : !llvm.struct<(ptr, i32)>
+ llvm.resume %lp : !llvm.struct<(ptr, i32)>
+}
>From c4fe68762d75c706cc78d9a2786431d1298cd38f Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sun, 13 Sep 2026 11:04:28 +0200
Subject: [PATCH 6/7] [MLIR][LLVM] Preserve segment metadata in mutable call
arguments
---
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 8 +-
...remove-dead-values-llvm-call-segments.mlir | 93 +++++++++++++++++++
.../remove-dead-values-return-segments.mlir | 4 +-
3 files changed, 99 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Transforms/remove-dead-values-llvm-call-segments.mlir
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index edd2d540132d0..fbe8daad35798 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1161,8 +1161,8 @@ Operation::operand_range CallOp::getArgOperands() {
}
MutableOperandRange CallOp::getArgOperandsMutable() {
- return MutableOperandRange(*this, getNumConsumedCalleeOperands(*this),
- getArgOperandsImpl(*this).size());
+ return getCalleeOperandsMutable().slice(getNumConsumedCalleeOperands(*this),
+ getArgOperandsImpl(*this).size());
}
/// Verify that an inlinable callsite of a debug-info-bearing function in a
@@ -1650,8 +1650,8 @@ Operation::operand_range InvokeOp::getArgOperands() {
}
MutableOperandRange InvokeOp::getArgOperandsMutable() {
- return MutableOperandRange(*this, getNumConsumedCalleeOperands(*this),
- getArgOperandsImpl(*this).size());
+ return getCalleeOperandsMutable().slice(getNumConsumedCalleeOperands(*this),
+ getArgOperandsImpl(*this).size());
}
LogicalResult InvokeOp::verify() {
diff --git a/mlir/test/Transforms/remove-dead-values-llvm-call-segments.mlir b/mlir/test/Transforms/remove-dead-values-llvm-call-segments.mlir
new file mode 100644
index 0000000000000..1be4247b521c2
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-llvm-call-segments.mlir
@@ -0,0 +1,93 @@
+// RUN: mlir-opt %s --split-input-file --remove-dead-values="canonicalize=0" | FileCheck %s
+// RUN: mlir-opt %s --split-input-file --remove-dead-values="canonicalize=0" --mlir-print-op-generic | FileCheck %s --check-prefix=GEN
+
+// Remove all call arguments. Keep the operand bundle in its own segment.
+// CHECK-LABEL: llvm.func internal @call_empty()
+// CHECK-LABEL: llvm.func @caller(
+// CHECK-SAME: %[[DEAD:[^:]+]]: i32, %[[TAG:[^:]+]]: i32
+// CHECK: llvm.call @call_empty() ["tag"(%[[TAG]] : i32)] : () -> ()
+// GEN: "llvm.call"
+// GEN-SAME: callee = @call_empty
+// GEN-SAME: op_bundle_sizes = array<i32: 1>
+// GEN-SAME: operandSegmentSizes = array<i32: 0, 1>
+llvm.func internal @call_empty(%dead: i32) attributes {sym_visibility = "private"} {
+ llvm.return
+}
+llvm.func @caller(%dead: i32, %tag: i32) {
+ llvm.call @call_empty(%dead) ["tag"(%tag : i32)] : (i32) -> ()
+ llvm.return
+}
+
+// -----
+
+// Remove disjoint arguments. Keep the middle argument and the bundle operand.
+// CHECK-LABEL: llvm.func internal @call_partial(
+// CHECK-SAME: %[[LIVE:[^:]+]]: i32)
+// CHECK: llvm.call @use(%[[LIVE]]) : (i32) -> ()
+// CHECK-LABEL: llvm.func @caller(
+// CHECK-SAME: %[[A:[^:]+]]: i32, %[[B:[^:]+]]: i32, %[[C:[^:]+]]: i32, %[[TAG:[^:]+]]: i32
+// CHECK: llvm.call @call_partial(%[[B]]) ["tag"(%[[TAG]] : i32)] : (i32) -> ()
+// GEN: "llvm.call"{{.*}}callee = @call_partial
+// GEN-SAME: op_bundle_sizes = array<i32: 1>
+// GEN-SAME: operandSegmentSizes = array<i32: 1, 1>
+llvm.func @use(i32)
+llvm.func internal @call_partial(%a: i32, %b: i32, %c: i32) attributes {sym_visibility = "private"} {
+ llvm.call @use(%b) : (i32) -> ()
+ llvm.return
+}
+llvm.func @caller(%a: i32, %b: i32, %c: i32, %tag: i32) {
+ llvm.call @call_partial(%a, %b, %c) ["tag"(%tag : i32)] : (i32, i32, i32) -> ()
+ llvm.return
+}
+
+// -----
+
+// Remove all invoke arguments when the other operand segments are empty.
+// CHECK-LABEL: llvm.func internal @invoke_empty()
+// CHECK-LABEL: llvm.func @caller(
+// CHECK: llvm.invoke @invoke_empty() to ^{{.*}} unwind ^{{.*}} : () -> ()
+// GEN: "llvm.invoke"
+// GEN-SAME: callee = @invoke_empty
+// GEN-SAME: operandSegmentSizes = array<i32: 0, 0, 0, 0>
+llvm.func @__gxx_personality_v0(...) -> i32
+llvm.func internal @invoke_empty(%dead: i32) attributes {sym_visibility = "private"} {
+ llvm.return
+}
+llvm.func @caller(%dead: i32) attributes {personality = @__gxx_personality_v0} {
+ llvm.invoke @invoke_empty(%dead) to ^normal unwind ^unwind : (i32) -> ()
+^normal:
+ llvm.return
+^unwind:
+ %lp = llvm.landingpad cleanup : !llvm.struct<(ptr, i32)>
+ llvm.resume %lp : !llvm.struct<(ptr, i32)>
+}
+
+// -----
+
+// Shrink the callee segment while keeping both successor segments and a bundle.
+// CHECK-LABEL: llvm.func internal @invoke_partial(
+// CHECK-SAME: %[[LIVE:[^:]+]]: i32)
+// CHECK: llvm.call @use(%[[LIVE]]) : (i32) -> ()
+// CHECK-LABEL: llvm.func @caller(
+// CHECK-SAME: %[[DEAD:[^:]+]]: i32, %[[LIVE:[^:]+]]: i32, %[[NORMAL:[^:]+]]: i32, %[[UNWIND:[^:]+]]: i32, %[[TAG:[^:]+]]: i32
+// CHECK: llvm.invoke @invoke_partial(%[[LIVE]]) to ^{{[^(]+}}(%[[NORMAL]] : i32) unwind ^{{[^(]+}}(%[[UNWIND]] : i32) ["tag"(%[[TAG]] : i32)] : (i32) -> ()
+// GEN: "llvm.invoke"
+// GEN-SAME: callee = @invoke_partial
+// GEN-SAME: op_bundle_sizes = array<i32: 1>
+// GEN-SAME: operandSegmentSizes = array<i32: 1, 1, 1, 1>
+llvm.func @__gxx_personality_v0(...) -> i32
+llvm.func @use(i32)
+llvm.func internal @invoke_partial(%dead: i32, %live: i32) attributes {sym_visibility = "private"} {
+ llvm.call @use(%live) : (i32) -> ()
+ llvm.return
+}
+llvm.func @caller(%dead: i32, %live: i32, %normal: i32, %unwind: i32, %tag: i32) attributes {personality = @__gxx_personality_v0} {
+ llvm.invoke @invoke_partial(%dead, %live) to ^normal(%normal : i32) unwind ^unwind(%unwind : i32) ["tag"(%tag : i32)] : (i32, i32) -> ()
+^normal(%n: i32):
+ llvm.call @use(%n) : (i32) -> ()
+ llvm.return
+^unwind(%u: i32):
+ %lp = llvm.landingpad cleanup : !llvm.struct<(ptr, i32)>
+ llvm.call @use(%u) : (i32) -> ()
+ llvm.resume %lp : !llvm.struct<(ptr, i32)>
+}
diff --git a/mlir/test/Transforms/remove-dead-values-return-segments.mlir b/mlir/test/Transforms/remove-dead-values-return-segments.mlir
index 46a16e528c291..312b6a5758c0a 100644
--- a/mlir/test/Transforms/remove-dead-values-return-segments.mlir
+++ b/mlir/test/Transforms/remove-dead-values-return-segments.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -split-input-file -remove-dead-values="canonicalize=0" | FileCheck %s
-// RUN: mlir-opt %s -split-input-file -remove-dead-values="canonicalize=1" | FileCheck %s
+// RUN: mlir-opt %s --split-input-file --remove-dead-values="canonicalize=0" | FileCheck %s
+// RUN: mlir-opt %s --split-input-file --remove-dead-values="canonicalize=1" | FileCheck %s
// Remove unused function results and update the return operand segments.
// Check with and without canonicalization to test the cleanup pass directly.
>From 6ff618b7ace8d2aa39857581658f3f63bf8a2554 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sun, 13 Sep 2026 11:45:05 +0200
Subject: [PATCH 7/7] [MLIR][LLVM] Explain segment metadata in mutable argument
slices
---
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index fbe8daad35798..747615069e281 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1161,6 +1161,8 @@ Operation::operand_range CallOp::getArgOperands() {
}
MutableOperandRange CallOp::getArgOperandsMutable() {
+ // Slice the generated range to retain its segment-size metadata. A raw
+ // range would not update operandSegmentSizes when arguments are erased.
return getCalleeOperandsMutable().slice(getNumConsumedCalleeOperands(*this),
getArgOperandsImpl(*this).size());
}
@@ -1650,6 +1652,8 @@ Operation::operand_range InvokeOp::getArgOperands() {
}
MutableOperandRange InvokeOp::getArgOperandsMutable() {
+ // Slice the generated range to retain its segment-size metadata. A raw
+ // range would not update operandSegmentSizes when arguments are erased.
return getCalleeOperandsMutable().slice(getNumConsumedCalleeOperands(*this),
getArgOperandsImpl(*this).size());
}
More information about the Mlir-commits
mailing list