[AArch64][SVE] Sink op into loop if it's used by PTEST and known to zero inactive lanes.

[sdesmalen]

Summary

This helps fold away the ptest instructions, which needs the knowledge on whether
the general predicate is known to zero the inactive lanes.

This fixes some PTEST regressions introduced by D129282.

[RosieSumpter]

Thanks @sdesmalen for fixing some of my regressions!

llvm/test/CodeGen/AArch64/sve-ptest-removal-sink.ll

9

Looks like %a and %b aren't used anywhere?

[paulwalker-arm]

Can you provide a C/C++ example for when this is a problem? I'm not saying there's no problem to solve but rather that I'm starting to wonder if early decisions about removing reinterpret intrinsics are now starting to bite us.

[sdesmalen]

In D129852#3655364, @paulwalker-arm wrote:

Can you provide a C/C++ example for when this is a problem? I'm not saying there's no problem to solve but rather that I'm starting to wonder if early decisions about removing reinterpret intrinsics are now starting to bite us.

https://godbolt.org/z/q6hKMsxKK

In this case it doesn't have anything to do with the reinterpret intrinsics, but rather the ptrue that gets hoisted out which is then no longer visible to (non-Global) ISel when it lowers the ptest.

[sdesmalen]

In D129852#3655395, @sdesmalen wrote:

In D129852#3655364, @paulwalker-arm wrote:

Can you provide a C/C++ example for when this is a problem? I'm not saying there's no problem to solve but rather that I'm starting to wonder if early decisions about removing reinterpret intrinsics are now starting to bite us.

https://godbolt.org/z/q6hKMsxKK

In this case it doesn't have anything to do with the reinterpret intrinsics, but rather the ptrue that gets hoisted out which is then no longer visible to (non-Global) ISel when it lowers the ptest.

Just to expand on this a little bit, I experimented with disabling the removal of the reinterpret intrinsics, but found that the convert.to.svbool call is still outside the loop at the point of reaching codegen, meaning that it still can't fold away the ptest instruction unless it actively tries to sink the operand.

[paulwalker-arm]

In D129852#3659230, @sdesmalen wrote:

In D129852#3655395, @sdesmalen wrote:

In D129852#3655364, @paulwalker-arm wrote:

Can you provide a C/C++ example for when this is a problem? I'm not saying there's no problem to solve but rather that I'm starting to wonder if early decisions about removing reinterpret intrinsics are now starting to bite us.

https://godbolt.org/z/q6hKMsxKK

In this case it doesn't have anything to do with the reinterpret intrinsics, but rather the ptrue that gets hoisted out which is then no longer visible to (non-Global) ISel when it lowers the ptest.

Just to expand on this a little bit, I experimented with disabling the removal of the reinterpret intrinsics, but found that the convert.to.svbool call is still outside the loop at the point of reaching codegen, meaning that it still can't fold away the ptest instruction unless it actively tries to sink the operand.

I can believe that because removing those optimisations just replace the implicit zeroing with explicit zeroing. Whereas, I'm trying to figure out if we need to tighten up the intrinsics across the board to stop hiding the fact that predicate setting intrinsics always set all the bits. There's nothing here for you to worry about because this patch is generally restoring previous behaviour. I'm just worried we'll see an ever increase list of places were we need to restore the zeroing information that we've deliberately hidden.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12793

My worry here is pulling something into a loop that's more expensive than what we're trying to remove. Would it be wrong to limit this to simple constant splats and ptrue intrinsics?

[sdesmalen]

Limit opcode to ptrue intrinsic.

[paulwalker-arm]

Please don't forget @RosieSumpter's unused parameters comment.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     case Intrinsic::aarch64_neon_sqdmulh:
     case Intrinsic::aarch64_neon_sqrdmulh:
       // Sink splats for index lane variants
       if (isSplatShuffle(II->getOperand(0)))
         Ops.push_back(&II->getOperandUse(0));
       if (isSplatShuffle(II->getOperand(1)))
         Ops.push_back(&II->getOperandUse(1));
       return !Ops.empty();
+    case Intrinsic::aarch64_sve_ptest_first:
+    case Intrinsic::aarch64_sve_ptest_last:
+      if (auto *IIOp = dyn_cast<IntrinsicInst>(II->getOperand(0)))
+        if (IIOp->getIntrinsicID() == Intrinsic::aarch64_sve_ptrue)
+          Ops.push_back(&II->getOperandUse(0));
+      return !Ops.empty();
     case Intrinsic::aarch64_sme_write_horiz:
     case Intrinsic::aarch64_sme_write_vert:
     case Intrinsic::aarch64_sme_writeq_horiz:
     case Intrinsic::aarch64_sme_writeq_vert: {
       auto *Idx = dyn_cast<Instruction>(II->getOperand(1));
       if (!Idx || Idx->getOpcode() != Instruction::Add)
         return false;
       Ops.push_back(&II->getOperandUse(1));

llvm/test/CodeGen/AArch64/sve-ptest-removal-sink.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 -mattr=+sve < %s | FileCheck %s
+
+;
+; Ensure that the %ptrue from the preheader is sunk into the loop such that the ptest is removed.
+;
+
+define void @test_sink_ptrue_into_ptest(i32 %n) {
+; CHECK-LABEL: test_sink_ptrue_into_ptest:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    whilelt p0.s, wzr, w0
+; CHECK-NEXT:    b.pl .LBB0_3
+; CHECK-NEXT:  // %bb.1: // %for.body.preheader
+; CHECK-NEXT:    mov w9, wzr
+; CHECK-NEXT:    cntw x8
+; CHECK-NEXT:  .LBB0_2: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    add w10, w9, w8
+; CHECK-NEXT:    whilelt p0.s, w9, w0
+; CHECK-NEXT:    mov w9, w10
+; CHECK-NEXT:    b.mi .LBB0_2
+; CHECK-NEXT:  .LBB0_3: // %exit
+; CHECK-NEXT:    ret
+entry:
+  %vscale = tail call i32 @llvm.vscale.i32()
+  %step = shl nuw nsw i32 %vscale, 2
+  %ptrue.ph = tail call <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 31)
+  %while.ph = tail call <vscale x 4 x i1> @llvm.aarch64.sve.whilelt.nxv4i1.i32(i32 0, i32 %n)
+  %ptest.ph = tail call i1 @llvm.aarch64.sve.ptest.first.nxv4i1(<vscale x 4 x i1> %ptrue.ph, <vscale x 4 x i1> %while.ph)
+  br i1 %ptest.ph, label %for.body, label %exit
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %for.body ]
+  %i.next = add i32 %i, %step
+  %while = call <vscale x 4 x i1> @llvm.aarch64.sve.whilelt.nxv4i1.i32(i32 %i, i32 %n)
+  %ptest = call i1 @llvm.aarch64.sve.ptest.first.nxv4i1(<vscale x 4 x i1> %ptrue.ph, <vscale x 4 x i1> %while)
+  br i1 %ptest, label %for.body, label %exit
+
+exit:
+  ret void
+}
+
+declare i32 @llvm.vscale.i32()
+declare <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 immarg)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.whilelt.nxv4i1.i32(i32, i32)
+declare i1 @llvm.aarch64.sve.ptest.first.nxv4i1(<vscale x 4 x i1>, <vscale x 4 x i1>)