[AArch64] Canonicalize SIGN_EXTEND to VSELECT

[NicolaLancellotti]

Summary

This patch canonicalizes SIGN_EXTEND to VSELECT to allow sext to transform into a predicated instruction like add, sub or mul.

[NicolaLancellotti]

I know you might be busy. When you have time, can you review this patch, please?

[dmgreen]

To add some extra context @NicolaLancellotti wrote this as an extension to D135596. It requires a lot more changes though, and my feeling is that we might be playing whack-a-mole with generic combines that all assume sext's. As we had the patch I suggested putting it up to see what others thought. So.. what do others think?

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5418

I don't think this is limited to scalable vectors?

5421

Could this use isConstantSplatVectorAllOnes?

5431

Could this be part of the if condition?

[paulwalker-arm]

Sorry for the delay, but it's likely to be later in the week if not next week before I'll have enough time to review properly.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5418

indeed - I'd like to see what effect this has on other targets for fixed vectors (where we have a lot more test coverage e.g. AVX512)

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5561

DAG.getAllOnesConstant ?

18364

!SplatVal.isAllOnes()

llvm/test/CodeGen/AArch64/predicated-add-sub-mul.ll

224

please can you pre-commit these tests with current (trunk) codegen and then rebase the patch to show the codegen diff

[NicolaLancellotti]

Superseded by: https://reviews.llvm.org/D138570

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (HandOpcode == ISD::VECTOR_SHUFFLE && Level < AfterLegalizeDAG) {
     // (logic_op (shuf (C, A), shuf (C, B))) --> shuf (C, logic_op (A, B))
     if (N0.getOperand(0) == N1.getOperand(0) && ShOp.getNode()) {
       SDValue Logic = DAG.getNode(LogicOpcode, DL, VT, N0.getOperand(1),
                                   N1.getOperand(1));
       return DAG.getVectorShuffle(VT, DL, ShOp, Logic, SVN0->getMask());
     }
   }
 
+  // logic_op (vselect (A, splat_vector(-1), splat_vector(0)),
+  //           vselect (B, splat_vector(-1), splat_vector(0))) -->
+  // vselect (logic_op (A, B), splat_vector(-1), splat_vector(0))
+  // where A and B are i1 scalar vectors
+  APInt SplatVal;
+  if (HandOpcode == ISD::VSELECT &&
+      N0->getOperand(0)->getValueType(0).isScalableVector() &&
+      N0->getOperand(0)->getValueType(0).getScalarType() == MVT::i1 &&
+      ISD::isConstantSplatVectorAllZeros(N0->getOperand(2).getNode()) &&
+      ISD::isConstantSplatVector(N0->getOperand(1).getNode(), SplatVal) &&
+      SplatVal.getSExtValue() == -1 &&
+      N1->getOperand(0)->getValueType(0).isScalableVector() &&
+      N1->getOperand(0)->getValueType(0).getScalarType() == MVT::i1 &&
+      ISD::isConstantSplatVectorAllZeros(N1->getOperand(2).getNode()) &&
+      ISD::isConstantSplatVector(N1->getOperand(1).getNode(), SplatVal) &&
+      SplatVal.getSExtValue() == -1) {
+
+    // If both operands have other uses, this transform would create extra
+    // instructions without eliminating anything.
+    if (!N0.hasOneUse() && !N1.hasOneUse())
+      return SDValue();
+
+    SDValue Logic = DAG.getNode(LogicOpcode, DL, XVT, X, Y);
+    return DAG.getNode(ISD::VSELECT, DL, VT, Logic, N0.getOperand(1),
+                       N0.getOperand(2));
+  }
+
   return SDValue();
 }
 
 /// Try to make (and/or setcc (LL, LR), setcc (RL, RR)) more efficient.
 SDValue DAGCombiner::foldLogicOfSetCCs(bool IsAnd, SDValue N0, SDValue N1,
                                        const SDLoc &DL) {
   SDValue LL, LR, RL, RR, N0CC, N1CC;
   if (!isSetCCEquivalent(N0, LL, LR, N0CC) ||

llvm/lib/Target/AArch64/AArch64ISelLowering.h

SDValue LowerCallResult(SDValue Chain, SDValue InFlag,
                           SmallVectorImpl<SDValue> &InVals, bool isThisReturn,
                           SDValue ThisVal) const;
 
   SDValue LowerLOAD(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSTORE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerStore128(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerABS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerZERO_EXTEND(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerSIGN_EXTEND(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerMGATHER(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerMSCATTER(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerMLOAD(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue LowerINTRINSIC_W_CHAIN(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINTRINSIC_WO_CHAIN(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (auto VT : {MVT::nxv16i8, MVT::nxv8i16, MVT::nxv4i32, MVT::nxv2i64}) {
       setOperationAction(ISD::VECREDUCE_AND, VT, Custom);
       setOperationAction(ISD::VECREDUCE_OR, VT, Custom);
       setOperationAction(ISD::VECREDUCE_XOR, VT, Custom);
       setOperationAction(ISD::VECREDUCE_UMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_UMAX, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SMIN, VT, Custom);
       setOperationAction(ISD::VECREDUCE_SMAX, VT, Custom);
       setOperationAction(ISD::ZERO_EXTEND, VT, Custom);
+      setOperationAction(ISD::SIGN_EXTEND, VT, Custom);
 
       setOperationAction(ISD::UMUL_LOHI, VT, Expand);
       setOperationAction(ISD::SMUL_LOHI, VT, Expand);
       setOperationAction(ISD::SELECT_CC, VT, Expand);
       setOperationAction(ISD::ROTL, VT, Expand);
       setOperationAction(ISD::ROTR, VT, Expand);
 
       setOperationAction(ISD::SADDSAT, VT, Legal);
SDValue AArch64TargetLowering::LowerZERO_EXTEND(SDValue Op, SelectionDAG &DAG) const {
 
   SDLoc DL = SDLoc(Op);
   EVT VT = Op->getValueType(0);
   SDValue Ones = DAG.getConstant(1, DL, VT);
   SDValue Zeros = DAG.getConstant(0, DL, VT);
   return DAG.getNode(ISD::VSELECT, DL, VT, Value, Ones, Zeros);
 }
 
+SDValue AArch64TargetLowering::LowerSIGN_EXTEND(SDValue Op, SelectionDAG &DAG) const {
+  assert(Op->getOpcode() == ISD::SIGN_EXTEND && "Expected SIGN_EXTEND");
+
+  if (Op.getValueType().isFixedLengthVector())
+    return LowerFixedLengthVectorIntExtendToSVE(Op, DAG);
+
+  // Try to lower to VSELECT to allow zext to transform into
+  // a predicated instruction like add, sub or mul.
+  SDValue Value = Op->getOperand(0);
+  if (!Value->getValueType(0).isScalableVector() ||
+      Value->getValueType(0).getScalarType() != MVT::i1)
+    return SDValue();
+
+  SDLoc DL = SDLoc(Op);
+  EVT VT = Op->getValueType(0);
+  SDValue MinusOnes = DAG.getConstant(APInt(VT.getScalarType().getScalarSizeInBits(), -1, true), DL, VT);
+  SDValue Zeros = DAG.getConstant(APInt(VT.getScalarType().getScalarSizeInBits(), 0, true), DL, VT);
+  return DAG.getNode(ISD::VSELECT, DL, VT, Value, MinusOnes, Zeros);
+}
+
 SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   LLVM_DEBUG(dbgs() << "Custom lowering: ");
   LLVM_DEBUG(Op.dump());
 
   switch (Op.getOpcode()) {
   default:
     llvm_unreachable("unimplemented operand");
case ISD::ATOMIC_LOAD_SUB:
     return LowerATOMIC_LOAD_SUB(Op, DAG);
   case ISD::ATOMIC_LOAD_AND:
     return LowerATOMIC_LOAD_AND(Op, DAG);
   case ISD::DYNAMIC_STACKALLOC:
     return LowerDYNAMIC_STACKALLOC(Op, DAG);
   case ISD::VSCALE:
     return LowerVSCALE(Op, DAG);
   case ISD::ANY_EXTEND:
-  case ISD::SIGN_EXTEND:
     return LowerFixedLengthVectorIntExtendToSVE(Op, DAG);
+  case ISD::SIGN_EXTEND:
+    return LowerSIGN_EXTEND(Op, DAG);
   case ISD::ZERO_EXTEND:
     return LowerZERO_EXTEND(Op, DAG);
   case ISD::SIGN_EXTEND_INREG: {
     // Only custom lower when ExtraVT has a legal byte based element type.
     EVT ExtraVT = cast<VTSDNode>(Op.getOperand(1))->getVT();
     EVT ExtraEltVT = ExtraVT.getVectorElementType();
     if ((ExtraEltVT != MVT::i8) && (ExtraEltVT != MVT::i16) &&
         (ExtraEltVT != MVT::i32) && (ExtraEltVT != MVT::i64))
static SDValue performVectorShiftCombine(SDNode *N,
   APInt DemandedMask = ~ShiftedOutBits;
 
   if (TLI.SimplifyDemandedBits(Op, DemandedMask, DCI))
     return SDValue(N, 0);
 
   return SDValue();
 }
 
+static bool isVSELECTLoweredFromSIGN_EXTEND(SDValue N) {
+  APInt SplatVal;
+  if (N->getOpcode() != ISD::VSELECT ||
+      !N->getOperand(0)->getValueType(0).isScalableVector() ||
+      N->getOperand(0)->getValueType(0).getScalarType() != MVT::i1 ||
+      !ISD::isConstantSplatVectorAllZeros(N->getOperand(2).getNode()) ||
+      !ISD::isConstantSplatVector(N->getOperand(1).getNode(), SplatVal) ||
+      SplatVal.getSExtValue() != -1)
+    return false;
+  return true;
+}
+
 static SDValue performSunpkloCombine(SDNode *N, SelectionDAG &DAG) {
   // sunpklo(sext(pred)) -> sext(extract_low_half(pred))
   // This transform works in partnership with performSetCCPunpkCombine to
   // remove unnecessary transfer of predicates into standard registers and back
-  if (N->getOperand(0).getOpcode() == ISD::SIGN_EXTEND &&
+  if ((N->getOperand(0).getOpcode() == ISD::SIGN_EXTEND &&
       N->getOperand(0)->getOperand(0)->getValueType(0).getScalarType() ==
-          MVT::i1) {
+          MVT::i1) || isVSELECTLoweredFromSIGN_EXTEND(N->getOperand(0))) {
     SDValue CC = N->getOperand(0)->getOperand(0);
     auto VT = CC->getValueType(0).getHalfNumVectorElementsVT(*DAG.getContext());
     SDValue Unpk = DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), VT, CC,
                                DAG.getVectorIdxConstant(0, SDLoc(N)));
     return DAG.getNode(ISD::SIGN_EXTEND, SDLoc(N), N->getValueType(0), Unpk);
   }
 
   return SDValue();
static SDValue performSetCCPunpkCombine(SDNode *N, SelectionDAG &DAG) {
   //   (sign_extend (extract_subvector (setcc_merge_zero ... pred ...))), 0, ne
   //   => extract_subvector (inner setcc_merge_zero)
   SDValue Pred = N->getOperand(0);
   SDValue LHS = N->getOperand(1);
   SDValue RHS = N->getOperand(2);
   ISD::CondCode Cond = cast<CondCodeSDNode>(N->getOperand(3))->get();
 
   if (Cond != ISD::SETNE || !isZerosVector(RHS.getNode()) ||
-      LHS->getOpcode() != ISD::SIGN_EXTEND)
+      (LHS->getOpcode() != ISD::SIGN_EXTEND && !isVSELECTLoweredFromSIGN_EXTEND(LHS)))
     return SDValue();
 
   SDValue Extract = LHS->getOperand(0);
   if (Extract->getOpcode() != ISD::EXTRACT_SUBVECTOR ||
       Extract->getValueType(0) != N->getValueType(0) ||
       Extract->getConstantOperandVal(1) != 0)
     return SDValue();
 
performSetccMergeZeroCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
   SDValue LHS = N->getOperand(1);
   SDValue RHS = N->getOperand(2);
   ISD::CondCode Cond = cast<CondCodeSDNode>(N->getOperand(3))->get();
 
   if (SDValue V = performSetCCPunpkCombine(N, DAG))
     return V;
 
   if (Cond == ISD::SETNE && isZerosVector(RHS.getNode()) &&
-      LHS->getOpcode() == ISD::SIGN_EXTEND &&
+      (LHS->getOpcode() == ISD::SIGN_EXTEND || isVSELECTLoweredFromSIGN_EXTEND(LHS)) &&
       LHS->getOperand(0)->getValueType(0) == N->getValueType(0)) {
     //    setcc_merge_zero(
     //       pred, extend(setcc_merge_zero(pred, ...)), != splat(0))
     // => setcc_merge_zero(pred, ...)
     if (LHS->getOperand(0)->getOpcode() == AArch64ISD::SETCC_MERGE_ZERO &&
         LHS->getOperand(0)->getOperand(0) == Pred)
       return LHS->getOperand(0);
 
AArch64TargetLowering::LowerFixedLengthVectorSelectToSVE(SDValue Op,
   SDValue Op1 = convertToScalableVector(DAG, ContainerVT, Op->getOperand(1));
   SDValue Op2 = convertToScalableVector(DAG, ContainerVT, Op->getOperand(2));
 
   // Convert the mask to a predicated (NOTE: We don't need to worry about
   // inactive lanes since VSELECT is safe when given undefined elements).
   EVT MaskVT = Op.getOperand(0).getValueType();
   EVT MaskContainerVT = getContainerForFixedLengthVector(DAG, MaskVT);
   auto Mask = convertToScalableVector(DAG, MaskContainerVT, Op.getOperand(0));
+  if (isVSELECTLoweredFromSIGN_EXTEND(Mask)) {
+    Mask = Mask->getOperand(0);
+  } else {
     Mask = DAG.getNode(ISD::TRUNCATE, DL,
                        MaskContainerVT.changeVectorElementType(MVT::i1), Mask);
-
-  auto ScalableRes = DAG.getNode(ISD::VSELECT, DL, ContainerVT,
-                                Mask, Op1, Op2);
-
+  }
+  auto ScalableRes = DAG.getNode(ISD::VSELECT, DL, ContainerVT, Mask, Op1, Op2);
   return convertFromScalableVector(DAG, VT, ScalableRes);
 }
 
 SDValue AArch64TargetLowering::LowerFixedLengthVectorSetccToSVE(
     SDValue Op, SelectionDAG &DAG) const {
   SDLoc DL(Op);
   EVT InVT = Op.getOperand(0).getValueType();
   EVT ContainerVT = getContainerForFixedLengthVector(DAG, InVT);

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

ConversionTbl[] = {
     { ISD::BITCAST, MVT::nxv4i16, MVT::nxv4f16, 0 },
     { ISD::BITCAST, MVT::nxv2i32, MVT::nxv2f32, 0 },
 
     // Zero extends from nxvmi1 to nxvmiN.
     { ISD::ZERO_EXTEND, MVT::nxv2i64, MVT::nxv2i1, 1 },
     { ISD::ZERO_EXTEND, MVT::nxv4i32, MVT::nxv4i1, 1 },
     { ISD::ZERO_EXTEND, MVT::nxv8i16, MVT::nxv8i1, 1 },
     { ISD::ZERO_EXTEND, MVT::nxv16i8, MVT::nxv16i1, 1 },
+
+    // Sign extends from nxvmi1 to nxvmiN.
+    { ISD::SIGN_EXTEND, MVT::nxv2i64, MVT::nxv2i1, 1 },
+    { ISD::SIGN_EXTEND, MVT::nxv4i32, MVT::nxv4i1, 1 },
+    { ISD::SIGN_EXTEND, MVT::nxv8i16, MVT::nxv8i1, 1 },
+    { ISD::SIGN_EXTEND, MVT::nxv16i8, MVT::nxv16i1, 1 },
   };
 
   if (const auto *Entry = ConvertCostTableLookup(ConversionTbl, ISD,
                                                  DstTy.getSimpleVT(),
                                                  SrcTy.getSimpleVT()))
     return AdjustCost(Entry->Cost);
 
   static const TypeConversionCostTblEntry FP16Tbl[] = {

llvm/lib/Target/AArch64/SVEInstrFormats.td

 
 multiclass sve_int_dup_imm_pred_zero_inst<
     bits<2> sz8_64, string asm, ZPRRegOp zprty, imm8_opt_lsl cpyimm,
     ValueType intty, ValueType predty, ValueType scalarty, ComplexPattern cpx> {
   def NAME : sve_int_dup_imm_pred<sz8_64, 0, asm, zprty, "/z",
                                   (ins PPRAny:$Pg, cpyimm:$imm)>;
   def : InstAlias<"mov $Zd, $Pg/z, $imm",
                   (!cast<Instruction>(NAME) zprty:$Zd, PPRAny:$Pg, cpyimm:$imm), 1>;
-  def : Pat<(intty (sext (predty PPRAny:$Ps1))),
-            (!cast<Instruction>(NAME) PPRAny:$Ps1, -1, 0)>;
   def : Pat<(intty (anyext (predty PPRAny:$Ps1))),
             (!cast<Instruction>(NAME) PPRAny:$Ps1, 1, 0)>;
   def : Pat<(vselect predty:$Pg,
                 (intty (splat_vector (scalarty (cpx i32:$imm, i32:$shift)))),
                 (intty (splat_vector (scalarty 0)))),
             (!cast<Instruction>(NAME) $Pg, $imm, $shift)>;
 }
 

llvm/test/CodeGen/AArch64/predicated-add-sub-mul.ll

 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %v = icmp ugt <vscale x 16 x i32> %a0, %a1
   %zero.extend = zext <vscale x 16 x i1> %v to <vscale x 16 x i32>
   %result = mul <vscale x 16 x i32> %zero.extend, %a0
   ret <vscale x 16 x i32> %result
 }
 
+define <vscale x 2 x i32> @sext.add2(<vscale x 2 x i32> %a0, <vscale x 2 x i32> %a1) #0 {
+; CHECK-LABEL: sext.add2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z2.d, z0.d
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z1.d, z1.d, #0xffffffff
+; CHECK-NEXT:    and z2.d, z2.d, #0xffffffff
+; CHECK-NEXT:    cmphi p0.d, p0/z, z2.d, z1.d
+; CHECK-NEXT:    mov z1.d, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    add z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 2 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 2 x i1> %v to <vscale x 2 x i32>
+  %result = add <vscale x 2 x i32> %zero.extend, %a0
+  ret <vscale x 2 x i32> %result
+}
+
+define <vscale x 4 x i32> @sext.add4(<vscale x 4 x i32> %a0, <vscale x 4 x i32> %a1) #0 {
+; CHECK-LABEL: sext.add4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p0.s, p0/z, z0.s, z1.s
+; CHECK-NEXT:    mov z1.s, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    add z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 4 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 4 x i1> %v to <vscale x 4 x i32>
+  %result = add <vscale x 4 x i32> %zero.extend, %a0
+  ret <vscale x 4 x i32> %result
+}
+
+define <vscale x 8 x i32> @sext.add8(<vscale x 8 x i32> %a0, <vscale x 8 x i32> %a1) #0 {
+; CHECK-LABEL: sext.add8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p1.s, p0/z, z1.s, z3.s
+; CHECK-NEXT:    cmphi p0.s, p0/z, z0.s, z2.s
+; CHECK-NEXT:    mov z2.s, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    add z0.s, p0/m, z0.s, z2.s
+; CHECK-NEXT:    add z1.s, p1/m, z1.s, z2.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 8 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 8 x i1> %v to <vscale x 8 x i32>
+  %result = add <vscale x 8 x i32> %zero.extend, %a0
+  ret <vscale x 8 x i32> %result
+}
+
+define <vscale x 16 x i32> @sext.add16(<vscale x 16 x i32> %a0, <vscale x 16 x i32> %a1) #0 {
+; CHECK-LABEL: sext.add16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p1.s, p0/z, z3.s, z7.s
+; CHECK-NEXT:    cmphi p2.s, p0/z, z2.s, z6.s
+; CHECK-NEXT:    cmphi p3.s, p0/z, z1.s, z5.s
+; CHECK-NEXT:    cmphi p0.s, p0/z, z0.s, z4.s
+; CHECK-NEXT:    mov z4.s, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    add z0.s, p0/m, z0.s, z4.s
+; CHECK-NEXT:    add z1.s, p3/m, z1.s, z4.s
+; CHECK-NEXT:    add z2.s, p2/m, z2.s, z4.s
+; CHECK-NEXT:    add z3.s, p1/m, z3.s, z4.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 16 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 16 x i1> %v to <vscale x 16 x i32>
+  %result = add <vscale x 16 x i32> %zero.extend, %a0
+  ret <vscale x 16 x i32> %result
+}
+
+define <vscale x 2 x i32> @sext.sub2(<vscale x 2 x i32> %a0, <vscale x 2 x i32> %a1) #0 {
+; CHECK-LABEL: sext.sub2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z2.d, z0.d
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z1.d, z1.d, #0xffffffff
+; CHECK-NEXT:    and z2.d, z2.d, #0xffffffff
+; CHECK-NEXT:    cmphi p0.d, p0/z, z2.d, z1.d
+; CHECK-NEXT:    mov z1.d, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    sub z0.d, z1.d, z0.d
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 2 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 2 x i1> %v to <vscale x 2 x i32>
+  %result = sub <vscale x 2 x i32> %zero.extend, %a0
+  ret <vscale x 2 x i32> %result
+}
+
+define <vscale x 4 x i32> @sext.sub4(<vscale x 4 x i32> %a0, <vscale x 4 x i32> %a1) #0 {
+; CHECK-LABEL: sext.sub4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p0.s, p0/z, z0.s, z1.s
+; CHECK-NEXT:    mov z1.s, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    sub z0.s, z1.s, z0.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 4 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 4 x i1> %v to <vscale x 4 x i32>
+  %result = sub <vscale x 4 x i32> %zero.extend, %a0
+  ret <vscale x 4 x i32> %result
+}
+
+define <vscale x 8 x i32> @sext.sub8(<vscale x 8 x i32> %a0, <vscale x 8 x i32> %a1) #0 {
+; CHECK-LABEL: sext.sub8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p1.s, p0/z, z0.s, z2.s
+; CHECK-NEXT:    cmphi p0.s, p0/z, z1.s, z3.s
+; CHECK-NEXT:    mov z2.s, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z3.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    sub z0.s, z3.s, z0.s
+; CHECK-NEXT:    sub z1.s, z2.s, z1.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 8 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 8 x i1> %v to <vscale x 8 x i32>
+  %result = sub <vscale x 8 x i32> %zero.extend, %a0
+  ret <vscale x 8 x i32> %result
+}
+
+define <vscale x 16 x i32> @sext.sub16(<vscale x 16 x i32> %a0, <vscale x 16 x i32> %a1) #0 {
+; CHECK-LABEL: sext.sub16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p1.s, p0/z, z2.s, z6.s
+; CHECK-NEXT:    cmphi p2.s, p0/z, z0.s, z4.s
+; CHECK-NEXT:    cmphi p3.s, p0/z, z1.s, z5.s
+; CHECK-NEXT:    cmphi p0.s, p0/z, z3.s, z7.s
+; CHECK-NEXT:    mov z4.s, p0/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z5.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z6.s, p3/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z7.s, p2/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    sub z0.s, z7.s, z0.s
+; CHECK-NEXT:    sub z1.s, z6.s, z1.s
+; CHECK-NEXT:    sub z2.s, z5.s, z2.s
+; CHECK-NEXT:    sub z3.s, z4.s, z3.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 16 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 16 x i1> %v to <vscale x 16 x i32>
+  %result = sub <vscale x 16 x i32> %zero.extend, %a0
+  ret <vscale x 16 x i32> %result
+}
+
+define <vscale x 2 x i32> @sext.mul2(<vscale x 2 x i32> %a0, <vscale x 2 x i32> %a1) #0 {
+; CHECK-LABEL: sext.mul2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z2.d, z0.d
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    and z1.d, z1.d, #0xffffffff
+; CHECK-NEXT:    and z2.d, z2.d, #0xffffffff
+; CHECK-NEXT:    cmphi p1.d, p0/z, z2.d, z1.d
+; CHECK-NEXT:    mov z1.d, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mul z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 2 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 2 x i1> %v to <vscale x 2 x i32>
+  %result = mul <vscale x 2 x i32> %zero.extend, %a0
+  ret <vscale x 2 x i32> %result
+}
+
+define <vscale x 4 x i32> @sext.mul4(<vscale x 4 x i32> %a0, <vscale x 4 x i32> %a1) #0 {
+; CHECK-LABEL: sext.mul4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p1.s, p0/z, z0.s, z1.s
+; CHECK-NEXT:    mov z1.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mul z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 4 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 4 x i1> %v to <vscale x 4 x i32>
+  %result = mul <vscale x 4 x i32> %zero.extend, %a0
+  ret <vscale x 4 x i32> %result
+}
+
+define <vscale x 8 x i32> @sext.mul8(<vscale x 8 x i32> %a0, <vscale x 8 x i32> %a1) #0 {
+; CHECK-LABEL: sext.mul8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmphi p1.s, p0/z, z0.s, z2.s
+; CHECK-NEXT:    cmphi p2.s, p0/z, z1.s, z3.s
+; CHECK-NEXT:    mov z2.s, p2/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z3.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mul z0.s, p0/m, z0.s, z3.s
+; CHECK-NEXT:    mul z1.s, p0/m, z1.s, z2.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 8 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 8 x i1> %v to <vscale x 8 x i32>
+  %result = mul <vscale x 8 x i32> %zero.extend, %a0
+  ret <vscale x 8 x i32> %result
+}
+
+define <vscale x 16 x i32> @sext.mul16(<vscale x 16 x i32> %a0, <vscale x 16 x i32> %a1) #0 {
+; CHECK-LABEL: sext.mul16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    addvl sp, sp, #-1
+; CHECK-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x08, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 8 * VG
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    str p4, [sp, #7, mul vl] // 2-byte Folded Spill
+; CHECK-NEXT:    cmphi p4.s, p0/z, z3.s, z7.s
+; CHECK-NEXT:    cmphi p2.s, p0/z, z0.s, z4.s
+; CHECK-NEXT:    mov z4.s, p4/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    ldr p4, [sp, #7, mul vl] // 2-byte Folded Reload
+; CHECK-NEXT:    cmphi p1.s, p0/z, z2.s, z6.s
+; CHECK-NEXT:    cmphi p3.s, p0/z, z1.s, z5.s
+; CHECK-NEXT:    mov z5.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z6.s, p3/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mov z7.s, p2/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    mul z1.s, p0/m, z1.s, z6.s
+; CHECK-NEXT:    mul z0.s, p0/m, z0.s, z7.s
+; CHECK-NEXT:    mul z2.s, p0/m, z2.s, z5.s
+; CHECK-NEXT:    mul z3.s, p0/m, z3.s, z4.s
+; CHECK-NEXT:    addvl sp, sp, #1
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 16 x i32> %a0, %a1
+  %zero.extend = sext <vscale x 16 x i1> %v to <vscale x 16 x i32>
+  %result = mul <vscale x 16 x i32> %zero.extend, %a0
+  ret <vscale x 16 x i32> %result
+}
+
 attributes #0 = { "target-features"="+sve" }