[llvm] [GlobalISel] Add iterative known-bits evaluation to GISelValueTracking (PR #221755)

[via llvm-commits]

https://github.com/daemonpilot updated https://github.com/llvm/llvm-project/pull/221755

>From 2e958743737b7f5ecccc86ce31fd2e6c9faaa277 Mon Sep 17 00:00:00 2001
From: daemonpilot <daemonpilot at hotmail.com>
Date: Mon, 7 Sep 2026 23:03:05 +0800
Subject: [PATCH] [GlobalISel] Add iterative known-bits evaluation to
 GISelValueTracking

---
 .../CodeGen/GlobalISel/GISelValueTracking.h   | 24 ++++++
 .../CodeGen/GlobalISel/GISelValueTracking.cpp | 77 ++++++++++++++++++-
 2 files changed, 98 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/GISelValueTracking.h b/llvm/include/llvm/CodeGen/GlobalISel/GISelValueTracking.h
index d8352ca2b2003..05a9f95c5e1c0 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/GISelValueTracking.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/GISelValueTracking.h
@@ -38,6 +38,13 @@ class LLVM_ABI GISelValueTracking : public GISelChangeObserver {
   const DataLayout &DL;
   unsigned MaxDepth;
 
+  // The triple (Register, DemandedElts, Depth) used as worklist and cache key.
+  using WorkItem = std::tuple<Register, APInt, unsigned>;
+  // Items pending evaluation during the active top-level query.
+  SmallVector<WorkItem, 6> Stack;
+  // Memoised results for the current top-level query.
+  DenseMap<WorkItem, KnownBits> Results;
+
   void computeKnownBitsMin(Register Src0, Register Src1, KnownBits &Known,
                            const APInt &DemandedElts, unsigned Depth = 0);
 
@@ -56,6 +63,23 @@ class LLVM_ABI GISelValueTracking : public GISelChangeObserver {
                            FPClassTest InterestedClasses, KnownFPClass &Known,
                            unsigned Depth);
 
+  void computeKnownBits(Register R, KnownBits &Known, const APInt &DemandedElts,
+                        unsigned Depth = 0);
+
+  bool getKnownBitsResult(const Register &Reg, const APInt &DemandedElts,
+                          unsigned Depth, KnownBits &Known) {
+    auto It = Results.find({Reg, DemandedElts, Depth});
+    if (It == Results.end())
+      return false;
+    Known = It->second;
+    return true;
+  }
+
+  void setKnownBitsResult(const Register &Reg, const APInt &DemandedElts,
+                          unsigned Depth, const KnownBits &Known) {
+    Results[{Reg, DemandedElts, Depth}] = Known;
+  }
+
 public:
   GISelValueTracking(MachineFunction &MF, unsigned MaxDepth = 6);
   ~GISelValueTracking() override = default;
diff --git a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
index d8e94cecc5e2a..442aab8e5b25b 100644
--- a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
@@ -249,9 +249,9 @@ static KnownBits extractBits(unsigned BitWidth, const KnownBits &SrcOpKnown,
   return KnownBits::lshr(SrcOpKnown, OffsetKnown) & Mask;
 }
 
-void GISelValueTracking::computeKnownBitsImpl(Register R, KnownBits &Known,
-                                              const APInt &DemandedElts,
-                                              unsigned Depth) {
+void GISelValueTracking::computeKnownBits(Register R, KnownBits &Known,
+                                          const APInt &DemandedElts,
+                                          unsigned Depth) {
   MachineInstr &MI = *MRI.getVRegDef(R);
   unsigned Opcode = MI.getOpcode();
   LLT DstTy = MRI.getType(R);
@@ -1138,6 +1138,77 @@ void GISelValueTracking::computeKnownBitsImpl(Register R, KnownBits &Known,
   LLVM_DEBUG(dumpResult(MI, Known, Depth));
 }
 
+static void genUnknown(MachineRegisterInfo &MRI, const Register &Reg,
+                       KnownBits &Known) {
+  LLT Ty = MRI.getType(Reg);
+  if (!Ty.isValid()) {
+    Known = KnownBits();
+    return;
+  }
+  unsigned BitWidth = Ty.getScalarSizeInBits();
+  Known = KnownBits(BitWidth);
+}
+
+/// Evaluate a known-bits query with an explicit worklist instead of recursive
+/// descent.
+void GISelValueTracking::computeKnownBitsImpl(Register R, KnownBits &Known,
+                                              const APInt &DemandedElts,
+                                              unsigned Depth) {
+  const bool TopLevel = Stack.empty();
+
+  // Nested queries only consult the per-query cache. If the result is not
+  // available yet, enqueue the request and return an unknown placeholder.
+  if (!TopLevel) {
+    if (!getKnownBitsResult(R, DemandedElts, Depth, Known)) {
+      Stack.push_back({R, DemandedElts, Depth});
+      genUnknown(MRI, R, Known);
+    }
+    return;
+  }
+
+  // Top-level queries drive evaluation iteratively until every queued item has
+  // either been computed or found in the cache.
+  Stack.push_back({R, DemandedElts, Depth});
+  while (!Stack.empty()) {
+    WorkItem Item = Stack.back();
+    size_t StackSize = Stack.size();
+    const Register &ItemReg = std::get<0>(Item);
+    const APInt &ItemDemandedElts = std::get<1>(Item);
+    const unsigned ItemDepth = std::get<2>(Item);
+    KnownBits ItemKnown;
+
+    if (getKnownBitsResult(ItemReg, ItemDemandedElts, ItemDepth, ItemKnown)) {
+      Stack.pop_back();
+      continue;
+    }
+
+    // Evaluate this item with the per-instruction known-bits logic. Dependent
+    // queries issued from there re-enter this worklist driver and take the
+    // `!TopLevel` path to enqueue more work.
+    computeKnownBits(ItemReg, ItemKnown, ItemDemandedElts, ItemDepth);
+
+    // If evaluating this item did not queue more work, its dependencies are
+    // resolved and the result can be memoized immediately.
+    if (Stack.size() == StackSize) {
+      assert((std::get<0>(Stack.back()) == ItemReg &&
+              std::get<1>(Stack.back()) == ItemDemandedElts &&
+              std::get<2>(Stack.back()) == ItemDepth) &&
+             "The item we just evaluated must still be the top one.");
+
+      setKnownBitsResult(ItemReg, ItemDemandedElts, ItemDepth, ItemKnown);
+      Stack.pop_back();
+    }
+  }
+
+  // The original query must have been computed by the time the worklist is
+  // drained.
+  if (!getKnownBitsResult(R, DemandedElts, Depth, Known))
+    llvm_unreachable(
+        "Top level query must be in `results` after iteration is complete.");
+
+  Results.clear();
+}
+
 void GISelValueTracking::computeKnownFPClass(Register R, KnownFPClass &Known,
                                              FPClassTest InterestedClasses,
                                              unsigned Depth) {