[AArch64][SVE] Use PTRUE instruction for get_active_lane_mask intrinsic if the range is appropriate for predicator constant

[dtemirbulatov]

Summary

While get_active_lane_mask lowering it uses WHILELO instruction, but for constant range suitable for PTRUE then we could issue PTRUE instruction instead.

[peterwaller-arm]

llvm/test/CodeGen/AArch64/active_lane_mask.ll

486

Please can you add a test which shows what happens when the %n parameter of active.lane.mask goes out of the range that can be represented as a ptrue immediate. I expect this should fall back to the old codegen?

[paulwalker-arm]

This transformation is more subtle than you expect. See https://developer.arm.com/documentation/ddi0602/2022-09/SVE-Instructions/PTRUE--predicate---Initialise-predicate-from-named-constraint-, specifically the part that reads If the constraint specifies more elements than are available at the current vector length then all elements of the destination predicate are set to false..

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17640

Notwithstanding Paul's comment, I think something like this also needs to account for the mapping between PredPattern constants and the integer %n constant.

https://github.com/llvm/llvm-project/blob/2e999b7dd1934a44d38c3a753460f1e5a217e9a5/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h#L530-L556

Note that 1-8 have a one-to-one correspondance and above that they don't.

[dtemirbulatov]

Fixed remarks.

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17639–17640

Unless I'm missing something, this condition doesn't look right.

I suspect this combine could happen before legalization in which case you could have vscale x [big] x something, and the [big] you're testing against here could exceed the runtime vector length -- and then you're hitting the problem Paul pointed out that ptrue would return an empty predicate.

I take to account for this properly you're going to need to query the minimum vscale via SubTarget->getMinSVEVectorSizeInBits().

[dtemirbulatov]

Rebased, addressed comments.

[dtemirbulatov]

Add support for any range not just from 0.

[c-rhodes]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4755

unsigned

4755–4757

std::max(Subtarget->getMinSVEVectorSizeInBits(), 128u)

4759

NumActiveElems?

4759–4760

Op.getConstantOperandVal(2) - Op.getConstantOperandVal(1);

4762

this should be <=

4763

this should be the encoded pattern returned by getSVEPredPatternFromNumElements, see elsewhere in ISelLowering for examples.

llvm/test/CodeGen/AArch64/active_lane_mask.ll

478

these tests aren't sufficient, it would be good to have tests for

• the upper bound of active elements vl256, this can be tested with a <vscale x 16 x i1> predicate and minimum SVE vector width of 2048 specified via vscale_range(16, 16) function attribute. Also add a test for this vector width that exceeds 256 active elements.
• A non-zero %base argument for active.lane.mask.

495

active.lane.mask is an unsigned less-than comparison so this generates a predicate with 4 active lanes, which for a minimum vector width of 128 and esize of 32 can be represented with ptrue p0.s, vl4?

[dtemirbulatov]

Addressed remarks.

[dtemirbulatov]

llvm/test/CodeGen/AArch64/active_lane_mask.ll

495

correct.

[c-rhodes]

LGTM

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4758–4759

This doesn't look correct because while can take 64-bit operation. This code can truncate the result in such a way that a ptrue might incorrectly look like a viable replacement.

4763

It likely doesn't really matter but given getSVEPredPatternFromNumElements takes an unsigned (because it's linked to the interface used by getVectorMinNumElements) perhaps it's clean to perform the <= check before calling getSVEPredPatternFromNumElements, that way we know the truncation is not losing any precision.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4758–4759

Sorry I meant "while can take 64-bit operands".

[dtemirbulatov]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4758–4759

ok, as whilelo operands are i32, I think we have to zero extend those operands to unsigned by cast<ConstantSDNode>(Op)->getZExtValue().

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4758–4759

whilelo operands are not restricted to i32, which can be seen by looking at IntrinsicsAArch64.td. Although the definition suggests any integer type is allowed, that's just a deficiency of how the td file is parsed. Realistically these intrinsics only accept either i32 or i64 operands, which are tested by sve-intrinsics-while.ll.

Perhaps you can keep things as APInt up until the comparison, at which point you know an unsigned will be big enough to hold the length.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case Intrinsic::aarch64_neon_scalar_uqxtn: {
     if (Op.getValueType() == MVT::i32)
       return DAG.getNode(ISD::BITCAST, dl, MVT::i32,
                          DAG.getNode(ISD::INTRINSIC_WO_CHAIN, dl, MVT::f32,
                                      Op.getOperand(0),
                                      DAG.getNode(ISD::BITCAST, dl, MVT::f64,
                                                  Op.getOperand(1))));
     return SDValue();
   }
+  case Intrinsic::aarch64_sve_whilelo: {
+    if (isa<ConstantSDNode>(Op.getOperand(1)) &&
+        isa<ConstantSDNode>(Op.getOperand(2))) {
+      unsigned MinSVEVectorSize =
+          std::max(Subtarget->getMinSVEVectorSizeInBits(), 128u);
+      unsigned ElementSize = 128 / Op.getValueType().getVectorMinNumElements();
+      unsigned NumActiveElems =
+          Op.getConstantOperandVal(2) - Op.getConstantOperandVal(1);
+      Optional<unsigned> PredPattern =
+          getSVEPredPatternFromNumElements(NumActiveElems);
+      if ((PredPattern != None) &&
+          NumActiveElems <= (MinSVEVectorSize / ElementSize))
+        return getPTrue(DAG, dl, Op.getValueType(), *PredPattern);
+    }
+    return SDValue();
+  }
   case Intrinsic::aarch64_sve_sunpkhi:
     return DAG.getNode(AArch64ISD::SUNPKHI, dl, Op.getValueType(),
                        Op.getOperand(1));
   case Intrinsic::aarch64_sve_sunpklo:
     return DAG.getNode(AArch64ISD::SUNPKLO, dl, Op.getValueType(),
                        Op.getOperand(1));
   case Intrinsic::aarch64_sve_uunpkhi:
     return DAG.getNode(AArch64ISD::UUNPKHI, dl, Op.getValueType(),
if (VT.isFixedLengthVector()) {
       Res = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ExtVT, Res,
                         DAG.getConstant(0, DL, MVT::i64));
       Res = DAG.getNode(ISD::TRUNCATE, DL, VT, Res);
     }
     return Res;
   }
   case Intrinsic::aarch64_neon_vcvtfxs2fp:
   case Intrinsic::aarch64_neon_vcvtfxu2fp:
     return tryCombineFixedPointConvert(N, DCI, DAG);
   case Intrinsic::aarch64_neon_saddv:
     return combineAcrossLanesIntrinsic(AArch64ISD::SADDV, N, DAG);
   case Intrinsic::aarch64_neon_uaddv:
     return combineAcrossLanesIntrinsic(AArch64ISD::UADDV, N, DAG);
   case Intrinsic::aarch64_neon_sminv:
     return combineAcrossLanesIntrinsic(AArch64ISD::SMINV, N, DAG);
   case Intrinsic::aarch64_neon_uminv:
     return combineAcrossLanesIntrinsic(AArch64ISD::UMINV, N, DAG);
   case Intrinsic::aarch64_neon_smaxv:

llvm/test/CodeGen/AArch64/active_lane_mask.ll

 ; CHECK-NEXT:    umin v1.2s, v1.2s, v0.2s
 ; CHECK-NEXT:    and v0.8b, v3.8b, v0.8b
 ; CHECK-NEXT:    cmhi v0.2s, v0.2s, v1.2s
 ; CHECK-NEXT:    ret
   %active.lane.mask = call <2 x i1> @llvm.get.active.lane.mask.v2i1.i8(i8 %index, i8 %TC)
   ret <2 x i1> %active.lane.mask
 }
 
+define <vscale x 4 x i1> @lane_mask_nxv4i1_imm3() {
+; CHECK-LABEL: lane_mask_nxv4i1_imm3:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s, vl3
+; CHECK-NEXT:    ret
+entry:
+  %active.lane.mask = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 0, i64 3)
+  ret <vscale x 4 x i1> %active.lane.mask
+}
+
+define <vscale x 4 x i1> @lane_mask_nxv4i1_imm5() {
+; CHECK-LABEL: lane_mask_nxv4i1_imm5:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w8, #5
+; CHECK-NEXT:    whilelo p0.s, xzr, x8
+; CHECK-NEXT:    ret
+entry:
+  %active.lane.mask = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 0, i64 5)
+  ret <vscale x 4 x i1> %active.lane.mask
+}
+
+define <vscale x 4 x i1> @lane_mask_nxv4i1_imm4() {
+; CHECK-LABEL: lane_mask_nxv4i1_imm4:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s, vl4
+; CHECK-NEXT:    ret
+entry:
+  %active.lane.mask = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 10, i64 14)
+  ret <vscale x 4 x i1> %active.lane.mask
+}
+
+define <vscale x 16 x i1> @lane_mask_nxv16i1_imm10() {
+; CHECK-LABEL: lane_mask_nxv16i1_imm10:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w8, #10
+; CHECK-NEXT:    whilelo p0.b, xzr, x8
+; CHECK-NEXT:    ret
+entry:
+  %active.lane.mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 0, i64 10)
+  ret <vscale x 16 x i1> %active.lane.mask
+}
+
+define <vscale x 16 x i1> @lane_mask_nxv16i1_imm256() vscale_range(16, 16) {
+; CHECK-LABEL: lane_mask_nxv16i1_imm256:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.b, vl256
+; CHECK-NEXT:    ret
+entry:
+  %active.lane.mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 0, i64 256)
+  ret <vscale x 16 x i1> %active.lane.mask
+}
+
 
 declare <vscale x 32 x i1> @llvm.get.active.lane.mask.nxv32i1.i32(i32, i32)
 declare <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i32(i32, i32)
 declare <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i32(i32, i32)
 declare <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i32(i32, i32)
 declare <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i32(i32, i32)
 
 declare <vscale x 32 x i1> @llvm.get.active.lane.mask.nxv32i1.i64(i64, i64)