[llvm] [InstCombine, AArch64] Avoid vector Ext in case by-element operation variant apply for all elements (PR #140733)

[Ahmad Yasin via llvm-commits]

================
@@ -18602,6 +18602,33 @@ static SDValue performBuildShuffleExtendCombine(SDValue BV, SelectionDAG &DAG) {
     SeenZExtOrSExt = true;
   }
 
+  // Avoid the said use of vector SExt/ZExt in case all vector elements are
+  // consumed and each shuffle's mask uses same index, in order to permit use of
+  // indexed OP (e.g. MLA, MUL) variants
+  EVT ExtendType = Extend->getValueType(0);
+  if (ExtendType.isVector() && !ExtendType.isScalableVT()) {
+    const int NumElements = ExtendType.getVectorNumElements();
+    SmallBitVector UsedElements(NumElements, false);
+    for (auto UI = Extend.getNode()->use_begin(),
+              UE = Extend.getNode()->use_end();
+         UI != UE; ++UI) {
+      SDNode *User = UI->getUser();
+      if (User->getOpcode() == ISD::VECTOR_SHUFFLE &&
+          User->getOperand(0) == Extend) {
----------------
ayasin-a wrote:

The loop is looking for all shuffles whose 1st operand is the Extend SDValue, then we do exist early (through the `break`) on first non-splat shuffle.
Once we traversed all such shuffles, we allow the optimization if the splats cover all indices of the Extend vector. For example, we observe splats to indices 0, 1, 2 and 3 for a vector of 4 elements.
Not sure if possible to inverts the condition in the said if-statement..


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