[VPlan] Share PreservesUniformity logic between isSingleScalar and isUniformAcrossVFsAndUFs

Extract the PreservesUniformity logic from isSingleScalar into a shared
static helper function. Update isUniformAcrossVFsAndUFs to use this
logic for VPWidenRecipe and VPInstruction, so that any opcode that
preserves uniformity is considered uniform-across-vf-and-uf if its
operands are.

This unifies the uniformity checking logic and makes it easier to extend
in the future.

This should effectively by NFC currently.

1 files changed, 30 insertions, 22 deletions

diff --git a/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp b/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
index cffc40960e47..939216fe162a 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
@@ -150,24 +150,26 @@ const SCEV *vputils::getSCEVExprForVPValue(const VPValue *V,
       .Default([&SE](const VPRecipeBase *) { return SE.getCouldNotCompute(); });
 }
 
-bool vputils::isSingleScalar(const VPValue *VPV) {
-  auto PreservesUniformity = [](unsigned Opcode) -> bool {
-    if (Instruction::isBinaryOp(Opcode) || Instruction::isCast(Opcode))
-      return true;
-    switch (Opcode) {
-    case Instruction::GetElementPtr:
-    case Instruction::ICmp:
-    case Instruction::FCmp:
-    case Instruction::Select:
-    case VPInstruction::Not:
-    case VPInstruction::Broadcast:
-    case VPInstruction::PtrAdd:
-      return true;
-    default:
-      return false;
-    }
-  };
+/// Returns true if \p Opcode preserves uniformity, i.e., if all operands are
+/// uniform, the result will also be uniform.
+static bool preservesUniformity(unsigned Opcode) {
+  if (Instruction::isBinaryOp(Opcode) || Instruction::isCast(Opcode))
+    return true;
+  switch (Opcode) {
+  case Instruction::GetElementPtr:
+  case Instruction::ICmp:
+  case Instruction::FCmp:
+  case Instruction::Select:
+  case VPInstruction::Not:
+  case VPInstruction::Broadcast:
+  case VPInstruction::PtrAdd:
+    return true;
+  default:
+    return false;
+  }
+}
 
+bool vputils::isSingleScalar(const VPValue *VPV) {
   // A live-in must be uniform across the scope of VPlan.
   if (VPV->isLiveIn())
     return true;
@@ -179,19 +181,19 @@ bool vputils::isSingleScalar(const VPValue *VPV) {
     // lanes.
     if (RegionOfR && RegionOfR->isReplicator())
       return false;
-    return Rep->isSingleScalar() || (PreservesUniformity(Rep->getOpcode()) &&
+    return Rep->isSingleScalar() || (preservesUniformity(Rep->getOpcode()) &&
                                      all_of(Rep->operands(), isSingleScalar));
   }
   if (isa<VPWidenGEPRecipe, VPDerivedIVRecipe, VPBlendRecipe,
           VPWidenSelectRecipe>(VPV))
     return all_of(VPV->getDefiningRecipe()->operands(), isSingleScalar);
   if (auto *WidenR = dyn_cast<VPWidenRecipe>(VPV)) {
-    return PreservesUniformity(WidenR->getOpcode()) &&
+    return preservesUniformity(WidenR->getOpcode()) &&
            all_of(WidenR->operands(), isSingleScalar);
   }
   if (auto *VPI = dyn_cast<VPInstruction>(VPV))
     return VPI->isSingleScalar() || VPI->isVectorToScalar() ||
-           (PreservesUniformity(VPI->getOpcode()) &&
+           (preservesUniformity(VPI->getOpcode()) &&
             all_of(VPI->operands(), isSingleScalar));
   if (isa<VPPartialReductionRecipe>(VPV))
     return false;
@@ -234,9 +236,15 @@ bool vputils::isUniformAcrossVFsAndUFs(VPValue *V) {
                 isa<AssumeInst, StoreInst>(R->getUnderlyingInstr())) &&
                all_of(R->operands(), isUniformAcrossVFsAndUFs);
       })
+      .Case<VPWidenRecipe>([](const auto *R) {
+        return preservesUniformity(R->getOpcode()) &&
+               all_of(R->operands(), isUniformAcrossVFsAndUFs);
+      })
       .Case<VPInstruction>([](const auto *VPI) {
-        return VPI->isScalarCast() &&
-               isUniformAcrossVFsAndUFs(VPI->getOperand(0));
+        return (VPI->isScalarCast() &&
+                isUniformAcrossVFsAndUFs(VPI->getOperand(0))) ||
+               (preservesUniformity(VPI->getOpcode()) &&
+                all_of(VPI->operands(), isUniformAcrossVFsAndUFs));
       })
       .Case<VPWidenCastRecipe>([](const auto *R) {
         // A cast is uniform according to its operand.