[CodeGen][AArch64][SVE] Fold [rdffr, ptest] => rdffrs; bugfix for optimizePTestInstr

[peterwaller-arm]

Summary

When a ptest is used to set flags from the output of rdffr, the ptest
can be eliminated, using a flags-setting rdffrs instead.

Additionally, check that nothing consumes flags between rdffr and ptest;
this case appears to have been missed previously.

• There is no unpredicated RDFFRS instruction.
• If substituting RDFFR_PP, require that the mask argument of the PTEST matches that of the RDFFR_PP.
• Move some precondition code up inside optimizePTestInstr, so that it covers the new code paths for RDFFR which return earlier.
  • Only consider RDFFR, PTEST in same basic block.
  • Check for other flag setting instructions between the two, abort if found.
  • Drop an old TODO comment about removing dead PTEST instructions.

RDFFR_P to follow in later patch.

[peterwaller-arm]

Fix clang-format issues.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1369–1372

Is this true? I don't see how you know the value of p0 just prior to rdffr. There's nothing to say p0 has been defined at this point nor that it contains the same value as it will by the time the ptest is hit.

Although not exactly the same this is what the

auto *RDFFRMask = MRI->getUniqueVRegDef(RDFFR->getOperand(1).getReg());
if (Mask != RDFFRMask)
  return false;

code for the AArch64::RDFFR_PPz case is protecting against.

Perhaps this case is worth a separate patch?

1394–1396

Not saying there's anything wrong but I just wondered why you cannot fall through here? Is the following code specific to the other case blocks? I only ask because RDFFR_PPz->RDFFRS_PPz looks very similar to BRKN_PPzP->BRKNS_PPzP.

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1369–1372

Great catch, agree that seems problematic.

If the only use of p1.b after the rdffr is at the ptest, can we think of it as optimizing the ptest instead (as the comment is written)? e.g. hasOneNonDBGUser(rdffr.Operand) && getUniqueVRegDef(ptest.Pred) == rdffr.

If that is a simple enough solution and seems obviously correct to you, I'll keep it in-patch, otherwise split.

1394–1396

I think you are correct that this could fall through.

It ended up like this because I first implemented RDFFR_P, which AIUI cannot fall through because it grows an operand. Then I used the same implementation strategy for RDFFR_PP. When I came to consider falling through I concluded that the non-fall-through code was more straightforward to read/reason about. It didn't seem to me a too much extra code (we 'spend 4 extra lines of obvious code' to 'get to a return true 30 lines of code earlier' which seemed a worthwhile trade).

I'm guessing that updating the existing MI is done to avoid additional allocations and reduce the total code slightly? Additionally it has to add NZCV as an operand and mark it live, none of which looks necessary if the instruction has been built.

What is your preference in this case? Shall I change to fall-through?

[peterwaller-arm]

• Update for code review.
• Drop RDFFR_P - harder than I thought, for a later patch.

[peterwaller-arm]

Update commit message.

[paulwalker-arm]

Looks good to me. I've a few comments that you're free to ignore.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1313–1321

I have no proof on this so feel free to ignore if you disagree.

Hoisting the Pred->getParent() != PTest->getParent() part offers a faster route to termination so I'm good with that.

However, I feel iterating across the BB's instructions may not and so that part might be better to keep until after we've decided Mask and Pred are interesting.

1373

Given there's only a single use, which matches the style of the previous blocks, I don't see much value in creating this alias to Pred. Especially as the comment describes the situation.

[peterwaller-arm]

• Update for code review
• Factor loop with pre-existing areCFlagsAccessedBetweenInstrs.
• Add test case for flags used between Pred and PTest, and fix bug where this was not covered.

[peterwaller-arm]

Found and fixed additional bug here.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1313–1321

Opted to leave the code as is for now and change it in next patch.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1388–1392

I believe this case is handled within areCFlagsAccessedBetweenInstrs, making this redundant code that can be removed?

[peterwaller-arm]

Drop redundant same-bb check.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

bool AArch64InstrInfo::optimizePTestInstr(
   auto *Mask = MRI->getUniqueVRegDef(MaskReg);
   auto *Pred = MRI->getUniqueVRegDef(PredReg);
   auto NewOp = Pred->getOpcode();
   bool OpChanged = false;
 
   unsigned MaskOpcode = Mask->getOpcode();
   unsigned PredOpcode = Pred->getOpcode();
   bool PredIsPTestLike = isPTestLikeOpcode(PredOpcode);
   bool PredIsWhileLike = isWhileOpcode(PredOpcode);
 
   if (isPTrueOpcode(MaskOpcode) && (PredIsPTestLike || PredIsWhileLike)) {
     // For PTEST(PTRUE, OTHER_INST), PTEST is redundant when PTRUE doesn't
     // deactivate any lanes OTHER_INST might set.
     uint64_t MaskElementSize = getElementSizeForOpcode(MaskOpcode);
     uint64_t PredElementSize = getElementSizeForOpcode(PredOpcode);
 
     // Must be an all active predicate of matching element size.
     if ((PredElementSize != MaskElementSize) ||
         (Mask->getOperand(1).getImm() != 31))
       return false;
 
     // Fallthough to simply remove the PTEST.
   } else if ((Mask == Pred) && (PredIsPTestLike || PredIsWhileLike)) {
     // For PTEST(PG, PG), PTEST is redundant when PG is the result of an
     // instruction that sets the flags as PTEST would.
case AArch64::BRKN_PPzP: {
       auto *PredMask = MRI->getUniqueVRegDef(Pred->getOperand(1).getReg());
       if (Mask != PredMask)
         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)
+        return false;
+
+      NewOp = AArch64::RDFFRS_PPz;
+      OpChanged = true;
+      break;
+    }
     default:
       // Bail out if we don't recognize the input
       return false;
     }
   }
 
   const TargetRegisterInfo *TRI = &getRegisterInfo();
 
-  // If the predicate is in a different block (possibly because its been
-  // hoisted out), then assume the flags are set in between statements.
-  if (Pred->getParent() != PTest->getParent())
+  // If another instruction between Pred and PTest accesses flags, don't remove
+  // the ptest or update the earlier instruction to modify them.
+  if (areCFlagsAccessedBetweenInstrs(Pred, PTest, TRI))
     return false;
 
-  // If another instruction between the propagation and test sets the
-  // flags, don't remove the ptest.
-  MachineBasicBlock::iterator I = Pred, E = PTest;
-  ++I; // Skip past the predicate op itself.
-  for (; I != E; ++I) {
-    const MachineInstr &Inst = *I;
-
-    // TODO: If the ptest flags are unused, we could still remove it.
-    if (Inst.modifiesRegister(AArch64::NZCV, TRI))
-      return false;
-  }
-
   // If we pass all the checks, it's safe to remove the PTEST and use the flags
   // as they are prior to PTEST. Sometimes this requires the tested PTEST
   // operand to be replaced with an equivalent instruction that also sets the
   // flags.
   Pred->setDesc(get(NewOp));
   PTest->eraseFromParent();
   if (OpChanged) {
     bool succeeded = UpdateOperandRegClass(*Pred);
     (void)succeeded;
     assert(succeeded && "Operands have incompatible register classes!");
     Pred->addRegisterDefined(AArch64::NZCV, TRI);

llvm/test/CodeGen/AArch64/sve-ptest-removal-rdffr.mir

+# RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve -run-pass=peephole-opt -verify-machineinstrs %s -o - | FileCheck %s
+# Test that RDFFR followed by PTEST is replaced with RDFFRS.
+---
+# CHECK-LABEL: name:{{\s*}} substitute_rdffr_pp_with_rdffrs_pp
+name:                       substitute_rdffr_pp_with_rdffrs_pp
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $ffr, $p0
+    %0:ppr_3b = COPY $p0
+
+    ; CHECK: RDFFRS_PPz
+    ; CHECK-NOT: PTEST
+    %1:ppr_3b = RDFFR_PPz %0:ppr_3b
+    PTEST_PP killed %0:ppr_3b, killed %1:ppr_3b, implicit-def $nzcv
+
+    ; Consume nzcv
+    %2:gpr32 = COPY $wzr
+    %3:gpr32 = CSINCWr killed %2, $wzr, 0, implicit $nzcv
+    $w0 = COPY %3
+    RET_ReallyLR implicit $w0
+...
+---
+# CHECK-LABEL: name:{{\s*}} fail_to_substitute_rdffr_pp_with_rdffrs_pp_differing_mask
+name:                       fail_to_substitute_rdffr_pp_with_rdffrs_pp_differing_mask
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $ffr, $p0, $p1
+    %0:ppr_3b = COPY $p0
+    %1:ppr_3b = COPY $p1
+
+    ; CHECK: RDFFR_PPz
+    ; CHECK: PTEST
+    %2:ppr_3b = RDFFR_PPz %0:ppr_3b
+    PTEST_PP killed %1:ppr_3b, killed %2:ppr_3b, implicit-def $nzcv
+
+    ; Consume nzcv
+    %3:gpr32 = COPY $wzr
+    %4:gpr32 = CSINCWr killed %3, $wzr, 0, implicit $nzcv
+    $w0 = COPY %4
+    RET_ReallyLR implicit $w0
+...
+---
+# CHECK-LABEL: name:{{\s*}} fail_to_substitute_rdffr_pp_with_rdffrs_pp_nzcv_clobbered
+name:                       fail_to_substitute_rdffr_pp_with_rdffrs_pp_nzcv_clobbered
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $ffr, $p0, $x0
+    %0:ppr_3b = COPY $p0
+
+    ; CHECK: RDFFR_PPz
+    ; CHECK-NEXT: ADDSXrr
+    ; CHECK-NEXT: PTEST_PP
+    %1:ppr_3b = RDFFR_PPz %0:ppr_3b
+    ; Clobber nzcv
+    $x0 = ADDSXrr $x0, $x0, implicit-def $nzcv
+    PTEST_PP killed %0:ppr_3b, killed %1:ppr_3b, implicit-def $nzcv
+
+    ; Consume nzcv
+    %2:gpr32 = COPY $wzr
+    %3:gpr32 = CSINCWr killed %2, $wzr, 0, implicit $nzcv
+    $w0 = COPY %3
+    RET_ReallyLR implicit $w0
+...
+---
+# CHECK-LABEL: name:{{\s*}} fail_to_substitute_rdffr_pp_with_rdffrs_pp_nzcv_flags_used_between
+name:                       fail_to_substitute_rdffr_pp_with_rdffrs_pp_nzcv_flags_used_between
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $ffr, $p0, $x0
+    %0:ppr_3b = COPY $p0
+
+    $wzr = SUBSWri $w0, 0, 0, implicit-def $nzcv
+
+    ; CHECK: RDFFR_PPz
+    ; CHECK-NEXT: CSINCWr
+    ; CHECK-NEXT: PTEST_PP
+    %1:ppr_3b = RDFFR_PPz %0:ppr_3b
+    ; Consume nzcv
+    %2:gpr32 = CSINCWr $wzr, $wzr, 0, implicit $nzcv
+    PTEST_PP killed %0:ppr_3b, killed %1:ppr_3b, implicit-def $nzcv
+
+    ; Consume nzcv
+    %3:gpr32 = COPY $wzr
+    %4:gpr32 = CSINCWr killed %3, $wzr, 0, implicit $nzcv
+    $w0 = ORRWrs %4, %2, 1
+    RET_ReallyLR implicit $w0