[Mlir-commits] [mlir] [mlir][dataflow] Drop the firstIndex argument of old visitNonControlFlowArguments, add new visitNonControlFlowArguments API and use it in IntegerRangeAnalysis (PR #175210)
lonely eagle
llvmlistbot at llvm.org
Sat Jan 24 01:00:57 PST 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/175210
>From 0afc994588c912b2395b123d7743a702c7c471cd Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 24 Jan 2026 08:18:58 +0000
Subject: [PATCH 1/2] drop first index and add visitNonControlFlowArguments API
for no RegionBranchOpInterface op.
---
.../Analysis/DataFlow/IntegerRangeAnalysis.h | 13 ++-
.../mlir/Analysis/DataFlow/SparseAnalysis.h | 47 ++++++---
.../DataFlow/IntegerRangeAnalysis.cpp | 98 +++++++++----------
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 27 +++--
4 files changed, 106 insertions(+), 79 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
index e549a56a6f960..9a7031c7a4fc5 100644
--- a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
@@ -66,11 +66,14 @@ class IntegerRangeAnalysis
/// function calls `InferIntRangeInterface` to provide values for block
/// arguments or tries to reduce the range on loop induction variables with
/// known bounds.
- void
- visitNonControlFlowArguments(Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs,
- ArrayRef<IntegerValueRangeLattice *> argLattices,
- unsigned firstIndex) override;
+ void visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ValueRange successorInputs,
+ ArrayRef<IntegerValueRangeLattice *> argLattices) override;
+
+ void visitNonControlFlowArguments(
+ Operation *op, Region *const region, ValueRange successorInputs,
+ ArrayRef<IntegerValueRangeLattice *> argLattices) override;
};
/// Succeeds if an op can be converted to its unsigned equivalent without
diff --git a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
index 02f699de06f99..cb050a45496fa 100644
--- a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
@@ -215,8 +215,15 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
/// of loops).
virtual void visitNonControlFlowArgumentsImpl(
Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs, ArrayRef<AbstractSparseLattice *> argLattices,
- unsigned firstIndex) = 0;
+ ValueRange successorInputs,
+ ArrayRef<AbstractSparseLattice *> argLattices) = 0;
+
+ /// Given an operation with region non-control-flow, the lattices of the entry
+ /// block arguments, compute the lattice values for block arguments.(ex. the
+ /// block argument of gpu.launch).
+ virtual void visitNonControlFlowArgumentsImpl(
+ Operation *op, Region *const region, ValueRange arguments,
+ ArrayRef<AbstractSparseLattice *> argLattices) = 0;
/// Get the lattice element of a value.
virtual AbstractSparseLattice *getLatticeElement(Value value) = 0;
@@ -325,16 +332,21 @@ class SparseForwardDataFlowAnalysis
/// operands, and a region successor, compute the lattice values for block
/// arguments that are not accounted for by the branching control flow (ex.
/// the bounds of loops). By default, this method marks all such lattice
- /// elements as having reached a pessimistic fixpoint. `firstIndex` is the
- /// index of the first element of `argLattices` that is set by control-flow.
+ /// elements as having reached a pessimistic fixpoint.
virtual void visitNonControlFlowArguments(Operation *op,
const RegionSuccessor &successor,
ValueRange successorInputs,
- ArrayRef<StateT *> argLattices,
- unsigned firstIndex) {
- setAllToEntryStates(argLattices.take_front(firstIndex));
- setAllToEntryStates(
- argLattices.drop_front(firstIndex + successorInputs.size()));
+ ArrayRef<StateT *> argLattices) {
+ setAllToEntryStates(argLattices);
+ }
+
+ /// Given an operation with region non-control-flow, the lattices of the entry
+ /// block arguments, compute the lattice values for block arguments.(ex. the
+ /// block argument of gpu.launch).
+ virtual void visitNonControlFlowArguments(Operation *op, Region *const region,
+ ValueRange successorInputs,
+ ArrayRef<StateT *> argLattices) {
+ setAllToEntryStates(argLattices);
}
protected:
@@ -385,14 +397,23 @@ class SparseForwardDataFlowAnalysis
}
void visitNonControlFlowArgumentsImpl(
Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs, ArrayRef<AbstractSparseLattice *> argLattices,
- unsigned firstIndex) override {
+ ValueRange successorInputs,
+ ArrayRef<AbstractSparseLattice *> argLattices) override {
visitNonControlFlowArguments(
op, successor, successorInputs,
{reinterpret_cast<StateT *const *>(argLattices.begin()),
- argLattices.size()},
- firstIndex);
+ argLattices.size()});
}
+
+ virtual void visitNonControlFlowArgumentsImpl(
+ Operation *op, Region *const region, ValueRange successorInputs,
+ ArrayRef<AbstractSparseLattice *> argLattices) override {
+ visitNonControlFlowArguments(
+ op, region, successorInputs,
+ {reinterpret_cast<StateT *const *>(argLattices.begin()),
+ argLattices.size()});
+ }
+
void setToEntryState(AbstractSparseLattice *lattice) override {
return setToEntryState(reinterpret_cast<StateT *>(lattice));
}
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index 012d8384d3098..c0a0f2b63c327 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -139,47 +139,7 @@ LogicalResult IntegerRangeAnalysis::visitOperation(
void IntegerRangeAnalysis::visitNonControlFlowArguments(
Operation *op, const RegionSuccessor &successor, ValueRange successorInputs,
- ArrayRef<IntegerValueRangeLattice *> argLattices, unsigned firstIndex) {
- if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
- LDBG() << "Inferring ranges for "
- << OpWithFlags(op, OpPrintingFlags().skipRegions());
-
- auto argRanges = llvm::map_to_vector(op->getOperands(), [&](Value value) {
- return getLatticeElementFor(getProgramPointAfter(op), value)->getValue();
- });
-
- auto joinCallback = [&](Value v, const IntegerValueRange &attrs) {
- auto arg = dyn_cast<BlockArgument>(v);
- if (!arg)
- return;
- if (!llvm::is_contained(successor.getSuccessor()->getArguments(), arg))
- return;
-
- LDBG() << "Inferred range " << attrs;
- IntegerValueRangeLattice *lattice = argLattices[arg.getArgNumber()];
- IntegerValueRange oldRange = lattice->getValue();
-
- ChangeResult changed = lattice->join(attrs);
-
- // Catch loop results with loop variant bounds and conservatively make
- // them [-inf, inf] so we don't circle around infinitely often (because
- // the dataflow analysis in MLIR doesn't attempt to work out trip counts
- // and often can't).
- bool isYieldedValue = llvm::any_of(v.getUsers(), [](Operation *op) {
- return op->hasTrait<OpTrait::IsTerminator>();
- });
- if (isYieldedValue && !oldRange.isUninitialized() &&
- !(lattice->getValue() == oldRange)) {
- LDBG() << "Loop variant loop result detected";
- changed |= lattice->join(IntegerValueRange::getMaxRange(v));
- }
- propagateIfChanged(lattice, changed);
- };
-
- inferrable.inferResultRangesFromOptional(argRanges, joinCallback);
- return;
- }
-
+ ArrayRef<IntegerValueRangeLattice *> argLattices) {
/// Given a lower bound, upper bound, or step from a LoopLikeInterface return
/// the lower/upper bound for that result if possible.
auto getLoopBoundFromFold = [&](OpFoldResult loopBound, Type boundType,
@@ -204,18 +164,12 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
// Infer bounds for loop arguments that have static bounds
if (auto loop = dyn_cast<LoopLikeOpInterface>(op)) {
- std::optional<llvm::SmallVector<Value>> maybeIvs =
- loop.getLoopInductionVars();
- if (!maybeIvs) {
- return SparseForwardDataFlowAnalysis ::visitNonControlFlowArguments(
- op, successor, successorInputs, argLattices, firstIndex);
- }
// This shouldn't be returning nullopt if there are indunction variables.
SmallVector<OpFoldResult> lowerBounds = *loop.getLoopLowerBounds();
SmallVector<OpFoldResult> upperBounds = *loop.getLoopUpperBounds();
SmallVector<OpFoldResult> steps = *loop.getLoopSteps();
- for (auto [iv, lowerBound, upperBound, step] :
- llvm::zip_equal(*maybeIvs, lowerBounds, upperBounds, steps)) {
+ for (auto [iv, lowerBound, upperBound, step, ivEntry] : llvm::zip_equal(
+ successorInputs, lowerBounds, upperBounds, steps, argLattices)) {
Block *block = iv.getParentBlock();
APInt min = getLoopBoundFromFold(lowerBound, iv.getType(), block,
/*getUpper=*/false);
@@ -237,7 +191,6 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
// resulting range is meaningless and should not be used in further
// inferences.
if (max.sge(min)) {
- IntegerValueRangeLattice *ivEntry = getLatticeElement(iv);
auto ivRange = ConstantIntRanges::fromSigned(min, max);
propagateIfChanged(ivEntry, ivEntry->join(IntegerValueRange{ivRange}));
}
@@ -246,5 +199,48 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
}
return SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
- op, successor, successorInputs, argLattices, firstIndex);
+ op, successor, successorInputs, argLattices);
}
+
+void IntegerRangeAnalysis::visitNonControlFlowArguments(
+ Operation *op, Region *const region, ValueRange successorInputs,
+ ArrayRef<IntegerValueRangeLattice *> argLattices) {
+ if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
+ LDBG() << "Inferring ranges for "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ auto argRanges = llvm::map_to_vector(op->getOperands(), [&](Value value) {
+ return getLatticeElementFor(getProgramPointAfter(op), value)->getValue();
+ });
+
+ auto joinCallback = [&](Value v, const IntegerValueRange &attrs) {
+ auto arg = dyn_cast<BlockArgument>(v);
+ if (!arg)
+ return;
+ if (!llvm::is_contained(region->getArguments(), arg))
+ return;
+
+ LDBG() << "Inferred range " << attrs;
+ IntegerValueRangeLattice *lattice = argLattices[arg.getArgNumber()];
+ IntegerValueRange oldRange = lattice->getValue();
+
+ ChangeResult changed = lattice->join(attrs);
+
+ // Catch loop results with loop variant bounds and conservatively make
+ // them [-inf, inf] so we don't circle around infinitely often (because
+ // the dataflow analysis in MLIR doesn't attempt to work out trip counts
+ // and often can't).
+ bool isYieldedValue = llvm::any_of(v.getUsers(), [](Operation *op) {
+ return op->hasTrait<OpTrait::IsTerminator>();
+ });
+ if (isYieldedValue && !oldRange.isUninitialized() &&
+ !(lattice->getValue() == oldRange)) {
+ LDBG() << "Loop variant loop result detected";
+ changed |= lattice->join(IntegerValueRange::getMaxRange(v));
+ }
+ propagateIfChanged(lattice, changed);
+ };
+
+ inferrable.inferResultRangesFromOptional(argRanges, joinCallback);
+ return;
+ }
+}
\ No newline at end of file
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 7dad9676e7e53..3a87c4efae370 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -187,8 +187,8 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) {
// Otherwise, we can't reason about the data-flow.
return visitNonControlFlowArgumentsImpl(
- block->getParentOp(), RegionSuccessor(block->getParent()), ValueRange(),
- argLattices, /*firstIndex=*/0);
+ block->getParentOp(), block->getParent(), block->getParent()->getArguments(),
+ argLattices);
}
// Iterate over the predecessors of the non-entry block.
@@ -309,23 +309,30 @@ void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
assert(inputs.size() == operands->size() &&
"expected the same number of successor inputs as operands");
+ auto valueToLattices = [&](Value v) { return getLatticeElement(v); };
unsigned firstIndex = 0;
if (inputs.size() != lattices.size()) {
if (!point->isBlockStart()) {
if (!inputs.empty())
firstIndex = cast<OpResult>(inputs.front()).getResultNumber();
- visitNonControlFlowArgumentsImpl(
- branch, RegionSuccessor::parent(),
- branch->getResults().slice(firstIndex, inputs.size()), lattices,
- firstIndex);
+ SmallVector<Value> nonSuccessorInputs =
+ branch.getNonSuccessorInputs(RegionSuccessor::parent());
+ SmallVector<AbstractSparseLattice *> nonSuccessorInputLattices =
+ llvm::map_to_vector(nonSuccessorInputs, valueToLattices);
+ visitNonControlFlowArgumentsImpl(branch, RegionSuccessor::parent(),
+ nonSuccessorInputs,
+ nonSuccessorInputLattices);
} else {
if (!inputs.empty())
firstIndex = cast<BlockArgument>(inputs.front()).getArgNumber();
Region *region = point->getBlock()->getParent();
- visitNonControlFlowArgumentsImpl(
- branch, RegionSuccessor(region),
- region->getArguments().slice(firstIndex, inputs.size()), lattices,
- firstIndex);
+ SmallVector<Value> nonSuccessorInputs =
+ branch.getNonSuccessorInputs(RegionSuccessor(region));
+ SmallVector<AbstractSparseLattice *> nonSuccessorInputLattices =
+ llvm::map_to_vector(nonSuccessorInputs, valueToLattices);
+ visitNonControlFlowArgumentsImpl(branch, RegionSuccessor(region),
+ nonSuccessorInputs,
+ nonSuccessorInputLattices);
}
}
>From 73789e07e3dead9c90bf746ea67c84a0ba78da39 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 24 Jan 2026 08:28:34 +0000
Subject: [PATCH 2/2] fix nit.
---
.../mlir/Analysis/DataFlow/IntegerRangeAnalysis.h | 9 +++++----
mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h | 10 +++++-----
mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp | 2 +-
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 4 ++--
4 files changed, 13 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
index 9a7031c7a4fc5..157c24d40e74e 100644
--- a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
@@ -63,16 +63,17 @@ class IntegerRangeAnalysis
/// Visit block arguments or operation results of an operation with region
/// control-flow for which values are not defined by region control-flow. This
- /// function calls `InferIntRangeInterface` to provide values for block
- /// arguments or tries to reduce the range on loop induction variables with
+ /// function tries to reduce the range on loop induction variables with
/// known bounds.
void visitNonControlFlowArguments(
Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs,
+ ValueRange nonSuccessorInputs,
ArrayRef<IntegerValueRangeLattice *> argLattices) override;
+ /// This function calls `InferIntRangeInterface` to provide values for block
+ /// arguments which is not RegionBranchOpInterface, (ex, gpu.launch).
void visitNonControlFlowArguments(
- Operation *op, Region *const region, ValueRange successorInputs,
+ Operation *op, Region *const region, ValueRange nonSuccessorInputs,
ArrayRef<IntegerValueRangeLattice *> argLattices) override;
};
diff --git a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
index cb050a45496fa..e1d66c230d3cb 100644
--- a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
@@ -215,7 +215,7 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
/// of loops).
virtual void visitNonControlFlowArgumentsImpl(
Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs,
+ ValueRange noSuccessorInputs,
ArrayRef<AbstractSparseLattice *> argLattices) = 0;
/// Given an operation with region non-control-flow, the lattices of the entry
@@ -397,19 +397,19 @@ class SparseForwardDataFlowAnalysis
}
void visitNonControlFlowArgumentsImpl(
Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs,
+ ValueRange noSuccessorInputs,
ArrayRef<AbstractSparseLattice *> argLattices) override {
visitNonControlFlowArguments(
- op, successor, successorInputs,
+ op, successor, noSuccessorInputs,
{reinterpret_cast<StateT *const *>(argLattices.begin()),
argLattices.size()});
}
virtual void visitNonControlFlowArgumentsImpl(
- Operation *op, Region *const region, ValueRange successorInputs,
+ Operation *op, Region *const region, ValueRange arguments,
ArrayRef<AbstractSparseLattice *> argLattices) override {
visitNonControlFlowArguments(
- op, region, successorInputs,
+ op, region, arguments,
{reinterpret_cast<StateT *const *>(argLattices.begin()),
argLattices.size()});
}
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index c0a0f2b63c327..286fd88352bc7 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -243,4 +243,4 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
inferrable.inferResultRangesFromOptional(argRanges, joinCallback);
return;
}
-}
\ No newline at end of file
+}
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 3a87c4efae370..187c0d646dc7f 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -187,8 +187,8 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) {
// Otherwise, we can't reason about the data-flow.
return visitNonControlFlowArgumentsImpl(
- block->getParentOp(), block->getParent(), block->getParent()->getArguments(),
- argLattices);
+ block->getParentOp(), block->getParent(),
+ block->getParent()->getArguments(), argLattices);
}
// Iterate over the predecessors of the non-entry block.
More information about the Mlir-commits
mailing list