[AArch64][SVE] Fix BRKNS bug in optimizePTestInstr

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.

Reviewed By: paulwalker-arm

Differential Revision: https://reviews.llvm.org/D135655

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 3994c82..638851c 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1351,8 +1351,9 @@ bool AArch64InstrInfo::optimizePTestInstr(
       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;

diff --git a/llvm/lib/Target/AArch64/SVEInstrFormats.td b/llvm/lib/Target/AArch64/SVEInstrFormats.td
index 9fb834b..b301e07 100644 (file)
--- a/llvm/lib/Target/AArch64/SVEInstrFormats.td
+++ b/llvm/lib/Target/AArch64/SVEInstrFormats.td
@@ -7987,6 +7987,7 @@ class sve_int_brkn<bit S, string asm>
 
   let Constraints = "$Pdm = $_Pdm";
   let Defs = !if(S, [NZCV], []);
+  let ElementSize = ElementSizeB;
 }
 
 multiclass sve_int_brkn<bits<1> opc, string asm, SDPatternOperator op> {

diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll b/llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll
index bd7ed49..718cc13 100644 (file)
--- a/llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll
+++ b/llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll
@@ -30,7 +30,8 @@ define i32 @brkb(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a) {
 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)
@@ -39,6 +40,19 @@ define i32 @brkn(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i
   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) {
@@ -84,3 +98,4 @@ declare <vscale x 16 x i1> @llvm.aarch64.sve.brkpb.z.nxv16i1(<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)