[AArch64] Canonicalize ZERO_EXTEND to VSELECT

[NicolaLancellotti]

Summary

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

[NicolaLancellotti]

This update removes an if statement that prevents the canonicalization

[paulwalker-arm]

Given you want to canonicalise to vselect is it worth doing this during lowering and removing the existing zext isel pattern? that way we can be sure all instances use the same idiom? This looks like it would be equally useful for sign_extend's of predicates also?

[NicolaLancellotti]

Address comments

[NicolaLancellotti]

In D135596#3848253, @paulwalker-arm wrote:

Given you want to canonicalise to vselect is it worth doing this during lowering and removing the existing zext isel pattern? that way we can be sure all instances use the same idiom? This looks like it would be equally useful for sign_extend's of predicates also?

Paul, thank you very much for the review.
I moved the code so that the canonicalisation is done during the lowering, and I removed the existing zext isel pattern.
You are right, it would be useful for sign_extend too, I'll create another patch later.

[paulwalker-arm]

A couple of nits but it looks like the move to lowering has caused some side effects whereby we now return "Invalid cost" for predicate based llvm.experimental.vector.splice.nxv16i1. I think we're just missing some entries in AArch64TTIImpl::getCastInstrCost like

{ 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 },

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5545

Can use DL, VT here as above.

5765–5766

Please can you move these two lines into LowerZERO_EXTEND so all the zext lowering has the same entry point.

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

515

Does (SVEDup0) work here as a shortcut?

[NicolaLancellotti]

Addressed comments

[NicolaLancellotti]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

515

I think so, I've just replaced zext with vselect, because of the new lowering.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5548–5549

This should all fit on one line now. You probably just need to rerun clang-format.

[NicolaLancellotti]

Files have been reformatted.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

SDValue LowerCallResult(SDValue Chain, SDValue InFlag,
                           const SDLoc &DL, SelectionDAG &DAG,
                           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 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_ADD, VT, Custom);
       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::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);
if (SDValue Cmp = emitConjunction(DAG, Cond, CC)) {
     SDValue CCVal = DAG.getConstant(CC, dl, MVT::i32);
     return DAG.getNode(AArch64ISD::BRCOND, dl, MVT::Other, Chain, Dest, CCVal,
                        Cmp);
   }
 
   return SDValue();
 }
 
+SDValue AArch64TargetLowering::LowerZERO_EXTEND(SDValue Op, SelectionDAG &DAG) const {
+  assert(Op->getOpcode() == ISD::ZERO_EXTEND && "Expected ZERO_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 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::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   LLVM_DEBUG(dbgs() << "Custom lowering: ");
   LLVM_DEBUG(Op.dump());
 
   switch (Op.getOpcode()) {
   default:
     llvm_unreachable("unimplemented operand");
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   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:
-  case ISD::ZERO_EXTEND:
     return LowerFixedLengthVectorIntExtendToSVE(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))
       return SDValue();
 
     return LowerToPredicatedOp(Op, DAG,

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVEorSME] in {
   defm CLZ_ZPmZ  : sve_int_un_pred_arit_1<   0b001, "clz",  AArch64clz_mt>;
   defm CNT_ZPmZ  : sve_int_un_pred_arit_1<   0b010, "cnt",  AArch64cnt_mt>;
   defm CNOT_ZPmZ : sve_int_un_pred_arit_1<   0b011, "cnot", AArch64cnot_mt>;
   defm NOT_ZPmZ  : sve_int_un_pred_arit_1<   0b110, "not",  AArch64not_mt>;
   defm FABS_ZPmZ : sve_int_un_pred_arit_1_fp<0b100, "fabs", AArch64fabs_mt>;
   defm FNEG_ZPmZ : sve_int_un_pred_arit_1_fp<0b101, "fneg", AArch64fneg_mt>;
 
   // zext(cmpeq(x, splat(0))) -> cnot(x)
-  def : Pat<(nxv16i8 (zext (nxv16i1 (AArch64setcc_z (nxv16i1 (SVEAllActive):$Pg), nxv16i8:$Op2, (SVEDup0), SETEQ)))),
+  def : Pat<(nxv16i8 (vselect (nxv16i1 (AArch64setcc_z (nxv16i1 (SVEAllActive):$Pg), nxv16i8:$Op2, (SVEDup0), SETEQ)), (nxv16i8 (splat_vector (i32 1))), (nxv16i8 (splat_vector (i32 0))))),
             (CNOT_ZPmZ_B $Op2, $Pg, $Op2)>;
-  def : Pat<(nxv8i16 (zext (nxv8i1 (AArch64setcc_z (nxv8i1 (SVEAllActive):$Pg), nxv8i16:$Op2, (SVEDup0), SETEQ)))),
+  def : Pat<(nxv8i16 (vselect (nxv8i1 (AArch64setcc_z (nxv8i1 (SVEAllActive):$Pg), nxv8i16:$Op2, (SVEDup0), SETEQ)), (nxv8i16 (splat_vector (i32 1))), (nxv8i16 (splat_vector (i32 0))))),
             (CNOT_ZPmZ_H $Op2, $Pg, $Op2)>;
-  def : Pat<(nxv4i32 (zext (nxv4i1 (AArch64setcc_z (nxv4i1 (SVEAllActive):$Pg), nxv4i32:$Op2, (SVEDup0), SETEQ)))),
+  def : Pat<(nxv4i32 (vselect (nxv4i1 (AArch64setcc_z (nxv4i1 (SVEAllActive):$Pg), nxv4i32:$Op2, (SVEDup0), SETEQ)), (nxv4i32 (splat_vector (i32 1))), (nxv4i32 (splat_vector (i32 0))))),
             (CNOT_ZPmZ_S $Op2, $Pg, $Op2)>;
-  def : Pat<(nxv2i64 (zext (nxv2i1 (AArch64setcc_z (nxv2i1 (SVEAllActive):$Pg), nxv2i64:$Op2, (SVEDup0), SETEQ)))),
+  def : Pat<(nxv2i64 (vselect (nxv2i1 (AArch64setcc_z (nxv2i1 (SVEAllActive):$Pg), nxv2i64:$Op2, (SVEDup0), SETEQ)), (nxv2i64 (splat_vector (i64 1))), (nxv2i64 (splat_vector (i64 0))))),
             (CNOT_ZPmZ_D $Op2, $Pg, $Op2)>;
 
   defm SMAX_ZPmZ : sve_int_bin_pred_arit_1<0b000, "smax", "SMAX_ZPZZ", AArch64smax_m1, DestructiveBinaryComm>;
   defm UMAX_ZPmZ : sve_int_bin_pred_arit_1<0b001, "umax", "UMAX_ZPZZ", AArch64umax_m1, DestructiveBinaryComm>;
   defm SMIN_ZPmZ : sve_int_bin_pred_arit_1<0b010, "smin", "SMIN_ZPZZ", AArch64smin_m1, DestructiveBinaryComm>;
   defm UMIN_ZPmZ : sve_int_bin_pred_arit_1<0b011, "umin", "UMIN_ZPZZ", AArch64umin_m1, DestructiveBinaryComm>;
   defm SABD_ZPmZ : sve_int_bin_pred_arit_1<0b100, "sabd", "SABD_ZPZZ", int_aarch64_sve_sabd, DestructiveBinaryComm>;
   defm UABD_ZPmZ : sve_int_bin_pred_arit_1<0b101, "uabd", "UABD_ZPZZ", int_aarch64_sve_uabd, DestructiveBinaryComm>;

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

ConversionTbl[] = {
     { ISD::BITCAST, MVT::nxv2f16, MVT::nxv2i16, 0 },
     { ISD::BITCAST, MVT::nxv4f16, MVT::nxv4i16, 0 },
     { ISD::BITCAST, MVT::nxv2f32, MVT::nxv2i32, 0 },
 
     // Bitcasts from integer to float
     { ISD::BITCAST, MVT::nxv2i16, MVT::nxv2f16, 0 },
     { 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 },
   };
 
   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 (zext (predty PPRAny:$Ps1))),
-            (!cast<Instruction>(NAME) PPRAny:$Ps1, 1, 0)>;
   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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux"
+
+define <vscale x 2 x i32> @zext.add2(<vscale x 2 x i32> %a0, <vscale x 2 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 = zext <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> @zext.add4(<vscale x 4 x i32> %a0, <vscale x 4 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 = zext <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> @zext.add8(<vscale x 8 x i32> %a0, <vscale x 8 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 = zext <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> @zext.add16(<vscale x 16 x i32> %a0, <vscale x 16 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 = zext <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> @zext.sub2(<vscale x 2 x i32> %a0, <vscale x 2 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; CHECK-NEXT:    sub z0.d, z1.d, z0.d
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 2 x i32> %a0, %a1
+  %zero.extend = zext <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> @zext.sub4(<vscale x 4 x i32> %a0, <vscale x 4 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; CHECK-NEXT:    sub z0.s, z1.s, z0.s
+; CHECK-NEXT:    ret
+  %v = icmp ugt <vscale x 4 x i32> %a0, %a1
+  %zero.extend = zext <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> @zext.sub8(<vscale x 8 x i32> %a0, <vscale x 8 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; CHECK-NEXT:    mov z3.s, p1/z, #1 // =0x1
+; 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 = zext <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> @zext.sub16(<vscale x 16 x i32> %a0, <vscale x 16 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; CHECK-NEXT:    mov z5.s, p1/z, #1 // =0x1
+; CHECK-NEXT:    mov z6.s, p3/z, #1 // =0x1
+; CHECK-NEXT:    mov z7.s, p2/z, #1 // =0x1
+; 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 = zext <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> @zext.mul2(<vscale x 2 x i32> %a0, <vscale x 2 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 = zext <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> @zext.mul4(<vscale x 4 x i32> %a0, <vscale x 4 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 = zext <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> @zext.mul8(<vscale x 8 x i32> %a0, <vscale x 8 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; CHECK-NEXT:    mov z3.s, p1/z, #1 // =0x1
+; 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 = zext <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> @zext.mul16(<vscale x 16 x i32> %a0, <vscale x 16 x i32> %a1) #0 {
+; CHECK-LABEL: zext.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 // =0x1
+; 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 // =0x1
+; CHECK-NEXT:    mov z6.s, p3/z, #1 // =0x1
+; CHECK-NEXT:    mov z7.s, p2/z, #1 // =0x1
+; 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 = 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
+}
+
+attributes #0 = { "target-features"="+sve" }