[AArch64][SVE] Fix BRKNS bug in optimizePTestInstr

[c-rhodes]

Summary

The BRKNS instruction is unlike the other instructions that set flags
since it has an all active implicit predicate, so the existing

PTEST(PG, BRKN(PG, A, B)) -> BRKNS(PG, A, B)

in AArch64InstrInfo::optimizePTestInstr is incorrect, however

PTEST(PTRUE_B(31), BRKN(PG, A, B)) -> BRKNS(PG, A, B)

is correct.

Spotted by @paulwalker-arm in D134946.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1359

What about "(MaskOpcode != AArch64::PTRUE_B)"?

1359–1360

With the above I doubt we need this.

[c-rhodes]

Replace checks for is ptrue opcode + assert with single check for PTRUE_B opcode.

[c-rhodes]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1359

What about "(MaskOpcode != AArch64::PTRUE_B)"?

that makes more sense, done

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

case AArch64::BRKPB_PPzPP: {
 
       // Switch to the new opcode
       NewOp = Pred->getOpcode() == AArch64::BRKB_PPzP ? AArch64::BRKBS_PPzP
                                                       : AArch64::BRKPBS_PPzPP;
       OpChanged = true;
       break;
     }
     case AArch64::BRKN_PPzP: {
-      auto *PredMask = MRI->getUniqueVRegDef(Pred->getOperand(1).getReg());
-      if (Mask != PredMask)
+      // PTEST(PTRUE_B(31), BRKN(PG, A, B)) -> BRKNS(PG, A, B).
+      if ((MaskOpcode != AArch64::PTRUE_B) ||
+          (Mask->getOperand(1).getImm() != 31))
         return false;
 
       NewOp = AArch64::BRKNS_PPzP;
       OpChanged = true;
       break;
     }
     case AArch64::RDFFR_PPz: {
       // rdffr   p1.b, PredMask=p0/z <--- Definition of Pred
       // ptest   Mask=p0, Pred=p1.b  <--- If equal masks, remove this and use
       //                                  `rdffrs p1.b, p0/z` above.
       auto *PredMask = MRI->getUniqueVRegDef(Pred->getOperand(1).getReg());
       if (Mask != PredMask)

llvm/lib/Target/AArch64/SVEInstrFormats.td

: I<(outs PPR8:$Pdm), (ins PPRAny:$Pg, PPR8:$Pn, PPR8:$_Pdm),
   let Inst{13-10} = Pg;
   let Inst{9}     = 0b0;
   let Inst{8-5}   = Pn;
   let Inst{4}     = 0b0;
   let Inst{3-0}   = Pdm;
 
   let Constraints = "$Pdm = $_Pdm";
   let Defs = !if(S, [NZCV], []);
+  let ElementSize = ElementSizeB;
 }
 
 multiclass sve_int_brkn<bits<1> opc, string asm, SDPatternOperator op> {
   def NAME : sve_int_brkn<opc, asm>;
 
   def : SVE_3_Op_Pat<nxv16i1, op, nxv16i1, nxv16i1, nxv16i1, !cast<Instruction>(NAME)>;
 }
 

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

; CHECK-NEXT:    ret
   %2 = tail call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %1)
   %conv = zext i1 %2 to i32
   ret i32 %conv
 }
 
 define i32 @brkn(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b) {
 ; CHECK-LABEL: brkn:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    brkns p2.b, p0/z, p1.b, p2.b
+; CHECK-NEXT:    brkn p2.b, p0/z, p1.b, p2.b
+; CHECK-NEXT:    ptest p0, p2.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.brkn.z.nxv16i1(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b)
   %2 = tail call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %1)
   %conv = zext i1 %2 to i32
   ret i32 %conv
 }
 
+define i32 @brkn_all_active(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b) {
+; CHECK-LABEL: brkn_all_active:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    brkns p2.b, p0/z, p1.b, p2.b
+; CHECK-NEXT:    cset w0, ne
+; CHECK-NEXT:    ret
+  %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.brkn.z.nxv16i1(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b)
+  %2 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
+  %3 = tail call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> %2, <vscale x 16 x i1> %1)
+  %conv = zext i1 %3 to i32
+  ret i32 %conv
+}
+
 ; Test that ptest instruction is not removed when using a non-flag setting brk
 
 define i32 @brkpb_neg(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b) {
 ; CHECK-LABEL: brkpb_neg:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    brkpb p0.b, p0/z, p1.b, p2.b
 ; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
; CHECK-NEXT:    ret
   %conv = zext i1 %2 to i32
   ret i32 %conv
 }
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.brkpb.z.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>, <vscale x 16 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.brkb.z.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.brkn.z.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>, <vscale x 16 x i1>)
 declare i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
+declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32)