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

[Florian Hahn via llvm-commits]

================
@@ -18602,6 +18602,30 @@ 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 (SDNode *User : Extend.getNode()->users()) {
+      if (User->getOpcode() == ISD::VECTOR_SHUFFLE &&
+          User->getOperand(0) == Extend) {
+        ArrayRef<int> Mask = cast<ShuffleVectorSDNode>(User)->getMask();
+        const int Idx = Mask[0];
+        if (Idx >= NumElements)
+          continue;
+        if (llvm::all_of(Mask, [Idx](int M) { return M == Idx; }))
+          UsedElements.set(Idx);
+        else
+          break; // early loop exit to help performance
+      }
----------------
fhahn wrote:


Early exit are generally preferred in LLVM to reduce nesting level and simplify the code (https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code). E.g. something like 
```suggestion
       if (User->getOpcode() != ISD::VECTOR_SHUFFLE ||  
          User->getOperand(0) != Extend)
        break;

      ArrayRef<int> Mask = cast<ShuffleVectorSDNode>(User)->getMask();
      const int Idx = Mask[0];
      if (Idx >= NumElements || !all_of(Mask, [Idx](int M) { return M == Idx; }))
        break;
      UsedElements.set(Idx);
```

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