[Mlir-commits] [mlir] [mlir] Add support for Static Single Information types (PR #210435)
Eric Schweitz
llvmlistbot at llvm.org
Mon Jul 27 12:58:26 PDT 2026
https://github.com/schweitzpgi updated https://github.com/llvm/llvm-project/pull/210435
>From 175c76a23ad07db770a7002bb92f17030148481b Mon Sep 17 00:00:00 2001
From: Eric Schweitz <eschweitz at nvidia.com>
Date: Fri, 17 Jul 2026 14:12:04 -0700
Subject: [PATCH 1/3] [mlir] Add support for Static Single Information types
For type systems with SSI types (static single information), the
path information from the def to the use should not be considered
redundant, and it is a logical error to eliminate it.
This PR adds a type trait that can be added to classify a type as
being SSI, teaches the canonicalizer to not consider values of these
types as "redundant" block arguments, and adds regression tests.
Signed-off-by: Eric Schweitz <eschweitz at nvidia.com>
---
mlir/include/mlir/IR/Types.h | 10 ++++
mlir/lib/Transforms/Utils/RegionUtils.cpp | 4 ++
.../canonicalize-redundant-ssi-type.mlir | 47 +++++++++++++++++++
mlir/test/lib/Dialect/Test/TestTypeDefs.td | 8 ++++
4 files changed, 69 insertions(+)
create mode 100644 mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
diff --git a/mlir/include/mlir/IR/Types.h b/mlir/include/mlir/IR/Types.h
index 70599b448f76e..d54dbbe7f2f4d 100644
--- a/mlir/include/mlir/IR/Types.h
+++ b/mlir/include/mlir/IR/Types.h
@@ -252,6 +252,16 @@ template <typename ConcreteType, template <typename> class TraitType>
using TraitBase = detail::StorageUserTraitBase<ConcreteType, TraitType>;
} // namespace TypeTrait
+namespace TypeTrait {
+/// Types with this trait opt out of redundant block argument simplification
+/// during region simplification. Block arguments of these types are not
+/// eliminated when all predecessors pass the same value. Use this to preserve
+/// SSI (Static Single Information) form where block arguments carry semantic
+/// meaning beyond their SSA def-use structure.
+template <typename ConcreteType>
+struct SSIType : public TraitBase<ConcreteType, SSIType> {};
+} // namespace TypeTrait
+
//===----------------------------------------------------------------------===//
// TypeInterface
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index ae3ea83758dab..313ddcad6221f 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -15,6 +15,7 @@
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -1095,6 +1096,9 @@ static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
// Go through the arguments of the block.
for (auto [argIdx, blockOperand] : llvm::enumerate(block.getArguments())) {
+ // Block arguments with SSIType are excluded from redundant arg elimination.
+ if (blockOperand.getType().hasTrait<TypeTrait::SSIType>())
+ continue;
bool sameArg = true;
Value commonValue;
diff --git a/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir b/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
new file mode 100644
index 0000000000000..4d5f06efb9588
--- /dev/null
+++ b/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
@@ -0,0 +1,47 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file \
+// RUN: -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}))' \
+// RUN: | FileCheck %s
+
+// Verify that a block argument with SSIType is NOT removed by
+// dropRedundantArguments, even when all predecessors pass the same value.
+
+// CHECK-LABEL: func @redundant_ssi_arg_preserved
+// CHECK: cf.br ^bb1(%{{.*}} : !test.ssi_type)
+// CHECK: ^bb1(%{{.*}}: !test.ssi_type):
+// CHECK-NEXT: return
+func.func @redundant_ssi_arg_preserved(%arg0: !test.ssi_type) {
+ cf.br ^succ(%arg0 : !test.ssi_type)
+^succ(%0: !test.ssi_type):
+ return
+}
+
+// -----
+
+// Verify that a normal (non-SSI) redundant block argument IS removed by
+// dropRedundantArguments, confirming the baseline behavior we are opting out
+// of.
+
+// CHECK-LABEL: func @redundant_normal_arg_eliminated
+// CHECK: cf.br ^bb1
+// CHECK: ^bb1:
+// CHECK-NEXT: return
+func.func @redundant_normal_arg_eliminated(%arg0: f32) {
+ cf.br ^succ(%arg0 : f32)
+^succ(%0: f32):
+ return
+}
+
+// -----
+
+// Verify that SSI and non-SSI args can coexist: the non-SSI dead arg is
+// eliminated while the SSI arg is preserved.
+
+// CHECK-LABEL: func @mixed_args_selective_preservation
+// CHECK: "test.br"(%{{.*}})[^bb1]
+// CHECK: ^bb1(%{{.*}}: !test.ssi_type):
+// CHECK-NEXT: return
+func.func @mixed_args_selective_preservation(%a: !test.ssi_type, %b: f32) {
+ "test.br"(%a, %b)[^succ] : (!test.ssi_type, f32) -> ()
+^succ(%0: !test.ssi_type, %1: f32):
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index 9481aa1b35efd..9384e34704ceb 100644
--- a/mlir/test/lib/Dialect/Test/TestTypeDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
@@ -189,6 +189,14 @@ def TestTypeWithTrait : Test_Type<"TestTypeWithTrait", [TestTypeTrait]> {
let mnemonic = "test_type_with_trait";
}
+def SSITypeTrait : NativeTypeTrait<"SSIType">;
+
+// A type that carries SSIType trait, opting block arguments of this type out
+// of block argument simplification during region simplification.
+def TestSSIType : Test_Type<"TestSSI", [SSITypeTrait]> {
+ let mnemonic = "ssi_type";
+}
+
// Type with assembly format.
def TestTypeWithFormat : Test_Type<"TestTypeWithFormat"> {
let parameters = (
>From f369e277b68413aab4dbed4854ce9636f0c492e5 Mon Sep 17 00:00:00 2001
From: Eric Schweitz <eschweitz at nvidia.com>
Date: Mon, 27 Jul 2026 10:08:28 -0700
Subject: [PATCH 2/3] Add guards so that SCF dialect will also respect SSI
constraints. Add regression tests. Update the cf test for dead values.
---
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 10 ++
.../canonicalize-redundant-ssi-type.mlir | 23 +++-
.../Transforms/canonicalize-ssi-type-scf.mlir | 105 ++++++++++++++++++
3 files changed, 136 insertions(+), 2 deletions(-)
create mode 100644 mlir/test/Transforms/canonicalize-ssi-type-scf.mlir
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 475969a7c0d09..ba26715796e6f 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -736,6 +736,11 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
// Nothing to do for successor inputs that are already dead.
if (value.use_empty())
continue;
+ // Do not fold SSI-typed successor inputs across control-flow paths.
+ // Each block argument of SSI type is a distinct information point even
+ // when all predecessors forward the same dominating value.
+ if (value.getType().hasTrait<TypeTrait::SSIType>())
+ continue;
// Nothing to do for successor inputs that may have multiple reachable
// values.
llvm::SmallDenseSet<Value> reachableValues;
@@ -1046,6 +1051,11 @@ struct RemoveDuplicateSuccessorInputUses : public RewritePattern {
// Total complexity: O(n * k * max(log k, log n)). For each input, sorting
// the signature costs O(k log k) and the std::map lookup costs O(k log n).
for (Value input : inputs) {
+ // Do not deduplicate SSI-typed successor inputs. Each represents a
+ // distinct control-flow information point even when all predecessors
+ // forward identical values.
+ if (input.getType().hasTrait<TypeTrait::SSIType>())
+ continue;
// Gather the predecessor value for each predecessor (region branch
// point) and sort them to form this input's signature.
Signature sig;
diff --git a/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir b/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
index 4d5f06efb9588..d052979e54067 100644
--- a/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
+++ b/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
@@ -7,10 +7,27 @@
// CHECK-LABEL: func @redundant_ssi_arg_preserved
// CHECK: cf.br ^bb1(%{{.*}} : !test.ssi_type)
-// CHECK: ^bb1(%{{.*}}: !test.ssi_type):
+// CHECK: ^bb1(%[[V:.*]]: !test.ssi_type):
+// CHECK-NEXT: "test.use"(%[[V]])
// CHECK-NEXT: return
func.func @redundant_ssi_arg_preserved(%arg0: !test.ssi_type) {
cf.br ^succ(%arg0 : !test.ssi_type)
+^succ(%0: !test.ssi_type):
+ "test.use"(%0) : (!test.ssi_type) -> ()
+ return
+}
+
+// -----
+
+// Verify that a dead SSI block argument IS removed -- deadness is safe to
+// fold; only collapsing a live arg across control-flow paths is forbidden.
+
+// CHECK-LABEL: func @dead_ssi_arg_eliminated
+// CHECK: cf.br ^bb1
+// CHECK: ^bb1:
+// CHECK-NEXT: return
+func.func @dead_ssi_arg_eliminated(%arg0: !test.ssi_type) {
+ cf.br ^succ(%arg0 : !test.ssi_type)
^succ(%0: !test.ssi_type):
return
}
@@ -38,10 +55,12 @@ func.func @redundant_normal_arg_eliminated(%arg0: f32) {
// CHECK-LABEL: func @mixed_args_selective_preservation
// CHECK: "test.br"(%{{.*}})[^bb1]
-// CHECK: ^bb1(%{{.*}}: !test.ssi_type):
+// CHECK: ^bb1(%[[V:.*]]: !test.ssi_type):
+// CHECK-NEXT: "test.use"(%[[V]])
// CHECK-NEXT: return
func.func @mixed_args_selective_preservation(%a: !test.ssi_type, %b: f32) {
"test.br"(%a, %b)[^succ] : (!test.ssi_type, f32) -> ()
^succ(%0: !test.ssi_type, %1: f32):
+ "test.use"(%0) : (!test.ssi_type) -> ()
return
}
diff --git a/mlir/test/Transforms/canonicalize-ssi-type-scf.mlir b/mlir/test/Transforms/canonicalize-ssi-type-scf.mlir
new file mode 100644
index 0000000000000..fae43caed89ed
--- /dev/null
+++ b/mlir/test/Transforms/canonicalize-ssi-type-scf.mlir
@@ -0,0 +1,105 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file \
+// RUN: -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}))' \
+// RUN: | FileCheck %s
+
+// Tests for SSIType preservation through SCF control-flow canonicalization
+// patterns registered via RegionBranchOpInterface:
+//
+// - MakeRegionBranchOpSuccessorInputsDead: must NOT fold a live SSI iter_arg
+// to its unique reachable dominating value across control-flow paths.
+// - RemoveDuplicateSuccessorInputUses: must NOT replace one SSI iter_arg with
+// another that has an identical operand signature.
+// - RemoveDeadRegionBranchOpSuccessorInputs: MAY remove a dead SSI iter_arg
+// (deadness is safe; only cross-path collapsing is forbidden).
+
+// -----
+
+// Verify that a live SSI iter_arg is NOT folded to its unique reachable value
+// by MakeRegionBranchOpSuccessorInputsDead, even though every predecessor
+// (loop entry and back-edge yield) forwards the same dominating value %x.
+
+// CHECK-LABEL: func @scf_for_live_ssi_not_folded
+// CHECK: %[[R:.*]] = scf.for {{.*}} iter_args(%[[ARG:.*]] = %{{.*}}) -> (!test.ssi_type)
+// CHECK: "test.use"(%[[ARG]])
+// CHECK: "test.use"(%[[R]])
+func.func @scf_for_live_ssi_not_folded(%x: !test.ssi_type,
+ %lb: index, %ub: index, %step: index) {
+ %r = scf.for %i = %lb to %ub step %step iter_args(%arg1 = %x)
+ -> (!test.ssi_type) {
+ "test.use"(%arg1) : (!test.ssi_type) -> ()
+ scf.yield %x : !test.ssi_type
+ }
+ "test.use"(%r) : (!test.ssi_type) -> ()
+ return
+}
+
+// -----
+
+// Verify that two live SSI iter_args with an identical operand signature are
+// NOT deduplicated by RemoveDuplicateSuccessorInputUses. Each represents a
+// distinct SSI definition point even though all predecessors forward %x for
+// both.
+
+// CHECK-LABEL: func @scf_for_duplicate_ssi_args_not_deduped
+// CHECK: scf.for {{.*}} iter_args(%[[A0:.*]] = %{{.*}}, %[[A1:.*]] = %{{.*}})
+// CHECK: "test.use"(%[[A0]])
+// CHECK: "test.use"(%[[A1]])
+// CHECK: "test.use"(%{{.*}})
+// CHECK: "test.use"(%{{.*}})
+func.func @scf_for_duplicate_ssi_args_not_deduped(%x: !test.ssi_type,
+ %lb: index, %ub: index,
+ %step: index) {
+ %r0, %r1 = scf.for %i = %lb to %ub step %step
+ iter_args(%arg0 = %x, %arg1 = %x) -> (!test.ssi_type, !test.ssi_type) {
+ "test.use"(%arg0) : (!test.ssi_type) -> ()
+ "test.use"(%arg1) : (!test.ssi_type) -> ()
+ scf.yield %x, %x : !test.ssi_type, !test.ssi_type
+ }
+ "test.use"(%r0) : (!test.ssi_type) -> ()
+ "test.use"(%r1) : (!test.ssi_type) -> ()
+ return
+}
+
+// -----
+
+// Verify that a dead SSI iter_arg IS removed by
+// RemoveDeadRegionBranchOpSuccessorInputs. Deadness is safe to fold; only
+// collapsing a live arg across control-flow paths is forbidden.
+
+// CHECK-LABEL: func @scf_for_dead_ssi_iter_arg_removed
+// CHECK-NOT: iter_args
+func.func @scf_for_dead_ssi_iter_arg_removed(%x: !test.ssi_type,
+ %lb: index, %ub: index,
+ %step: index) {
+ scf.for %i = %lb to %ub step %step iter_args(%arg1 = %x)
+ -> (!test.ssi_type) {
+ scf.yield %x : !test.ssi_type
+ }
+ return
+}
+
+// -----
+
+// Verify the same live-SSI preservation for scf.while: neither the
+// before-region arg nor the after-region arg should be folded to %x by
+// MakeRegionBranchOpSuccessorInputsDead.
+
+// CHECK-LABEL: func @scf_while_live_ssi_not_folded
+// CHECK: %[[R:.*]] = scf.while (%[[BA:.*]] = %{{.*}}) : (!test.ssi_type) -> !test.ssi_type
+// CHECK: "test.use"(%[[BA]])
+// CHECK: scf.condition(%{{.*}}) %[[BA]]
+// CHECK: ^bb0(%[[AA:.*]]: !test.ssi_type):
+// CHECK: "test.use"(%[[AA]])
+// CHECK: return %[[R]]
+func.func @scf_while_live_ssi_not_folded(%x: !test.ssi_type,
+ %cond: i1) -> !test.ssi_type {
+ %res = scf.while (%arg0 = %x) : (!test.ssi_type) -> !test.ssi_type {
+ "test.use"(%arg0) : (!test.ssi_type) -> ()
+ scf.condition(%cond) %arg0 : !test.ssi_type
+ } do {
+ ^bb0(%arg1: !test.ssi_type):
+ "test.use"(%arg1) : (!test.ssi_type) -> ()
+ scf.yield %x : !test.ssi_type
+ }
+ return %res : !test.ssi_type
+}
>From adaa93cd245178fd922ce57edf355289294cfde4 Mon Sep 17 00:00:00 2001
From: Eric Schweitz <eschweitz at nvidia.com>
Date: Mon, 27 Jul 2026 12:57:49 -0700
Subject: [PATCH 3/3] Fix CHECK lines and add more cases to the regression
tests.
---
.../canonicalize-redundant-ssi-type.mlir | 75 ++++++++++++++-----
1 file changed, 57 insertions(+), 18 deletions(-)
diff --git a/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir b/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
index d052979e54067..a8118048f7a55 100644
--- a/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
+++ b/mlir/test/Transforms/canonicalize-redundant-ssi-type.mlir
@@ -4,30 +4,32 @@
// Verify that a block argument with SSIType is NOT removed by
// dropRedundantArguments, even when all predecessors pass the same value.
+// Uses "test.br" (not cf.br) so block merging doesn't obscure the result.
// CHECK-LABEL: func @redundant_ssi_arg_preserved
-// CHECK: cf.br ^bb1(%{{.*}} : !test.ssi_type)
+// CHECK: "test.br"(%{{.*}})[^bb1] : (!test.ssi_type) -> ()
// CHECK: ^bb1(%[[V:.*]]: !test.ssi_type):
-// CHECK-NEXT: "test.use"(%[[V]])
-// CHECK-NEXT: return
-func.func @redundant_ssi_arg_preserved(%arg0: !test.ssi_type) {
- cf.br ^succ(%arg0 : !test.ssi_type)
+// CHECK-NEXT: %[[W:.*]] = "test.use"(%[[V]])
+// CHECK-NEXT: return %[[W]]
+func.func @redundant_ssi_arg_preserved(%arg0: !test.ssi_type) -> !test.ssi_type {
+ "test.br"(%arg0)[^succ] : (!test.ssi_type) -> ()
^succ(%0: !test.ssi_type):
- "test.use"(%0) : (!test.ssi_type) -> ()
- return
+ %1 = "test.use"(%0) : (!test.ssi_type) -> !test.ssi_type
+ return %1 : !test.ssi_type
}
// -----
// Verify that a dead SSI block argument IS removed -- deadness is safe to
// fold; only collapsing a live arg across control-flow paths is forbidden.
+// Uses "test.br" so block merging doesn't obscure the result.
// CHECK-LABEL: func @dead_ssi_arg_eliminated
-// CHECK: cf.br ^bb1
+// CHECK: "test.br"()[^bb1]
// CHECK: ^bb1:
// CHECK-NEXT: return
func.func @dead_ssi_arg_eliminated(%arg0: !test.ssi_type) {
- cf.br ^succ(%arg0 : !test.ssi_type)
+ "test.br"(%arg0)[^succ] : (!test.ssi_type) -> ()
^succ(%0: !test.ssi_type):
return
}
@@ -39,13 +41,11 @@ func.func @dead_ssi_arg_eliminated(%arg0: !test.ssi_type) {
// of.
// CHECK-LABEL: func @redundant_normal_arg_eliminated
-// CHECK: cf.br ^bb1
-// CHECK: ^bb1:
// CHECK-NEXT: return
-func.func @redundant_normal_arg_eliminated(%arg0: f32) {
+func.func @redundant_normal_arg_eliminated(%arg0: f32) -> f32 {
cf.br ^succ(%arg0 : f32)
^succ(%0: f32):
- return
+ return %0 : f32
}
// -----
@@ -56,11 +56,50 @@ func.func @redundant_normal_arg_eliminated(%arg0: f32) {
// CHECK-LABEL: func @mixed_args_selective_preservation
// CHECK: "test.br"(%{{.*}})[^bb1]
// CHECK: ^bb1(%[[V:.*]]: !test.ssi_type):
-// CHECK-NEXT: "test.use"(%[[V]])
-// CHECK-NEXT: return
-func.func @mixed_args_selective_preservation(%a: !test.ssi_type, %b: f32) {
+// CHECK-NEXT: %[[W:.*]] = "test.use"(%[[V]])
+// CHECK-NEXT: return %[[W]]
+func.func @mixed_args_selective_preservation(%a: !test.ssi_type, %b: f32) -> (!test.ssi_type, f32){
"test.br"(%a, %b)[^succ] : (!test.ssi_type, f32) -> ()
^succ(%0: !test.ssi_type, %1: f32):
- "test.use"(%0) : (!test.ssi_type) -> ()
- return
+ %2 = "test.use"(%0) : (!test.ssi_type) -> !test.ssi_type
+ return %2, %1 : !test.ssi_type, f32
+}
+
+// -----
+
+// CHECK-LABEL: func @ssi_preserved_block_merged(
+// CHECK-SAME: %[[V:.*]]: !test.ssi_type)
+// CHECK-NEXT: %[[W:.*]] = "test.use"(%[[V]])
+// CHECK-NEXT: return %[[W]]
+func.func @ssi_preserved_block_merged(%arg0: !test.ssi_type) -> !test.ssi_type {
+ cf.br ^succ(%arg0 : !test.ssi_type)
+^succ(%0: !test.ssi_type):
+ %1 = "test.use"(%0) : (!test.ssi_type) -> !test.ssi_type
+ return %1 : !test.ssi_type
+}
+
+// -----
+
+// CHECK-LABEL: func.func @ssi_preserved_cond_br(
+// CHECK-SAME: %[[ARG0:.*]]: i1,
+// CHECK-SAME: %[[ARG1:.*]]: !test.ssi_type)
+// CHECK: cf.cond_br %[[ARG0]], ^bb1(%[[ARG1]] : !test.ssi_type), ^bb2(%[[ARG1]] : !test.ssi_type)
+// CHECK: ^bb1(%[[VAL_0:[^:]*]]:
+// CHECK: %[[VAL_1:.*]] = "test.use1"(%[[VAL_0]])
+// CHECK: cf.br ^bb3(%[[VAL_1]]
+// CHECK: ^bb2(%[[VAL_2:[^:]*]]:
+// CHECK: %[[VAL_3:.*]] = "test.use2"(%[[VAL_2]])
+// CHECK: cf.br ^bb3(%[[VAL_3]]
+// CHECK: ^bb3(%[[VAL_4:[^:]*]]:
+// CHECK: return %[[VAL_4]]
+func.func @ssi_preserved_cond_br(%arg0: i1, %arg1: !test.ssi_type) -> !test.ssi_type {
+ cf.cond_br %arg0, ^succ0(%arg1 : !test.ssi_type), ^succ1(%arg1 : !test.ssi_type)
+^succ0(%0: !test.ssi_type):
+ %1 = "test.use1"(%0) : (!test.ssi_type) -> !test.ssi_type
+ cf.br ^exit(%1 : !test.ssi_type)
+^succ1(%2: !test.ssi_type):
+ %3 = "test.use2"(%2) : (!test.ssi_type) -> !test.ssi_type
+ cf.br ^exit(%3 : !test.ssi_type)
+^exit(%4: !test.ssi_type):
+ return %4 : !test.ssi_type
}
More information about the Mlir-commits
mailing list