[llvm] [SelectionDAG] Improve wide integer squaring (PR #226403)

[Craig Topper via llvm-commits]

================
@@ -12634,6 +12634,141 @@ void TargetLowering::forceExpandMultiply(SelectionDAG &DAG, const SDLoc &dl,
   }
 }
 
+SDValue TargetLowering::expandWideSquare(SDNode *N, SelectionDAG &DAG) const {
+  assert(N->getOpcode() == ISD::MUL && N->getOperand(0) == N->getOperand(1) &&
+         "Expected a squaring MUL");
+  EVT VT = N->getValueType(0);
+  if (!VT.isScalarInteger())
+    return SDValue();
+  LLVMContext &Ctx = *DAG.getContext();
+  SDLoc dl(N);
+
+  // The limbs have the legal type that VT is ultimately expanded to.
+  EVT LimbVT = getTypeToExpandTo(Ctx, VT);
+  unsigned U = LimbVT.getSizeInBits();
+  unsigned Bits = VT.getSizeInBits();
+  if (Bits < 4 * U)
+    return SDValue();
+  unsigned R = Bits / U; // number of result limbs
+
+  bool HasUMUL_LOHI = isOperationLegalOrCustom(ISD::UMUL_LOHI, LimbVT);
+  bool HasMULHU = isOperationLegalOrCustom(ISD::MULHU, LimbVT);
+  if (!HasUMUL_LOHI && !HasMULHU)
+    return SDValue();
+
+  // X is usually zero-extended (the way to request the full product), so
+  // square it at its real width: only the limbs that may be nonzero take part.
+  SDValue X = N->getOperand(0);
+  unsigned NumLimbs =
+      divideCeil(Bits - DAG.computeKnownBits(X).countMinLeadingZeros(), U);
----------------
topperc wrote:

Use DAG.computeKnownBits(X).countMaxActiveBits()?

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