let getInstructionUniformity hook wrap isSourceOfDivergence/isAlwaysUniform

Author: PankajDwivedi-25

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -2000,7 +2000,8 @@ class TargetTransformInfo {
   /// Returns true if GEP should not be used to index into vectors for this
   /// target.
   LLVM_ABI bool allowVectorElementIndexingUsingGEP() const;
-  InstructionUniformity getInstructionUniformity(const Instruction &I) const;
+
+  InstructionUniformity getInstructionUniformity(const Value *V) const;
 
 private:
   std::unique_ptr<const TargetTransformInfoImplBase> TTIImpl;

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -1153,8 +1153,11 @@ class TargetTransformInfoImplBase {
       SmallVectorImpl<std::pair<StringRef, int64_t>> &LB) const {}
 
   virtual bool allowVectorElementIndexingUsingGEP() const { return true; }
-  virtual InstructionUniformity
-  getInstructionUniformity(const Instruction &I) const {
+
+  // New API for uniformity classification
+  // Targets should override this to provide target-specific uniformity analysis
+  // The default implementation returns Default (conservative behavior)
+  virtual InstructionUniformity getInstructionUniformity(const Value *V) const {
     return InstructionUniformity::Default;
   }
 

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -1518,8 +1518,8 @@ bool TargetTransformInfo::allowVectorElementIndexingUsingGEP() const {
 }
 
 InstructionUniformity
-TargetTransformInfo::getInstructionUniformity(const Instruction &I) const {
-  return TTIImpl->getInstructionUniformity(I);
+TargetTransformInfo::getInstructionUniformity(const Value *V) const {
+  return TTIImpl->getInstructionUniformity(V);
 }
 
 TargetTransformInfoImplBase::~TargetTransformInfoImplBase() = default;

llvm/lib/Analysis/UniformityAnalysis.cpp

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

llvm/lib/CodeGen/MachineUniformityAnalysis.cpp

@@ -148,15 +148,11 @@ bool llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::isDivergentUse(
   return isTemporalDivergent(*UseInstr->getParent(), *DefInstr);
 }
 
+// This can be defined later depending on use of the MachineUniformityAnalysis.
 template <>
 bool GenericUniformityAnalysisImpl<MachineSSAContext>::isOperandUniform(
-    const MachineInstr &I, InstructionUniformity IU) const {
-  switch (IU) {
-  case InstructionUniformity::AnyOfFirstTwoUseOp:
-    return !isDivergentUse(I.getOperand(1)) || !isDivergentUse(I.getOperand(2));
-  default:
-    return false;
-  }
+    const MachineInstr &MI, InstructionUniformity IU) const {
+  return false;
 }
 
 // This ensures explicit instantiation of

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

@@ -1575,16 +1575,31 @@ unsigned GCNTTIImpl::getNumberOfParts(Type *Tp) const {
   return BaseT::getNumberOfParts(Tp);
 }
 
+// New API that wraps the old isSourceOfDivergence and isAlwaysUniform APIs
+// with additional support for new uniformity classifications
 InstructionUniformity
-GCNTTIImpl::getInstructionUniformity(const Instruction &I) const {
-  if (const auto *II = dyn_cast<IntrinsicInst>(&I)) {
-    switch (II->getIntrinsicID()) {
+GCNTTIImpl::getInstructionUniformity(const Value *V) const {
+  // Check for new special cases first (permlane16/permlanex16)
+  // These need operand-dependent uniformity analysis
+  if (const IntrinsicInst *Intrinsic = dyn_cast<IntrinsicInst>(V)) {
+    switch (Intrinsic->getIntrinsicID()) {
     case Intrinsic::amdgcn_permlane16:
     case Intrinsic::amdgcn_permlanex16:
+      // Result value can be uniform if either of first two operands are uniform
       return InstructionUniformity::AnyOfFirstTwoUseOp;
     default:
       break;
     }
   }
+
+  // Delegate to old APIs for backward compatibility
+  if (isAlwaysUniform(V))
+    return InstructionUniformity::AlwaysUniform;
+
+  // Check if source of divergence
+  if (isSourceOfDivergence(V))
+    return InstructionUniformity::NeverUniform;
+
+  // Default behavior
   return InstructionUniformity::Default;
 }

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

@@ -302,8 +302,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 Instruction &I) const override;
+
+  InstructionUniformity getInstructionUniformity(const Value *V) const override;
 };
 
 } // end namespace llvm

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp

@@ -635,3 +635,16 @@ void NVPTXTTIImpl::collectKernelLaunchBounds(
   if (MaxNTID.size() > 2)
     LB.push_back({"maxntidz", MaxNTID[2]});
 }
+
+// New API that wraps the old isSourceOfDivergence API
+// NVPTX doesn't have isAlwaysUniform, so we only delegate to
+// isSourceOfDivergence
+InstructionUniformity
+NVPTXTTIImpl::getInstructionUniformity(const Value *V) const {
+  // Delegate to old API for backward compatibility
+  if (isSourceOfDivergence(V))
+    return InstructionUniformity::NeverUniform;
+
+  // Default behavior
+  return InstructionUniformity::Default;
+}

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h

@@ -195,6 +195,7 @@ class NVPTXTTIImpl final : public BasicTTIImplBase<NVPTXTTIImpl> {
     // Self-referential globals are not supported.
     return false;
   }
+  InstructionUniformity getInstructionUniformity(const Value *V) const override;
 };
 
 } // end namespace llvm