[llvm] [LV] Only fold outer extend into narrow mul if it cannot wrap. (PR #229230)

[Sander de Smalen via llvm-commits]

================
@@ -3789,6 +3789,24 @@ tryToMatchAndCreateExtendedReduction(VPReductionRecipe *Red, VPCostContext &Ctx,
   return nullptr;
 }
 
+using ExtendKind = TTI::PartialReductionExtendKind;
+
+static ExtendKind getPartialReductionExtendKind(VPWidenCastRecipe *Cast) {
+  return TTI::getPartialReductionExtendKind(Cast->getOpcode());
+}
+
+/// Returns true if the narrow \p Mul in ext(mul(ext(A), ext(B))) or
+/// ext(mul(ext(A), C)) cannot wrap.
+static bool canFoldOuterExtendIntoMul(ExtendKind ExtKind, VPWidenRecipe *Mul,
+                                      Type *LHSSrcTy, Type *RHSSrcTy) {
+  if (Mul->getScalarType()->getScalarSizeInBits() >=
+      LHSSrcTy->getScalarSizeInBits() + RHSSrcTy->getScalarSizeInBits())
+    return true;
+  if (ExtKind == ExtendKind::PR_ZeroExtend)
+    return Mul->hasNoUnsignedWrap() || Mul->hasNoSignedWrap();
----------------
sdesmalen-arm wrote:

This comment doesn't really explain _why_ either NUW or NSW is sufficient, it just describes what the two lines below already express. I was thinking more about something like:

```
// When the source values are zero-extended, the values are unsigned. When an unsigned
// value doesn't wrap in the signed domain it also can't wrap in the unsigned domain.
```

That makes me wonder why the `mul` doesn't already have the `nuw` flag, and actually think that's just a quirk/omission from one of the tests, since instcombine seems to add `nuw`: https://godbolt.org/z/ah1vefc1c

So, I think you can just remove the `|| Mul->hasNoSignedWrap()` if you manually add `nuw` to the LV test.

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