[Mlir-commits] [mlir] [MLIR] Preserve operand segment sizes when removing dead inputs (PR #223154)
Ivan Butygin
llvmlistbot at llvm.org
Sat Sep 12 08:27:27 PDT 2026
https://github.com/Hardcode84 created https://github.com/llvm/llvm-project/pull/223154
`RemoveDeadRegionBranchOpSuccessorInputs` was calling `eraseOperands` and ignoring potential `AttrSizedOperandSegments`. Trying to use `populateRegionBranchOpInterfaceCanonicalizationPatterns` on downstream control flow op with it resulted in broken IR. Update the segments sizes if present.
Code was vibed and then edited, tests were vibed and not particularly edited.
>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] [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);
More information about the Mlir-commits
mailing list