[VPlan] Move out canNarrowOps (NFC).

[fhahn]

Move definition of canNarrowOps out to static function, to make it easier to extend + generalize

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

Changes

Move definition of canNarrowOps out to static function, to make it easier to extend + generalize

---

Full diff: https://github.com/llvm/llvm-project/pull/167309.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (+28-15)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index b319fbc7a78c0..210d372fba66e 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -4142,6 +4142,32 @@ static bool canNarrowLoad(VPWidenRecipe *WideMember0, unsigned OpIdx,
   return false;
 }
 
+static bool canNarrowOps(ArrayRef<VPValue *> Ops) {
+  SmallVector<VPValue *> Ops0;
+  auto *WideMember0 = dyn_cast<VPWidenRecipe>(Ops[0]);
+  if (!WideMember0)
+    return false;
+
+  for (unsigned Idx = 0; Idx != WideMember0->getNumOperands(); ++Idx) {
+    SmallVector<VPValue *> Ops0;
+    for (const auto &[I, V] : enumerate(Ops)) {
+      auto *R = dyn_cast_or_null<VPWidenRecipe>(V->getDefiningRecipe());
+      if (!R || R->getOpcode() != WideMember0->getOpcode() ||
+          R->getNumOperands() > 2)
+        return false;
+
+      Ops0.push_back(R->getOperand(Idx));
+    }
+    if (any_of(enumerate(Ops0), [WideMember0, Idx](const auto &P) {
+          const auto &[OpIdx, OpV] = P;
+          return !canNarrowLoad(WideMember0, Idx, OpV, OpIdx);
+        }))
+      return false;
+  }
+
+  return true;
+}
+
 /// Returns true if \p IR is a full interleave group with factor and number of
 /// members both equal to \p VF. The interleave group must also access the full
 /// vector width \p VectorRegWidth.
@@ -4313,22 +4339,9 @@ void VPlanTransforms::narrowInterleaveGroups(VPlan &Plan, ElementCount VF,
 
     // Check if all values feeding InterleaveR are matching wide recipes, which
     // operands that can be narrowed.
-    auto *WideMember0 = dyn_cast_or_null<VPWidenRecipe>(
-        InterleaveR->getStoredValues()[0]->getDefiningRecipe());
-    if (!WideMember0)
+    //
+    if (!canNarrowOps(InterleaveR->getStoredValues()))
       return;
-    for (const auto &[I, V] : enumerate(InterleaveR->getStoredValues())) {
-      auto *R = dyn_cast_or_null<VPWidenRecipe>(V->getDefiningRecipe());
-      if (!R || R->getOpcode() != WideMember0->getOpcode() ||
-          R->getNumOperands() > 2)
-        return;
-      if (any_of(enumerate(R->operands()),
-                 [WideMember0, Idx = I](const auto &P) {
-                   const auto &[OpIdx, OpV] = P;
-                   return !canNarrowLoad(WideMember0, OpIdx, OpV, Idx);
-                 }))
-        return;
-    }
     StoreGroups.push_back(InterleaveR);
   }
 

[artagnon]

Stray //?

[david-arm]

Looks like this doesn't quite match the old code in narrowInterleaveGroups, i.e.

  ... = dyn_cast_or_null<VPWidenRecipe>(
        InterleaveR->getStoredValues()[0]->getDefiningRecipe());

Does it matter? Just a bit worried it's not actually NFC and changing behaviour in a subtle way.

[fhahn]

This is regarding dyn_cast_or_null<VPWidenRecipe>(....->getDefiningRecipe()) vs dyn_cast<VPWidenRecipe>(...), right?

They should be equivalent, VPWidenRecipe::classof has an overload for VPValue *, which takes care of handling getDefiningRecipe return nullptr

[david-arm]

Where has this extra for loop come from? It wasn't in the original code.

[fhahn]

The previous code did not have 2 nested loops, but was processing R->operands() per iteration to check if all operands for a given position could be narrowed to memory operations.

The change re-orders the processing, to check if all first/second operands can be narrowed together.

Previously the code only handled wide ops with loads directly (which meant that they only had a single operand).

The movement here removes the restriction, but is still NFC, as the legality checks are just limited to one level of recursion at the moment.

I restructured the code a bit, to make this a bit clearer hopefully.

[artagnon]

LGTM, thanks!