[DAGCombiner] Handle type-promoted constants in UDIV exact lowering (#169949)

Builds up on the solution proposed for
#169491 and applies it for UDIV
exact as well.

Author: SavchenkoValeriy

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -6403,7 +6403,6 @@ static SDValue BuildExactUDIV(const TargetLowering &TLI, SDNode *N,
                               const SDLoc &dl, SelectionDAG &DAG,
                               SmallVectorImpl<SDNode *> &Created) {
   EVT VT = N->getValueType(0);
-  EVT SVT = VT.getScalarType();
   EVT ShVT = TLI.getShiftAmountTy(VT, DAG.getDataLayout());
   EVT ShSVT = ShVT.getScalarType();
 
@@ -6413,6 +6412,8 @@ static SDValue BuildExactUDIV(const TargetLowering &TLI, SDNode *N,
   auto BuildUDIVPattern = [&](ConstantSDNode *C) {
     if (C->isZero())
       return false;
+
+    EVT CT = C->getValueType(0);
     APInt Divisor = C->getAPIntValue();
     unsigned Shift = Divisor.countr_zero();
     if (Shift) {
@@ -6422,14 +6423,15 @@ static SDValue BuildExactUDIV(const TargetLowering &TLI, SDNode *N,
     // Calculate the multiplicative inverse modulo BW.
     APInt Factor = Divisor.multiplicativeInverse();
     Shifts.push_back(DAG.getConstant(Shift, dl, ShSVT));
-    Factors.push_back(DAG.getConstant(Factor, dl, SVT));
+    Factors.push_back(DAG.getConstant(Factor, dl, CT));
     return true;
   };
 
   SDValue Op1 = N->getOperand(1);
 
   // Collect all magic values from the build vector.
-  if (!ISD::matchUnaryPredicate(Op1, BuildUDIVPattern))
+  if (!ISD::matchUnaryPredicate(Op1, BuildUDIVPattern, /*AllowUndefs=*/false,
+                                /*AllowTruncation=*/true))
     return SDValue();
 
   SDValue Shift, Factor;

llvm/test/CodeGen/AArch64/udiv-by-const-promoted-ops.ll

@@ -76,3 +76,24 @@ define <16 x i16> @urem_v16i16_by_255(<16 x i16> %x) {
   %rem = urem <16 x i16> %x, splat (i16 255)
   ret <16 x i16> %rem
 }
+
+define <8 x i16> @udiv_exact_v8i16_by_255(<8 x i16> %x) {
+; CHECK-LABEL: udiv_exact_v8i16_by_255:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mvni v1.8h, #1, lsl #8
+; CHECK-NEXT:    mul v0.8h, v0.8h, v1.8h
+; CHECK-NEXT:    ret
+  %div = udiv exact <8 x i16> %x, splat (i16 255)
+  ret <8 x i16> %div
+}
+
+define <16 x i16> @udiv_exact_v16i16_by_255(<16 x i16> %x) {
+; CHECK-LABEL: udiv_exact_v16i16_by_255:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mvni v2.8h, #1, lsl #8
+; CHECK-NEXT:    mul v0.8h, v0.8h, v2.8h
+; CHECK-NEXT:    mul v1.8h, v1.8h, v2.8h
+; CHECK-NEXT:    ret
+  %div = udiv exact <16 x i16> %x, splat (i16 255)
+  ret <16 x i16> %div
+}