[llvm] [UniformityAnalysis] Track uniform values for conservative divergence queries at IR level (PR #180509)

[Pankaj Dwivedi via llvm-commits]

https://github.com/PankajDwivedi-25 updated https://github.com/llvm/llvm-project/pull/180509

>From eb9212fb24f352c432cd1a49242dd4b1b8aa4000 Mon Sep 17 00:00:00 2001
From: padivedi <padivedi at amd.com>
Date: Mon, 9 Feb 2026 17:07:32 +0530
Subject: [PATCH] track uniform values at SSA level

---
 llvm/include/llvm/ADT/GenericUniformityImpl.h | 25 +++++++++++++++++--
 llvm/include/llvm/ADT/GenericUniformityInfo.h |  1 +
 llvm/lib/Analysis/UniformityAnalysis.cpp      | 17 +++++++++++++
 .../lib/CodeGen/MachineUniformityAnalysis.cpp | 12 +++++++++
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ADT/GenericUniformityImpl.h b/llvm/include/llvm/ADT/GenericUniformityImpl.h
index 2db76a1ad9b13..ec06b27cfa4f8 100644
--- a/llvm/include/llvm/ADT/GenericUniformityImpl.h
+++ b/llvm/include/llvm/ADT/GenericUniformityImpl.h
@@ -51,6 +51,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SparseBitVector.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/IR/Argument.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "uniformity"
@@ -375,6 +376,10 @@ template <typename ContextT> class GenericUniformityAnalysisImpl {
   /// Divergence is seeded by calls to \p markDivergent.
   void compute();
 
+  /// \brief Populate UniformValues set after divergence analysis completes.
+  /// This enables safe uniformity queries for transformation passes.
+  void finalizeUniformValues();
+
   /// \brief Whether any value was marked or analyzed to be divergent.
   bool hasDivergence() const { return !DivergentValues.empty(); }
 
@@ -392,7 +397,19 @@ template <typename ContextT> class GenericUniformityAnalysisImpl {
   };
 
   /// \brief Whether \p Val is divergent at its definition.
-  bool isDivergent(ConstValueRefT V) const { return DivergentValues.count(V); }
+  bool isDivergent(ConstValueRefT V) const {
+    // For IR: Constants and GlobalValues are never divergent.
+    if constexpr (!std::is_same<InstructionT, MachineInstr>::value) {
+      if (!isa<InstructionT>(V) && !isa<Argument>(V))
+        return false;
+    }
+    // If UniformValues is empty (MIR, or before finalization), use original
+    // logic. If UniformValues is populated (IR after finalization), unknown
+    // values are conservatively treated as divergent.
+    if (UniformValues.empty())
+      return DivergentValues.count(V);
+    return !UniformValues.count(V);
+  }
 
   bool isDivergentUse(const UseT &U) const;
 
@@ -417,6 +434,10 @@ template <typename ContextT> class GenericUniformityAnalysisImpl {
   DenseSet<ConstValueRefT> DivergentValues;
   SmallPtrSet<const BlockT *, 32> DivergentTermBlocks;
 
+  // Known uniform values (populated after analysis by finalizeUniformValues).
+  // Values NOT in this set are conservatively treated as divergent.
+  DenseSet<ConstValueRefT> UniformValues;
+
   // Internal worklist for divergence propagation.
   std::vector<const InstructionT *> Worklist;
 
@@ -1107,7 +1128,7 @@ void GenericUniformityAnalysisImpl<ContextT>::compute() {
   // Initialize worklist.
   auto DivValuesCopy = DivergentValues;
   for (const auto DivVal : DivValuesCopy) {
-    assert(isDivergent(DivVal) && "Worklist invariant violated!");
+    assert(DivergentValues.count(DivVal) && "Worklist invariant violated!");
     pushUsers(DivVal);
   }
 
diff --git a/llvm/include/llvm/ADT/GenericUniformityInfo.h b/llvm/include/llvm/ADT/GenericUniformityInfo.h
index 9376fa6ee0bae..e40b96abb1bd7 100644
--- a/llvm/include/llvm/ADT/GenericUniformityInfo.h
+++ b/llvm/include/llvm/ADT/GenericUniformityInfo.h
@@ -52,6 +52,7 @@ template <typename ContextT> class GenericUniformityInfo {
   void compute() {
     DA->initialize();
     DA->compute();
+    DA->finalizeUniformValues();
   }
 
   /// Whether any divergence was detected.
diff --git a/llvm/lib/Analysis/UniformityAnalysis.cpp b/llvm/lib/Analysis/UniformityAnalysis.cpp
index b56534935d7c2..431ca32a741e7 100644
--- a/llvm/lib/Analysis/UniformityAnalysis.cpp
+++ b/llvm/lib/Analysis/UniformityAnalysis.cpp
@@ -50,6 +50,23 @@ template <> void llvm::GenericUniformityAnalysisImpl<SSAContext>::initialize() {
   }
 }
 
+template <>
+void llvm::GenericUniformityAnalysisImpl<SSAContext>::finalizeUniformValues() {
+  // Populate UniformValues with all values that were NOT marked divergent.
+  // This enables safe uniformity queries where unknown values (e.g., newly
+  // created instructions) are conservatively treated as divergent.
+  for (const Argument &Arg : F.args()) {
+    if (!DivergentValues.count(&Arg))
+      UniformValues.insert(&Arg);
+  }
+  for (const BasicBlock &BB : F) {
+    for (const Instruction &I : BB) {
+      if (!DivergentValues.count(&I))
+        UniformValues.insert(&I);
+    }
+  }
+}
+
 template <>
 void llvm::GenericUniformityAnalysisImpl<SSAContext>::pushUsers(
     const Value *V) {
diff --git a/llvm/lib/CodeGen/MachineUniformityAnalysis.cpp b/llvm/lib/CodeGen/MachineUniformityAnalysis.cpp
index dbadb67e1e6d2..e970d9a3f9167 100644
--- a/llvm/lib/CodeGen/MachineUniformityAnalysis.cpp
+++ b/llvm/lib/CodeGen/MachineUniformityAnalysis.cpp
@@ -68,6 +68,18 @@ void llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::initialize() {
   }
 }
 
+template <>
+void llvm::GenericUniformityAnalysisImpl<
+    MachineSSAContext>::finalizeUniformValues() {
+  // For MIR, we intentionally leave UniformValues empty.
+  // This preserves the original isDivergent() behavior where unknown registers
+  // are NOT treated as divergent. MIR passes like
+  // AMDGPUGlobalISelDivergenceLowering create new registers and query their
+  // divergence, expecting the original behavior. The conservative "unknown =
+  // divergent" behavior is only needed for IR-level passes like
+  // NaryReassociate.
+}
+
 template <>
 void llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::pushUsers(
     Register Reg) {