[Mlir-commits] [mlir] a986c09 - [MLIR] Preserve operand segment sizes when removing dead inputs (#223154)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 19 02:38:55 PDT 2026
Author: Ivan Butygin
Date: 2026-09-19T12:38:48+03:00
New Revision: a986c0984aa7ec77636fdc085f4228fec5a4a600
URL: https://github.com/llvm/llvm-project/commit/a986c0984aa7ec77636fdc085f4228fec5a4a600
DIFF: https://github.com/llvm/llvm-project/commit/a986c0984aa7ec77636fdc085f4228fec5a4a600.diff
LOG: [MLIR] Preserve operand segment sizes when removing dead inputs (#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.
Also fix LLVM dialect `CallOp` and `InvokeOp` `getArgOperandsMutable`
which were dropping segment sizes and a parser bug uncovered by the
tests.
Code was vibed and then edited, tests were vibed and not particularly
edited.
Added:
mlir/test/Dialect/LLVMIR/invoke-operand-segments.mlir
mlir/test/Transforms/region-branch-operand-segments.mlir
mlir/test/Transforms/remove-dead-values-llvm-call-segments.mlir
mlir/test/Transforms/remove-dead-values-return-segments.mlir
Modified:
mlir/include/mlir/IR/PatternMatch.h
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
mlir/lib/IR/PatternMatch.cpp
mlir/lib/Interfaces/ControlFlowInterfaces.cpp
mlir/lib/Transforms/RemoveDeadValues.cpp
mlir/test/lib/Dialect/Test/TestOpDefs.cpp
mlir/test/lib/Dialect/Test/TestOps.td
Removed:
################################################################################
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/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index e70d09bcdb0cc..e6b31fc3f9f8b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1160,8 +1160,10 @@ Operation::operand_range CallOp::getArgOperands() {
}
MutableOperandRange CallOp::getArgOperandsMutable() {
- return MutableOperandRange(*this, getNumConsumedCalleeOperands(*this),
- getArgOperandsImpl(*this).size());
+ // 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());
}
/// Verify that an inlinable callsite of a debug-info-bearing function in a
@@ -1649,8 +1651,10 @@ Operation::operand_range InvokeOp::getArgOperands() {
}
MutableOperandRange InvokeOp::getArgOperandsMutable() {
- return MutableOperandRange(*this, getNumConsumedCalleeOperands(*this),
- getArgOperandsImpl(*this).size());
+ // 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());
}
LogicalResult InvokeOp::verify() {
@@ -1798,15 +1802,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/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 0b080d2920230..8a24ffafd0b90 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -244,6 +244,37 @@ 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;
+
+ // TODO: Add an interface to update the operand segment-size property
+ // directly, without converting it to and from an attribute.
+ 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 475969a7c0d09..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, [&]() { op->eraseOperands(operands); });
+ rewriter.eraseOperands(op, operands);
}
// Erase block arguments.
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/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
diff erent 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)>
+}
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
diff er, 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/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
new file mode 100644
index 0000000000000..312b6a5758c0a
--- /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/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..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),
@@ -2797,6 +2802,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