Index: llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
===================================================================
--- llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1331,6 +1331,10 @@
 
     // Fallthough to simply remove the PTEST.
   } else {
+    auto *PredMask = MRI->getUniqueVRegDef(Pred->getOperand(1).getReg());
+    if (Mask != PredMask && (Mask != Pred))
+      return false;
+
     switch (Pred->getOpcode()) {
     case AArch64::BRKB_PPzP:
     case AArch64::BRKPB_PPzPP: {
@@ -1340,41 +1344,28 @@
       // Check to see if our mask is the same as the brkpb's. If
       // not the resulting flag bits may be different and we
       // can't remove the ptest.
-      auto *PredMask = MRI->getUniqueVRegDef(Pred->getOperand(1).getReg());
-      if (Mask != PredMask)
-        return false;
-
       // 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)
-        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;
     }
+
+    OpChanged = true;
   }
 
   const TargetRegisterInfo *TRI = &getRegisterInfo();
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
@@ -44,8 +44,7 @@
 define i32 @brkb2(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a) {
 ; CHECK-LABEL: brkb2:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    brkb p0.b, p0/z, p1.b
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    brkbs p0.b, p0/z, p1.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.brkb.z.nxv16i1(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a)
@@ -57,8 +56,7 @@
 define i32 @brkn2(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a, <vscale x 16 x i1> %b) {
 ; CHECK-LABEL: brkn2:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    brkn p2.b, p0/z, p1.b, p2.b
-; CHECK-NEXT:    ptest p2, p2.b
+; 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)
Index: llvm/test/CodeGen/AArch64/sve-ptest-removal-rdffr.mir
===================================================================
--- llvm/test/CodeGen/AArch64/sve-ptest-removal-rdffr.mir
+++ llvm/test/CodeGen/AArch64/sve-ptest-removal-rdffr.mir
@@ -19,7 +19,7 @@
     %3:gpr32 = CSINCWr killed %2, $wzr, 0, implicit $nzcv
     $w0 = COPY %3
     RET_ReallyLR implicit $w0
-# Mask and predicate are equal.
+# Test mask and predicate are equal.
 ...
 ---
 # CHECK-LABEL: name:{{\s*}} substitute_rdffr_pp_with_rdffrs_pp2
@@ -30,8 +30,8 @@
     liveins: $ffr, $p0
     %0:ppr_3b = COPY $p0
 
-    ; CHECK: RDFFR_PPz
-    ; CHECK: PTEST
+    ; CHECK: RDFFRS_PPz
+    ; CHECK-NOT: PTEST
     %1:ppr_3b = RDFFR_PPz %0:ppr_3b
     PTEST_PP killed %1:ppr_3b, killed %1:ppr_3b, implicit-def $nzcv