[llvm] [VectorCombine] Add support for zext/sext/trunc to shuffleToIdentity (PR #92696)

[Ramkumar Ramachandra via llvm-commits]

https://github.com/artagnon commented:

Hi,

I did a partial review of the existing code, and my suggestions are along the lines of the following diff. I think it would be nice if you could improve the existing code as an NFC patch before landing this patch.

```diff
diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 8573a8adf53b..52830d7ed462 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1686,7 +1686,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
       int M = SV->getMaskValue(Lane);
       if (M < 0)
         return {nullptr, PoisonMaskElem};
-      else if (M < (int)NumElts) {
+      if (static_cast<unsigned>(M) < NumElts) {
         V = SV->getOperand(0);
         Lane = M;
       } else {
@@ -1700,12 +1700,12 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
   auto GenerateInstLaneVectorFromOperand =
       [&LookThroughShuffles](ArrayRef<InstLane> Item, int Op) {
         SmallVector<InstLane> NItem;
-        for (InstLane V : Item) {
-          NItem.emplace_back(
-              !V.first
-                  ? InstLane{nullptr, PoisonMaskElem}
-                  : LookThroughShuffles(
-                        cast<Instruction>(V.first)->getOperand(Op), V.second));
+        for (InstLane IL : Item) {
+          auto [V, Lane] = IL;
+          // The caller is responsible for ensuring that this is always an Instruction.
+          auto *I = cast<Instruction>(V);
+          InstLane ToInsert = V ? LookThroughShuffles(I->getOperand(Op), Lane) : IL;
+          NItem.emplace_back(ToInsert);
         }
         return NItem;
       };
@@ -1721,57 +1721,66 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {

   while (!Worklist.empty()) {
     SmallVector<InstLane> Item = Worklist.pop_back_val();
-    if (++NumVisited > MaxInstrsToScan)
+    if (NumVisited++ > MaxInstrsToScan)
       return false;

+    auto [FrontV, FrontLane] = Item.front();
+
     // If we found an undef first lane then bail out to keep things simple.
-    if (!Item[0].first)
+    if (!FrontV)
       return false;

     // Look for an identity value.
-    if (Item[0].second == 0 && Item[0].first->getType() == Ty &&
-        all_of(drop_begin(enumerate(Item)), [&](const auto &E) {
-          return !E.value().first || (E.value().first == Item[0].first &&
-                                      E.value().second == (int)E.index());
+    if (!FrontLane && FrontV->getType() == Ty &&
+        all_of(drop_begin(enumerate(Item)), [Item](const auto &I) {
+          auto [FrontV, _] = Item.front();
+          auto [Idx, E] = I;
+          auto [V, Lane] = E;
+          return !V || (V == FrontV && Lane == (int)Idx);
         })) {
-      IdentityLeafs.insert(Item[0].first);
+      IdentityLeafs.insert(FrontV);
       continue;
     }
+
     // Look for a splat value.
-    if (all_of(drop_begin(Item), [&](InstLane &IL) {
-          return !IL.first ||
-                 (IL.first == Item[0].first && IL.second == Item[0].second);
+    if (all_of(drop_begin(Item), [Item](InstLane &IL) {
+          auto [FrontV, FrontLane] = Item.front();
+          auto [V, Lane] = IL;
+          return !V ||
+                 (V == FrontV && Lane == FrontLane);
         })) {
-      SplatLeafs.insert(Item[0].first);
+      SplatLeafs.insert(FrontV);
       continue;
     }

     // We need each element to be the same type of value, and check that each
     // element has a single use.
-    if (!all_of(drop_begin(Item), [&](InstLane IL) {
-          if (!IL.first)
+    if (!all_of(drop_begin(Item), [Item](InstLane IL) {
+          auto [FrontV, _] = Item.front();
+          auto [V, _Lane] = IL;
+          if (!V)
             return true;
-          if (auto *I = dyn_cast<Instruction>(IL.first); I && !I->hasOneUse())
+          if (auto *I = dyn_cast<Instruction>(V); I && !I->hasOneUse())
             return false;
-          if (IL.first->getValueID() != Item[0].first->getValueID())
+          if (V->getValueID() != FrontV->getValueID())
             return false;
-          if (isa<CallInst>(IL.first) && !isa<IntrinsicInst>(IL.first))
+          if (isa<CallInst>(V) && !isa<IntrinsicInst>(V))
             return false;
-          auto *II = dyn_cast<IntrinsicInst>(IL.first);
+          auto *II = dyn_cast<IntrinsicInst>(V);
           return !II ||
-                 (isa<IntrinsicInst>(Item[0].first) &&
+                 (isa<IntrinsicInst>(FrontV) &&
                   II->getIntrinsicID() ==
-                      cast<IntrinsicInst>(Item[0].first)->getIntrinsicID());
+                      cast<IntrinsicInst>(FrontV)->getIntrinsicID());
         }))
       return false;

     // Check the operator is one that we support. We exclude div/rem in case
     // they hit UB from poison lanes.
-    if (isa<BinaryOperator>(Item[0].first) &&
-        !cast<BinaryOperator>(Item[0].first)->isIntDivRem()) {
+    if (isa<BinaryOperator>(FrontV) &&
+        !cast<BinaryOperator>(FrontV)->isIntDivRem()) {
       Worklist.push_back(GenerateInstLaneVectorFromOperand(Item, 0));
       Worklist.push_back(GenerateInstLaneVectorFromOperand(Item, 1));
-    } else if (isa<UnaryOperator>(Item[0].first)) {
+    } else if (isa<UnaryOperator>(FrontV)) {
       Worklist.push_back(GenerateInstLaneVectorFromOperand(Item, 0));
     } else {
       return false;
@@ -1790,7 +1799,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
       return Item[0].first;
     }
     if (SplatLeafs.contains(Item[0].first)) {
-      if (auto ILI = dyn_cast<Instruction>(Item[0].first))
+      if (auto *ILI = dyn_cast<Instruction>(Item[0].first))
         Builder.SetInsertPoint(*ILI->getInsertionPointAfterDef());
       else if (isa<Argument>(Item[0].first))
         Builder.SetInsertPointPastAllocas(I.getParent()->getParent());
@@ -1803,7 +1812,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
     for (unsigned Idx = 0, E = I->getNumOperands(); Idx < E; Idx++)
       Ops[Idx] = Generate(GenerateInstLaneVectorFromOperand(Item, Idx));
     Builder.SetInsertPoint(I);
-    if (auto BI = dyn_cast<BinaryOperator>(I))
+    if (auto *BI = dyn_cast<BinaryOperator>(I))
       return Builder.CreateBinOp((Instruction::BinaryOps)BI->getOpcode(),
                                  Ops[0], Ops[1]);
     assert(isa<UnaryInstruction>(I) &&
--
2.45.1
```

https://github.com/llvm/llvm-project/pull/92696