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

[Matt Arsenault via llvm-commits]

================
@@ -27267,36 +27295,106 @@ 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();
-  if (!TLI.isOperationLegalOrCustom(BinOpcode, SubVT, LegalOperations))
+
+  // 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, and seed the source map by index.
+  EVT SubVT = N->getValueType(0);
+  if (VecVT.getSizeInBits() <= SubVT.getSizeInBits())
     return SDValue();
 
-  SDValue Sub0 = getSubVectorSrc(Bop0, Index, SubVT);
-  SDValue Sub1 = getSubVectorSrc(Bop1, Index, SubVT);
+  SmallMapVector<unsigned, std::tuple<SDNode *, SDValue, SDValue>, 4> Extracts;
+  // Scan each wide operand's chain once, filling the seeded indices' sources.
+  collectSubVectorSrcs(Bop0, SubVT, /*OpNo=*/0, Extracts);
+  collectSubVectorSrcs(Bop1, SubVT, /*OpNo=*/1, Extracts);
+
+  bool HasNonZeroExt = false;
+  bool HasNonExtUser = false;
+  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);
+    auto It = Extracts.find(Idx);
+    if (It == Extracts.end() || !std::get<1>(It->second) ||
+        !std::get<2>(It->second))
+      return SDValue();
+    SDNode *&ExtSubVec = std::get<0>(It->second);
+    if (!ExtSubVec) {
+      ExtSubVec = User;
+      if (Idx != 0)
+        HasNonZeroExt = true;
+    }
+  }
 
-  // TODO: We could handle the case where only 1 operand is being inserted by
-  //       creating an extract of the other operand, but that requires checking
-  //       number of uses and/or costs.
-  if (!Sub0 || !Sub1)
+  if (!TLI.isOperationLegalOrCustom(BinOpcode, SubVT, LegalOperations))
     return SDValue();
 
-  // We are inserting both operands of the wide binop only to extract back
-  // to the narrow vector size. Eliminate all of the insert/extract:
-  // ext (binop (ins ?, X, Index), (ins ?, Y, Index)), Index --> binop X, Y
-  return DAG.getNode(BinOpcode, DL, SubVT, Sub0, Sub1, BinOp->getFlags());
+  // Narrow for [SubVT, 0/undef,...,0/undef]: when the wide binop also has a
+  // non-extract user it survives, so narrowing only pays off if it folds for
+  // free to concat(narrow binop, 0/undef, ...). That holds when the sole
+  // extract is lane 0 (HasNonZeroExt rejects multi-real-lane cases such as
+  // [a,b,0,0,c,d]) and every other lane is 0/undef in both operands; otherwise
+  // a lane carries real data the wide binop must still compute. `V &&` keeps
+  // IsZeroOrUndef null-safe for lanes that were never sourced.
+  auto IsZeroOrUndef = [](SDValue V) {
+    return V && (V.isUndef() || ISD::isBuildVectorAllZeros(V.getNode()) ||
+                 isNullOrNullSplat(V));
+  };
----------------
arsenm wrote:

There should be an existing method to check is constant 0 or undef 

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