[PATCH] D135655: [AArch64][SVE] Fix BRKNS bug in optimizePTestInstr

[Cullen Rhodes via Phabricator via llvm-commits]

c-rhodes created this revision.
c-rhodes added reviewers: paulwalker-arm, peterwaller-arm, bsmith.
Herald added subscribers: psnobl, hiraditya, kristof.beyls, tschuett.
Herald added a reviewer: efriedma.
Herald added a project: All.
c-rhodes requested review of this revision.
Herald added a project: LLVM.

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 <https://reviews.llvm.org/D134946>.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D135655

Files:
  llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
  llvm/lib/Target/AArch64/SVEInstrFormats.td
  llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll


Index: llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll
===================================================================
--- llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll
+++ llvm/test/CodeGen/AArch64/sve-ptest-removal-brk.ll
@@ -30,7 +30,8 @@
 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 @@
   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.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)
Index: llvm/lib/Target/AArch64/SVEInstrFormats.td
===================================================================
--- llvm/lib/Target/AArch64/SVEInstrFormats.td
+++ llvm/lib/Target/AArch64/SVEInstrFormats.td
@@ -7987,6 +7987,7 @@
 
   let Constraints = "$Pdm = $_Pdm";
   let Defs = !if(S, [NZCV], []);
+  let ElementSize = ElementSizeB;
 }
 
 multiclass sve_int_brkn<bits<1> opc, string asm, SDPatternOperator op> {
Index: llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
===================================================================
--- llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1351,10 +1351,18 @@
       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)
+      uint64_t MaskElementSize = getElementSizeForOpcode(MaskOpcode);
+      uint64_t PredElementSize = getElementSizeForOpcode(PredOpcode);
+
+      // Must be an all active predicate of matching element size (.b).
+      if ((PredElementSize != MaskElementSize) || !isPTrueOpcode(MaskOpcode) ||
+          (Mask->getOperand(1).getImm() != 31))
         return false;
 
+      assert((MaskElementSize == AArch64::ElementSizeB) &&
+             "Unexpected element size!");
+
       NewOp = AArch64::BRKNS_PPzP;
       OpChanged = true;
       break;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D135655.466757.patch
Type: text/x-patch
Size: 3384 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20221011/59e7f72e/attachment.bin>