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

[Pankaj Dwivedi via llvm-commits]

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

This patch tracks uniform values in UniformityInfo so that unknown/new instructions are conservatively treated as divergent at the IR level. This enables passes like NaryReassociate to safely query divergence of newly created instructions. The conservative behavior is disabled for MIR.

>From 5fb8390d28f2b51dfc89ab6917646d31e4f4a204 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 | 24 +++++++++++++++++--
 llvm/include/llvm/ADT/GenericUniformityInfo.h |  1 +
 llvm/lib/Analysis/UniformityAnalysis.cpp      | 17 +++++++++++++
 .../lib/CodeGen/MachineUniformityAnalysis.cpp | 12 ++++++++++
 4 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ADT/GenericUniformityImpl.h b/llvm/include/llvm/ADT/GenericUniformityImpl.h
index 2db76a1ad9b13..2d0bc47b14fe3 100644
--- a/llvm/include/llvm/ADT/GenericUniformityImpl.h
+++ b/llvm/include/llvm/ADT/GenericUniformityImpl.h
@@ -375,6 +375,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 +396,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<Instruction>(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 +433,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 +1127,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) {