Adjust partial reduction type legality check

Author: NickGuy-Arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1586,23 +1586,41 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     }
 
     if (EnablePartialReduceNodes) {
+      // Mark known legal pairs as 'Legal' (these will expand to UDOT or SDOT).
+      // Other pairs will default to 'Expand'.
+      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i16, Legal);
+      setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i8, Legal);
+
+      auto CanSplitToLegalPartialReduce = [&](MVT AccTy, MVT InTy) {
+        while (true) {
+          switch (getTypeAction(AccTy)) {
+          case TargetLoweringBase::TypeLegal:
+            return isPartialReduceMLALegalOrCustom(AccTy, InTy);
+          case TargetLoweringBase::TypeSplitVector:
+            if (!InTy.getVectorElementCount().isKnownEven())
+              return false;
+            // Currently, we only implement spillting for partial reductions,
+            // which splits the result and both operands.
+            AccTy = AccTy.getHalfNumVectorElementsVT();
+            InTy = InTy.getHalfNumVectorElementsVT();
+            break;
+          default:
+            // Assume all other type pairs are expanded.
+            return false;
+          }
+        }
+      };
+
       for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
         for (MVT InnerVT : MVT::integer_scalable_vector_valuetypes()) {
-          // 1. Set all combinations where a type is illegal to "Legal"
-          // - These will be legalized to a legal type pair
-          // - Avoid expanding them too early (or preventing folds)
+          // Mark illegal type pairs that split to a legal pair as legal.
+          // This is needed as otherwise we may not apply useful combines.
           if (!isTypeLegal(VT) || !isTypeLegal(InnerVT)) {
-            setPartialReduceMLAAction(VT, InnerVT, Legal);
-            continue;
+            if (CanSplitToLegalPartialReduce(VT, InnerVT))
+              setPartialReduceMLAAction(VT, InnerVT, Legal);
           }
-          //  2. Set all legal combinations to "Expand"
-          // - Not all of these can be lowered (via a Legal or Custom lowering).
-          setPartialReduceMLAAction(VT, InnerVT, Expand);
         }
       }
-      // 3. Mark known legal pairs as 'Legal' (these will expand to USDOT).
-      setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i16, Legal);
-      setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i8, Legal);
     }
 
     // Firstly, exclude all scalable vector extending loads/truncating stores,

llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll

@@ -206,46 +206,59 @@ define <vscale x 4 x i64> @udot_8to64(<vscale x 4 x i64> %acc, <vscale x 16 x i8
 ;
 ; CHECK-NEWLOWERING-LABEL: udot_8to64:
 ; CHECK-NEWLOWERING:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-NEXT:    uunpklo z4.h, z3.b
-; CHECK-NEWLOWERING-NEXT:    uunpklo z5.h, z2.b
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z3.h, z3.b
+; CHECK-NEWLOWERING-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEWLOWERING-NEXT:    addvl sp, sp, #-2
+; CHECK-NEWLOWERING-NEXT:    str z9, [sp] // 16-byte Folded Spill
+; CHECK-NEWLOWERING-NEXT:    str z8, [sp, #1, mul vl] // 16-byte Folded Spill
+; CHECK-NEWLOWERING-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x10, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 16 * VG
+; CHECK-NEWLOWERING-NEXT:    .cfi_offset w29, -16
+; CHECK-NEWLOWERING-NEXT:    .cfi_escape 0x10, 0x48, 0x0a, 0x11, 0x70, 0x22, 0x11, 0x78, 0x92, 0x2e, 0x00, 0x1e, 0x22 // $d8 @ cfa - 16 - 8 * VG
+; CHECK-NEWLOWERING-NEXT:    .cfi_escape 0x10, 0x49, 0x0a, 0x11, 0x70, 0x22, 0x11, 0x70, 0x92, 0x2e, 0x00, 0x1e, 0x22 // $d9 @ cfa - 16 - 16 * VG
+; CHECK-NEWLOWERING-NEXT:    uunpklo z4.h, z2.b
+; CHECK-NEWLOWERING-NEXT:    uunpklo z5.h, z3.b
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.h, z2.b
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z3.h, z3.b
 ; CHECK-NEWLOWERING-NEXT:    ptrue p0.d
 ; CHECK-NEWLOWERING-NEXT:    uunpklo z6.s, z4.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z7.s, z5.h
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z4.s, z4.h
+; CHECK-NEWLOWERING-NEXT:    uunpklo z7.s, z5.h
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z5.s, z5.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z24.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z25.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z3.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z2.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z26.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z27.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z6.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z7.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z28.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z24.s, z2.h
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.s, z2.h
+; CHECK-NEWLOWERING-NEXT:    uunpklo z25.s, z3.h
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z3.s, z3.h
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z26.d, z6.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z6.d, z6.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z27.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z28.d, z7.s
 ; CHECK-NEWLOWERING-NEXT:    uunpklo z29.d, z5.s
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z4.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z7.d, z7.s
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z5.d, z5.s
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z27.d, z26.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z26.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z24.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z7.d, z6.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z6.d, z25.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z7.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    mul z27.d, z29.d, z28.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z28.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z25.d, z25.s
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z30.d, z24.s
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z31.d, z2.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z24.d, z24.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z2.d, z2.s
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z8.d, z25.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z25.d, z25.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z9.d, z3.s
+; CHECK-NEWLOWERING-NEXT:    mul z27.d, z27.d, z29.d
+; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z6.d, z28.d
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z3.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    mul z4.d, z5.d, z4.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z6.d, z26.d
-; CHECK-NEWLOWERING-NEXT:    movprfx z5, z27
-; CHECK-NEWLOWERING-NEXT:    mla z5.d, p0/m, z28.d, z7.d
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z25.d, z24.d
-; CHECK-NEWLOWERING-NEXT:    mad z2.d, p0/m, z3.d, z4.d
-; CHECK-NEWLOWERING-NEXT:    add z0.d, z5.d, z0.d
-; CHECK-NEWLOWERING-NEXT:    add z1.d, z2.d, z1.d
+; CHECK-NEWLOWERING-NEXT:    mul z4.d, z4.d, z5.d
+; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z26.d, z7.d
+; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z2.d, z9.d
+; CHECK-NEWLOWERING-NEXT:    movprfx z2, z27
+; CHECK-NEWLOWERING-NEXT:    mla z2.d, p0/m, z24.d, z25.d
+; CHECK-NEWLOWERING-NEXT:    ldr z9, [sp] // 16-byte Folded Reload
+; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z31.d, z3.d
+; CHECK-NEWLOWERING-NEXT:    movprfx z3, z4
+; CHECK-NEWLOWERING-NEXT:    mla z3.d, p0/m, z30.d, z8.d
+; CHECK-NEWLOWERING-NEXT:    ldr z8, [sp, #1, mul vl] // 16-byte Folded Reload
+; CHECK-NEWLOWERING-NEXT:    add z0.d, z2.d, z0.d
+; CHECK-NEWLOWERING-NEXT:    add z1.d, z3.d, z1.d
+; CHECK-NEWLOWERING-NEXT:    addvl sp, sp, #2
+; CHECK-NEWLOWERING-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
   %a.wide = zext <vscale x 16 x i8> %a to <vscale x 16 x i64>
@@ -269,46 +282,59 @@ define <vscale x 4 x i64> @sdot_8to64(<vscale x 4 x i64> %acc, <vscale x 16 x i8
 ;
 ; CHECK-NEWLOWERING-LABEL: sdot_8to64:
 ; CHECK-NEWLOWERING:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-NEXT:    sunpklo z4.h, z3.b
-; CHECK-NEWLOWERING-NEXT:    sunpklo z5.h, z2.b
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z3.h, z3.b
+; CHECK-NEWLOWERING-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEWLOWERING-NEXT:    addvl sp, sp, #-2
+; CHECK-NEWLOWERING-NEXT:    str z9, [sp] // 16-byte Folded Spill
+; CHECK-NEWLOWERING-NEXT:    str z8, [sp, #1, mul vl] // 16-byte Folded Spill
+; CHECK-NEWLOWERING-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x10, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 16 * VG
+; CHECK-NEWLOWERING-NEXT:    .cfi_offset w29, -16
+; CHECK-NEWLOWERING-NEXT:    .cfi_escape 0x10, 0x48, 0x0a, 0x11, 0x70, 0x22, 0x11, 0x78, 0x92, 0x2e, 0x00, 0x1e, 0x22 // $d8 @ cfa - 16 - 8 * VG
+; CHECK-NEWLOWERING-NEXT:    .cfi_escape 0x10, 0x49, 0x0a, 0x11, 0x70, 0x22, 0x11, 0x70, 0x92, 0x2e, 0x00, 0x1e, 0x22 // $d9 @ cfa - 16 - 16 * VG
+; CHECK-NEWLOWERING-NEXT:    sunpklo z4.h, z2.b
+; CHECK-NEWLOWERING-NEXT:    sunpklo z5.h, z3.b
 ; CHECK-NEWLOWERING-NEXT:    sunpkhi z2.h, z2.b
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z3.h, z3.b
 ; CHECK-NEWLOWERING-NEXT:    ptrue p0.d
 ; CHECK-NEWLOWERING-NEXT:    sunpklo z6.s, z4.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z7.s, z5.h
 ; CHECK-NEWLOWERING-NEXT:    sunpkhi z4.s, z4.h
+; CHECK-NEWLOWERING-NEXT:    sunpklo z7.s, z5.h
 ; CHECK-NEWLOWERING-NEXT:    sunpkhi z5.s, z5.h
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z24.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z25.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z3.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z2.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z26.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z27.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z6.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z7.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z28.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z24.s, z2.h
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z2.s, z2.h
+; CHECK-NEWLOWERING-NEXT:    sunpklo z25.s, z3.h
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z3.s, z3.h
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z26.d, z6.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z6.d, z6.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z27.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z28.d, z7.s
 ; CHECK-NEWLOWERING-NEXT:    sunpklo z29.d, z5.s
 ; CHECK-NEWLOWERING-NEXT:    sunpkhi z4.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z7.d, z7.s
 ; CHECK-NEWLOWERING-NEXT:    sunpkhi z5.d, z5.s
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z27.d, z26.d
-; CHECK-NEWLOWERING-NEXT:    sunpklo z26.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z24.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z7.d, z6.d
-; CHECK-NEWLOWERING-NEXT:    sunpklo z6.d, z25.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z7.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    mul z27.d, z29.d, z28.d
-; CHECK-NEWLOWERING-NEXT:    sunpklo z28.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z25.d, z25.s
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z30.d, z24.s
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z31.d, z2.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z24.d, z24.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z2.d, z2.s
+; CHECK-NEWLOWERING-NEXT:    sunpkhi z8.d, z25.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z25.d, z25.s
+; CHECK-NEWLOWERING-NEXT:    sunpklo z9.d, z3.s
+; CHECK-NEWLOWERING-NEXT:    mul z27.d, z27.d, z29.d
+; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z6.d, z28.d
 ; CHECK-NEWLOWERING-NEXT:    sunpkhi z3.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z2.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    mul z4.d, z5.d, z4.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z6.d, z26.d
-; CHECK-NEWLOWERING-NEXT:    movprfx z5, z27
-; CHECK-NEWLOWERING-NEXT:    mla z5.d, p0/m, z28.d, z7.d
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z25.d, z24.d
-; CHECK-NEWLOWERING-NEXT:    mad z2.d, p0/m, z3.d, z4.d
-; CHECK-NEWLOWERING-NEXT:    add z0.d, z5.d, z0.d
-; CHECK-NEWLOWERING-NEXT:    add z1.d, z2.d, z1.d
+; CHECK-NEWLOWERING-NEXT:    mul z4.d, z4.d, z5.d
+; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z26.d, z7.d
+; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z2.d, z9.d
+; CHECK-NEWLOWERING-NEXT:    movprfx z2, z27
+; CHECK-NEWLOWERING-NEXT:    mla z2.d, p0/m, z24.d, z25.d
+; CHECK-NEWLOWERING-NEXT:    ldr z9, [sp] // 16-byte Folded Reload
+; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z31.d, z3.d
+; CHECK-NEWLOWERING-NEXT:    movprfx z3, z4
+; CHECK-NEWLOWERING-NEXT:    mla z3.d, p0/m, z30.d, z8.d
+; CHECK-NEWLOWERING-NEXT:    ldr z8, [sp, #1, mul vl] // 16-byte Folded Reload
+; CHECK-NEWLOWERING-NEXT:    add z0.d, z2.d, z0.d
+; CHECK-NEWLOWERING-NEXT:    add z1.d, z3.d, z1.d
+; CHECK-NEWLOWERING-NEXT:    addvl sp, sp, #2
+; CHECK-NEWLOWERING-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
   %a.wide = sext <vscale x 16 x i8> %a to <vscale x 16 x i64>
@@ -790,11 +816,11 @@ define <vscale x 4 x i32> @not_udot(<vscale x 4 x i32> %acc, <vscale x 8 x i8> %
 ; CHECK-NEWLOWERING-NEXT:    and z1.h, z1.h, #0xff
 ; CHECK-NEWLOWERING-NEXT:    and z2.h, z2.h, #0xff
 ; CHECK-NEWLOWERING-NEXT:    ptrue p0.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z3.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z4.s, z1.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.s, z2.h
+; CHECK-NEWLOWERING-NEXT:    uunpklo z3.s, z1.h
+; CHECK-NEWLOWERING-NEXT:    uunpklo z4.s, z2.h
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z1.s, z1.h
-; CHECK-NEWLOWERING-NEXT:    mla z0.s, p0/m, z4.s, z3.s
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.s, z2.h
+; CHECK-NEWLOWERING-NEXT:    mla z0.s, p0/m, z3.s, z4.s
 ; CHECK-NEWLOWERING-NEXT:    mla z0.s, p0/m, z1.s, z2.s
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
@@ -824,11 +850,11 @@ define <vscale x 2 x i64> @not_udot_wide(<vscale x 2 x i64> %acc, <vscale x 4 x
 ; CHECK-NEWLOWERING-NEXT:    and z1.s, z1.s, #0xffff
 ; CHECK-NEWLOWERING-NEXT:    and z2.s, z2.s, #0xffff
 ; CHECK-NEWLOWERING-NEXT:    ptrue p0.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z3.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z4.d, z1.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.d, z2.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z3.d, z1.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z4.d, z2.s
 ; CHECK-NEWLOWERING-NEXT:    uunpkhi z1.d, z1.s
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z4.d, z3.d
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.d, z2.s
+; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z3.d, z4.d
 ; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z1.d, z2.d
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
@@ -1195,8 +1221,10 @@ define <vscale x 2 x i16> @udot_nxv8i8_promote (<vscale x 2 x i16> %acc, <vscale
 ;
 ; CHECK-NEWLOWERING-LABEL: udot_nxv8i8_promote:
 ; CHECK-NEWLOWERING:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-NEXT:    and z2.h, z2.h, #0xff
 ; CHECK-NEWLOWERING-NEXT:    and z1.h, z1.h, #0xff
+; CHECK-NEWLOWERING-NEXT:    and z2.h, z2.h, #0xff
+; CHECK-NEWLOWERING-NEXT:    mul z1.h, z1.h, z2.h
+; CHECK-NEWLOWERING-NEXT:    mov z2.h, #1 // =0x1
 ; CHECK-NEWLOWERING-NEXT:    udot z0.d, z1.h, z2.h
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
@@ -1229,9 +1257,11 @@ define <vscale x 2 x i16> @sdot_nxv8i8_promote (<vscale x 2 x i16> %acc, <vscale
 ; CHECK-NEWLOWERING-LABEL: sdot_nxv8i8_promote:
 ; CHECK-NEWLOWERING:       // %bb.0: // %entry
 ; CHECK-NEWLOWERING-NEXT:    ptrue p0.h
-; CHECK-NEWLOWERING-NEXT:    sxtb z2.h, p0/m, z2.h
 ; CHECK-NEWLOWERING-NEXT:    sxtb z1.h, p0/m, z1.h
-; CHECK-NEWLOWERING-NEXT:    sdot z0.d, z1.h, z2.h
+; CHECK-NEWLOWERING-NEXT:    sxtb z2.h, p0/m, z2.h
+; CHECK-NEWLOWERING-NEXT:    mul z1.h, z1.h, z2.h
+; CHECK-NEWLOWERING-NEXT:    mov z2.h, #1 // =0x1
+; CHECK-NEWLOWERING-NEXT:    udot z0.d, z1.h, z2.h
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
   %a.wide = sext <vscale x 8 x i8> %a to <vscale x 8 x i16>