[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