[llvm] [AMDGPU][X86][DAG] Avoid duplicate BinOp result from narrowing insert-extract sub-vector (PR #201056)

[Matt Arsenault via llvm-commits]

================
@@ -27269,36 +27301,113 @@ static SDValue narrowInsertExtractVectorBinOp(EVT SubVT, SDValue BinOp,
 
   EVT VecVT = BinOp.getValueType();
   SDValue Bop0 = BinOp.getOperand(0), Bop1 = BinOp.getOperand(1);
-  if (VecVT != Bop0.getValueType() || VecVT != Bop1.getValueType())
+  if (VecVT != Bop0.getValueType() || VecVT != Bop1.getValueType() ||
+      VecVT.isScalableVT())
     return SDValue();
+
+  // This fold only pays off when the wide binop disappears completely, so every
+  // user must be an extract_subvector. Require them all to extract N's type so
+  // a single chain scan serves every extract.
+  EVT SubVT = N->getValueType(0);
+  if (VecVT.getSizeInBits() <= SubVT.getSizeInBits())
+    return SDValue();
+
   if (!TLI.isOperationLegalOrCustom(BinOpcode, SubVT, LegalOperations))
     return SDValue();
 
-  SDValue Sub0 = getSubVectorSrc(Bop0, Index, SubVT);
-  SDValue Sub1 = getSubVectorSrc(Bop1, Index, SubVT);
+  // The wide binop splits into a small, dense set of SubVT-sized slots, so a
+  // flat table indexed by subvector position is all we need.
+  unsigned NumSubElts = SubVT.getVectorMinNumElements();
+  unsigned NumParts = VecVT.getVectorMinNumElements() / NumSubElts;
+  SmallVector<std::tuple<SDNode *, SDValue, SDValue>, 4> Slots(NumParts);
+  // Scan each wide operand's chain once, filling each slot's source.
+  collectSubVectorSrcs(Bop0, SubVT, /*OpNo=*/0, Slots);
+  collectSubVectorSrcs(Bop1, SubVT, /*OpNo=*/1, Slots);
+
+  bool HasNonZeroExt = false;
+  bool HasNonExtUser = false;
+  bool AllExtractsCheap = true;
+  for (SDNode *User : BinOp->users()) {
+    if (User->getOpcode() != ISD::EXTRACT_SUBVECTOR) {
+      HasNonExtUser = true;
+      continue;
+    }
+    if (User->getValueType(0) != SubVT ||
+        (User->getCombinerWorklistIndex() < 0 && User != N))
+      return SDValue();
+    unsigned Idx = User->getConstantOperandVal(1);
+    if (Idx % NumSubElts != 0 || Idx / NumSubElts >= NumParts)
+      return SDValue();
+    auto &Slot = Slots[Idx / NumSubElts];
----------------
arsenm wrote:

Can use structure binding here instead of all the std::gets 

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