[NFC][TTI] Introduce getInstructionUniformity API for uniformity analysis (#168903)

This patch introduces a new TargetTransformInfo hook
`getInstructionUniformity()`
that provides a unified interface for querying target-specific
uniformity
information about instructions and values.

The new hook returns an `InstructionUniformity` enum with three values:
- Default: Result is uniform if all operands are uniform (standard
propagation)
- AlwaysUniform: Result is always uniform regardless of operands
- NeverUniform: Result can never be assumed uniform

This API wraps the existing `isAlwaysUniform()` and
`isSourceOfDivergence()`
hooks, providing a single entry point for uniformity queries. Both LLVM
IR-level
(via TTI) and MIR-level (via TargetInstrInfo) uniformity analysis have
been
updated to use the new hook.

Target implementations:
- AMDGPU: Wraps existing `isAlwaysUniform()` and
`isSourceOfDivergence()` hooks
- NVPTX: Wraps existing `isSourceOfDivergence()` hook

This is an NFC change - all implementations return conservative defaults
or wrap
existing functionality.

Ref patch:#137639

Author: PankajDwivedi-25

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -24,6 +24,7 @@
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitmaskEnum.h"
+#include "llvm/ADT/Uniformity.h"
 #include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/Analysis/InterestingMemoryOperand.h"
 #include "llvm/IR/FMF.h"
@@ -511,16 +512,14 @@ class TargetTransformInfo {
   /// uniformity analysis and assume all values are uniform.
   LLVM_ABI bool hasBranchDivergence(const Function *F = nullptr) const;
 
-  /// Returns whether V is a source of divergence.
-  ///
-  /// This function provides the target-dependent information for
-  /// the target-independent UniformityAnalysis.
-  LLVM_ABI bool isSourceOfDivergence(const Value *V) const;
-
-  // Returns true for the target specific
-  // set of operations which produce uniform result
-  // even taking non-uniform arguments
-  LLVM_ABI bool isAlwaysUniform(const Value *V) const;
+  /// Get target-specific uniformity information for an instruction.
+  /// This allows targets to provide more fine-grained control over
+  /// uniformity analysis by specifying whether specific instructions
+  /// should always or never be considered uniform, or require custom
+  /// operand-based analysis.
+  /// \param V The value to query for uniformity information.
+  /// \return InstructionUniformity.
+  LLVM_ABI InstructionUniformity getInstructionUniformity(const Value *V) const;
 
   /// Query the target whether the specified address space cast from FromAS to
   /// ToAS is valid.

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -131,9 +131,9 @@ class TargetTransformInfoImplBase {
     return false;
   }
 
-  virtual bool isSourceOfDivergence(const Value *V) const { return false; }
-
-  virtual bool isAlwaysUniform(const Value *V) const { return false; }
+  virtual InstructionUniformity getInstructionUniformity(const Value *V) const {
+    return InstructionUniformity::Default;
+  }
 
   virtual bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const {
     return false;

llvm/include/llvm/CodeGen/BasicTTIImpl.h

@@ -410,10 +410,6 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
     return false;
   }
 
-  bool isSourceOfDivergence(const Value *V) const override { return false; }
-
-  bool isAlwaysUniform(const Value *V) const override { return false; }
-
   bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const override {
     return false;
   }

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -294,16 +294,14 @@ bool TargetTransformInfo::hasBranchDivergence(const Function *F) const {
   return TTIImpl->hasBranchDivergence(F);
 }
 
-bool TargetTransformInfo::isSourceOfDivergence(const Value *V) const {
+InstructionUniformity
+llvm::TargetTransformInfo::getInstructionUniformity(const Value *V) const {
+  // Calls with the NoDivergenceSource attribute are always uniform.
   if (const auto *Call = dyn_cast<CallBase>(V)) {
     if (Call->hasFnAttr(Attribute::NoDivergenceSource))
-      return false;
+      return InstructionUniformity::AlwaysUniform;
   }
-  return TTIImpl->isSourceOfDivergence(V);
-}
-
-bool llvm::TargetTransformInfo::isAlwaysUniform(const Value *V) const {
-  return TTIImpl->isAlwaysUniform(V);
+  return TTIImpl->getInstructionUniformity(V);
 }
 
 bool llvm::TargetTransformInfo::isValidAddrSpaceCast(unsigned FromAS,

llvm/lib/Analysis/UniformityAnalysis.cpp

@@ -31,15 +31,22 @@ bool llvm::GenericUniformityAnalysisImpl<SSAContext>::markDefsDivergent(
 
 template <> void llvm::GenericUniformityAnalysisImpl<SSAContext>::initialize() {
   for (auto &I : instructions(F)) {
-    if (TTI->isSourceOfDivergence(&I))
-      markDivergent(I);
-    else if (TTI->isAlwaysUniform(&I))
+    InstructionUniformity IU = TTI->getInstructionUniformity(&I);
+    switch (IU) {
+    case InstructionUniformity::AlwaysUniform:
       addUniformOverride(I);
+      continue;
+    case InstructionUniformity::NeverUniform:
+      markDivergent(I);
+      continue;
+    case InstructionUniformity::Default:
+      break;
+    }
   }
   for (auto &Arg : F.args()) {
-    if (TTI->isSourceOfDivergence(&Arg)) {
+    if (TTI->getInstructionUniformity(&Arg) ==
+        InstructionUniformity::NeverUniform)
       markDivergent(&Arg);
-    }
   }
 }
 

llvm/lib/CodeGen/MachineUniformityAnalysis.cpp

@@ -53,13 +53,16 @@ void llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::initialize() {
   for (const MachineBasicBlock &block : F) {
     for (const MachineInstr &instr : block) {
       auto uniformity = InstrInfo.getInstructionUniformity(instr);
-      if (uniformity == InstructionUniformity::AlwaysUniform) {
-        addUniformOverride(instr);
-        continue;
-      }
 
-      if (uniformity == InstructionUniformity::NeverUniform) {
+      switch (uniformity) {
+      case InstructionUniformity::AlwaysUniform:
+        addUniformOverride(instr);
+        break;
+      case InstructionUniformity::NeverUniform:
         markDivergent(instr);
+        break;
+      case InstructionUniformity::Default:
+        break;
       }
     }
   }

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

@@ -1651,3 +1651,14 @@ unsigned GCNTTIImpl::getNumberOfParts(Type *Tp) const {
   }
   return BaseT::getNumberOfParts(Tp);
 }
+
+InstructionUniformity
+GCNTTIImpl::getInstructionUniformity(const Value *V) const {
+  if (isAlwaysUniform(V))
+    return InstructionUniformity::AlwaysUniform;
+
+  if (isSourceOfDivergence(V))
+    return InstructionUniformity::NeverUniform;
+
+  return InstructionUniformity::Default;
+}

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

@@ -101,6 +101,14 @@ class GCNTTIImpl final : public BasicTTIImplBase<GCNTTIImpl> {
 
   std::pair<InstructionCost, MVT> getTypeLegalizationCost(Type *Ty) const;
 
+  /// \returns true if V might be divergent even when all of its operands
+  /// are uniform.
+  bool isSourceOfDivergence(const Value *V) const;
+
+  /// Returns true for the target specific set of operations which produce
+  /// uniform result even taking non-uniform arguments.
+  bool isAlwaysUniform(const Value *V) const;
+
 public:
   explicit GCNTTIImpl(const AMDGPUTargetMachine *TM, const Function &F);
 
@@ -174,8 +182,6 @@ class GCNTTIImpl final : public BasicTTIImplBase<GCNTTIImpl> {
                                      const Value *Op1) const override;
 
   bool isReadRegisterSourceOfDivergence(const IntrinsicInst *ReadReg) const;
-  bool isSourceOfDivergence(const Value *V) const override;
-  bool isAlwaysUniform(const Value *V) const override;
 
   bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const override {
     // Address space casts must cast between different address spaces.
@@ -302,6 +308,8 @@ class GCNTTIImpl final : public BasicTTIImplBase<GCNTTIImpl> {
   /// together under a single i32 value. Otherwise fall back to base
   /// implementation.
   unsigned getNumberOfParts(Type *Tp) const override;
+
+  InstructionUniformity getInstructionUniformity(const Value *V) const override;
 };
 
 } // end namespace llvm

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp

@@ -674,3 +674,11 @@ void NVPTXTTIImpl::collectKernelLaunchBounds(
   if (MaxNTID.size() > 2)
     LB.push_back({"maxntidz", MaxNTID[2]});
 }
+
+InstructionUniformity
+NVPTXTTIImpl::getInstructionUniformity(const Value *V) const {
+  if (isSourceOfDivergence(V))
+    return InstructionUniformity::NeverUniform;
+
+  return InstructionUniformity::Default;
+}

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

@@ -37,6 +37,10 @@ class NVPTXTTIImpl final : public BasicTTIImplBase<NVPTXTTIImpl> {
   const NVPTXSubtarget *getST() const { return ST; };
   const NVPTXTargetLowering *getTLI() const { return TLI; };
 
+  /// \returns true if the result of the value could potentially be
+  /// different across threads in a warp.
+  bool isSourceOfDivergence(const Value *V) const;
+
 public:
   explicit NVPTXTTIImpl(const NVPTXTargetMachine *TM, const Function &F)
       : BaseT(TM, F.getDataLayout()), ST(TM->getSubtargetImpl()),
@@ -46,8 +50,6 @@ class NVPTXTTIImpl final : public BasicTTIImplBase<NVPTXTTIImpl> {
     return true;
   }
 
-  bool isSourceOfDivergence(const Value *V) const override;
-
   unsigned getFlatAddressSpace() const override {
     return AddressSpace::ADDRESS_SPACE_GENERIC;
   }
@@ -201,6 +203,8 @@ class NVPTXTTIImpl final : public BasicTTIImplBase<NVPTXTTIImpl> {
     // Self-referential globals are not supported.
     return false;
   }
+
+  InstructionUniformity getInstructionUniformity(const Value *V) const override;
 };
 
 } // end namespace llvm