[llvm] [SelectionDAG][x86] Ensure vector reduction optimization (PR #144231)

[Phoebe Wang via llvm-commits]

================
@@ -47873,26 +47892,29 @@ static SDValue combineExtractVectorElt(SDNode *N, SelectionDAG &DAG,
     return DAG.getNode(X86ISD::MMX_MOVD2W, dl, MVT::i32,
                        InputVector.getOperand(0));
 
-  // Check whether this extract is the root of a sum of absolute differences
-  // pattern. This has to be done here because we really want it to happen
-  // pre-legalization,
-  if (SDValue SAD = combineBasicSADPattern(N, DAG, Subtarget))
-    return SAD;
-
-  if (SDValue VPDPBUSD = combineVPDPBUSDPattern(N, DAG, Subtarget))
-    return VPDPBUSD;
-
-  // Attempt to replace an all_of/any_of horizontal reduction with a MOVMSK.
-  if (SDValue Cmp = combinePredicateReduction(N, DAG, Subtarget))
-    return Cmp;
-
-  // Attempt to replace min/max v8i16/v16i8 reductions with PHMINPOSUW.
-  if (SDValue MinMax = combineMinMaxReduction(N, DAG, Subtarget))
-    return MinMax;
-
-  // Attempt to optimize ADD/FADD/MUL reductions with HADD, promotion etc..
-  if (SDValue V = combineArithReduction(N, DAG, Subtarget))
-    return V;
+  SDValue BinOpReduction;
+  bool IsPartialReduction = false;
+  if (auto Result = std::tie(BinOpReduction, IsPartialReduction);
+      (BinOpReduction = combineBasicSADPattern(N, DAG, Subtarget)) ||
+      (BinOpReduction = combineVPDPBUSDPattern(N, DAG, Subtarget)) ||
+      (BinOpReduction = combinePredicateReduction(N, DAG, Subtarget)) ||
+      ((Result = combineMinMaxReduction(N, DAG, Subtarget)), BinOpReduction) ||
+      ((Result = combineArithReduction(N, DAG, Subtarget)), BinOpReduction)) {
+    SDValue ExtractEltOperand = N->getOperand(0);
+    DCI.CombineTo(N, BinOpReduction);
+
+    if (!IsPartialReduction) {
+      // Replace also ExtractEltOperand.
+      // This is safe to do, because N resulted directly from a full reduction,
+      // which means all the elements are undefined except for the 0th element.
+      SDValue V =
+          DAG.getNode(ISD::SCALAR_TO_VECTOR, SDLoc(BinOpReduction),
+                      ExtractEltOperand->getValueType(0), BinOpReduction);
+      DCI.CombineTo(ExtractEltOperand.getNode(), V);
+    }
+
+    return SDValue(N, 0); // Return N so it doesn't get rechecked!
----------------
phoebewang wrote:

Can it be:
```
if (IsPartialReduction)
  return BinOpReduction;

...
```

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