[Mlir-commits] [mlir] [mlir][dataflow] Drop the firstIndex argument of old visitNonControlFlowArguments, add new visitNonControlFlowArguments API and use it in IntegerRangeAnalysis (PR #175210)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jan 9 09:29:51 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: lonely eagle (linuxlonelyeagle)
<details>
<summary>Changes</summary>
This PR primarily drops the `firstIndex` argument from the old `visitNonControlFlowArguments` function of `SparseForwardDataFlowAnalysis`.This argument actually originates from `AbstractSparseForwardDataFlowAnalysis`, but it is not very useful for the user.
Additionally, a dedicated `visitNonControlFlowArguments` API was added for accessing the IV of loop op. Now, we can obtain the IV's lattices through this API, without having to explicitly call getLatticeElement.I am currently using it in the IntegerRangeAnalysis.
---
Full diff: https://github.com/llvm/llvm-project/pull/175210.diff
5 Files Affected:
- (modified) mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h (+7-4)
- (modified) mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h (+19-10)
- (modified) mlir/include/mlir/Interfaces/ControlFlowInterfaces.h (+18)
- (modified) mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp (+15-18)
- (modified) mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp (+22-34)
``````````diff
diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
index 4975cedb282e4..a673a62dcf349 100644
--- a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
@@ -66,10 +66,13 @@ 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,
- ArrayRef<IntegerValueRangeLattice *> argLattices,
- unsigned firstIndex) override;
+ void visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ArrayRef<IntegerValueRangeLattice *> argLattices) override;
+
+ void visitNonControlFlowArguments(
+ const RegionSuccessor &successor,
+ 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 1bb42a246b701..012a6008c5e5b 100644
--- a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
@@ -215,7 +215,11 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
/// of loops).
virtual void visitNonControlFlowArgumentsImpl(
Operation *op, const RegionSuccessor &successor,
- ArrayRef<AbstractSparseLattice *> argLattices, unsigned firstIndex) = 0;
+ ArrayRef<AbstractSparseLattice *> argLattices) = 0;
+
+ virtual void visitNonControlFlowArgumentsImpl(
+ const RegionSuccessor &successor,
+ ArrayRef<AbstractSparseLattice *> argLattices) = 0;
/// Get the lattice element of a value.
virtual AbstractSparseLattice *getLatticeElement(Value value) = 0;
@@ -328,11 +332,11 @@ class SparseForwardDataFlowAnalysis
/// index of the first element of `argLattices` that is set by control-flow.
virtual void visitNonControlFlowArguments(Operation *op,
const RegionSuccessor &successor,
- ArrayRef<StateT *> argLattices,
- unsigned firstIndex) {
- setAllToEntryStates(argLattices.take_front(firstIndex));
- setAllToEntryStates(argLattices.drop_front(
- firstIndex + successor.getSuccessorInputs().size()));
+ ArrayRef<StateT *> argLattices) {}
+
+ virtual void visitNonControlFlowArguments(const RegionSuccessor &successor,
+ ArrayRef<StateT *> argLattices) {
+ setAllToEntryStates(argLattices);
}
protected:
@@ -383,13 +387,18 @@ class SparseForwardDataFlowAnalysis
}
void visitNonControlFlowArgumentsImpl(
Operation *op, const RegionSuccessor &successor,
- ArrayRef<AbstractSparseLattice *> argLattices,
- unsigned firstIndex) override {
+ ArrayRef<AbstractSparseLattice *> argLattices) override {
visitNonControlFlowArguments(
op, successor,
{reinterpret_cast<StateT *const *>(argLattices.begin()),
- argLattices.size()},
- firstIndex);
+ argLattices.size()});
+ }
+ void visitNonControlFlowArgumentsImpl(
+ const RegionSuccessor &successor,
+ ArrayRef<AbstractSparseLattice *> argLattices) override {
+ visitNonControlFlowArguments(
+ successor, {reinterpret_cast<StateT *const *>(argLattices.begin()),
+ argLattices.size()});
}
void setToEntryState(AbstractSparseLattice *lattice) override {
return setToEntryState(reinterpret_cast<StateT *>(lattice));
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
index b76c2891fad5a..76aa20f11d84b 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
@@ -228,6 +228,24 @@ class RegionSuccessor {
/// the current region.
ValueRange getSuccessorInputs() const { return inputs; }
+ ValueRange getRegionNonforwardedArguments() const {
+ if (isParent())
+ return {};
+ SmallVector<Value> nonForwardArguments;
+ MutableArrayRef<BlockArgument> arguments = getSuccessor()->getArguments();
+ if (arguments.empty())
+ return {};
+ ValueRange inputs = getSuccessorInputs();
+ if (inputs.empty())
+ return arguments;
+ for (BlockArgument argument : arguments) {
+ if (!llvm::is_contained(inputs, argument)) {
+ nonForwardArguments.push_back(argument);
+ }
+ }
+ return nonForwardArguments;
+ }
+
bool operator==(RegionSuccessor rhs) const {
return successor == rhs.successor && inputs == rhs.inputs;
}
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index a93e605445465..84e134171961a 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -139,7 +139,8 @@ LogicalResult IntegerRangeAnalysis::visitOperation(
void IntegerRangeAnalysis::visitNonControlFlowArguments(
Operation *op, const RegionSuccessor &successor,
- ArrayRef<IntegerValueRangeLattice *> argLattices, unsigned firstIndex) {
+ ArrayRef<IntegerValueRangeLattice *> argLattices) {
+ ValueRange inputs = successor.getSuccessorInputs();
if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
LDBG() << "Inferring ranges for "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
@@ -152,11 +153,13 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
auto arg = dyn_cast<BlockArgument>(v);
if (!arg)
return;
- if (!llvm::is_contained(successor.getSuccessor()->getArguments(), arg))
+ if (!llvm::is_contained(inputs, arg))
return;
LDBG() << "Inferred range " << attrs;
- IntegerValueRangeLattice *lattice = argLattices[arg.getArgNumber()];
+ unsigned latticeIdx =
+ std::distance(inputs.begin(), llvm::find(inputs, arg));
+ IntegerValueRangeLattice *lattice = argLattices[latticeIdx];
IntegerValueRange oldRange = lattice->getValue();
ChangeResult changed = lattice->join(attrs);
@@ -177,9 +180,13 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
};
inferrable.inferResultRangesFromOptional(argRanges, joinCallback);
- return;
}
+}
+void IntegerRangeAnalysis::visitNonControlFlowArguments(
+ const RegionSuccessor &successor,
+ ArrayRef<IntegerValueRangeLattice *> argLattices) {
+ Operation *op = successor.getSuccessor()->getParentOp();
/// 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 +211,13 @@ 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, argLattices, firstIndex);
- }
// This shouldn't be returning nullopt if there are indunction variables.
+ SmallVector<Value> ivs = successor.getRegionNonforwardedArguments();
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(ivs, lowerBounds, upperBounds, steps, argLattices)) {
Block *block = iv.getParentBlock();
APInt min = getLoopBoundFromFold(lowerBound, iv.getType(), block,
/*getUpper=*/false);
@@ -237,14 +239,9 @@ 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}));
}
}
- return;
}
-
- return SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
- op, successor, argLattices, firstIndex);
-}
+}
\ No newline at end of file
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index b9d861830dd38..9fb7364d9d662 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -183,9 +183,10 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) {
}
// Otherwise, we can't reason about the data-flow.
- return visitNonControlFlowArgumentsImpl(block->getParentOp(),
- RegionSuccessor(block->getParent()),
- argLattices, /*firstIndex=*/0);
+ return visitNonControlFlowArgumentsImpl(
+ block->getParentOp(),
+ RegionSuccessor(block->getParent(), block->getParent()->getArguments()),
+ argLattices);
}
// Iterate over the predecessors of the non-entry block.
@@ -307,29 +308,24 @@ void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
"expected the same number of successor inputs as operands");
unsigned firstIndex = 0;
- if (inputs.size() != lattices.size()) {
- if (!point->isBlockStart()) {
- if (!inputs.empty())
- firstIndex = cast<OpResult>(inputs.front()).getResultNumber();
- visitNonControlFlowArgumentsImpl(
- branch,
- RegionSuccessor(
- branch, branch->getResults().slice(firstIndex, inputs.size())),
- lattices, firstIndex);
- } 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);
- }
+ if (inputs.size() != lattices.size() && point->isBlockStart()) {
+ if (!inputs.empty())
+ firstIndex = cast<BlockArgument>(inputs.front()).getArgNumber();
+ Region *region = point->getBlock()->getParent();
+ RegionSuccessor regionSuccessor(
+ region, region->getArguments().slice(firstIndex, inputs.size()));
+ SmallVector<AbstractSparseLattice *> nonForwardArgs(
+ lattices.take_front(firstIndex));
+ nonForwardArgs.append(SmallVector<AbstractSparseLattice *>(
+ lattices.drop_front(firstIndex + inputs.size())));
+ visitNonControlFlowArgumentsImpl(
+ branch, regionSuccessor, lattices.slice(firstIndex, inputs.size()));
+ visitNonControlFlowArgumentsImpl(regionSuccessor, nonForwardArgs);
}
- for (auto it : llvm::zip(*operands, lattices.drop_front(firstIndex)))
- join(std::get<1>(it), *getLatticeElementFor(point, std::get<0>(it)));
+ for (auto [lattice, operand] :
+ llvm::zip(lattices.drop_front(firstIndex), *operands))
+ join(lattice, *getLatticeElementFor(point, operand));
}
}
@@ -612,16 +608,8 @@ void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
if (successor.isParent())
continue;
SmallVector<BlockArgument> noControlFlowArguments;
- MutableArrayRef<BlockArgument> arguments =
- successor.getSuccessor()->getArguments();
- ValueRange inputs = successor.getSuccessorInputs();
- for (BlockArgument argument : arguments) {
- // Visit blockArgument of RegionBranchOp which isn't "control
- // flow block arguments". For example, the IV of a loop.
- if (!llvm::is_contained(inputs, argument)) {
- noControlFlowArguments.push_back(argument);
- }
- }
+ for (Value arg : successor.getRegionNonforwardedArguments())
+ noControlFlowArguments.push_back(cast<BlockArgument>(arg));
visitNonControlFlowArguments(successor, noControlFlowArguments);
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/175210
More information about the Mlir-commits
mailing list