- Ensure both performFirstTrueTestVectorCombine & optimizePTestInstr only

kmclaughlin-arm · kmclaughlin-arm · commit e5d23b8215a0 · 2025-09-10T14:15:15.000Z
consider extracts/copies from the first result of whilecc_x2
- Add negative tests
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -20235,7 +20235,7 @@ performFirstTrueTestVectorCombine(SDNode *N,
 
   // Restricted the DAG combine to only cases where we're extracting from a
   // flag-setting operation.
-  if (!isPredicateCCSettingOp(N0))
+  if (!isPredicateCCSettingOp(N0) || N0.getResNo() != 0)
     return SDValue();
 
   // Extracts of lane 0 for SVE can be expressed as PTEST(Op, FIRST) ? 1 : 0
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1488,21 +1488,6 @@ AArch64InstrInfo::canRemovePTestInstr(MachineInstr *PTest, MachineInstr *Mask,
   bool PredIsPTestLike = isPTestLikeOpcode(PredOpcode);
   bool PredIsWhileLike = isWhileOpcode(PredOpcode);
 
-  uint64_t PredEltSize = 0;
-  if (PredIsWhileLike)
-    PredEltSize = getElementSizeForOpcode(PredOpcode);
-
-  if (Pred->isCopy()) {
-    // Instructions which return a multi-vector (e.g. WHILECC_x2) require copies
-    // before the branch to extract each subregister.
-    auto Op = Pred->getOperand(1);
-    if (Op.isReg() && Op.getReg().isVirtual() && Op.getSubReg() != 0) {
-      MachineInstr *DefMI = MRI->getVRegDef(Op.getReg());
-      PredIsWhileLike = isWhileOpcode(DefMI->getOpcode());
-      PredEltSize = getElementSizeForOpcode(DefMI->getOpcode());
-    }
-  }
-
   if (PredIsWhileLike) {
     // For PTEST(PG, PG), PTEST is redundant when PG is the result of a WHILEcc
     // instruction and the condition is "any" since WHILcc does an implicit
@@ -1514,7 +1499,8 @@ AArch64InstrInfo::canRemovePTestInstr(MachineInstr *PTest, MachineInstr *Mask,
     // redundant since WHILE performs an implicit PTEST with an all active
     // mask.
     if (isPTrueOpcode(MaskOpcode) && Mask->getOperand(1).getImm() == 31 &&
-        getElementSizeForOpcode(MaskOpcode) == PredEltSize)
+        getElementSizeForOpcode(MaskOpcode) ==
+        getElementSizeForOpcode(PredOpcode))
       return PredOpcode;
 
     return {};
@@ -1627,6 +1613,15 @@ bool AArch64InstrInfo::optimizePTestInstr(
     const MachineRegisterInfo *MRI) const {
   auto *Mask = MRI->getUniqueVRegDef(MaskReg);
   auto *Pred = MRI->getUniqueVRegDef(PredReg);
+
+  if (Pred->isCopy()) {
+    // Instructions which return a multi-vector (e.g. WHILECC_x2) require copies
+    // before the branch to extract each subregister.
+    auto Op = Pred->getOperand(1);
+    if (Op.isReg() && Op.getSubReg() == AArch64::psub0)
+      Pred = MRI->getUniqueVRegDef(Op.getReg());
+  }
+
   unsigned PredOpcode = Pred->getOpcode();
   auto NewOp = canRemovePTestInstr(PTest, Mask, Pred, MRI);
   if (!NewOp)
diff --git a/llvm/test/CodeGen/AArch64/sve-cmp-folds.ll b/llvm/test/CodeGen/AArch64/sve-cmp-folds.ll
@@ -316,6 +316,21 @@ define i1 @whilelt_x2_first(i64 %next, i64 %end) {
   ret i1 %bit
 }
 
+; Do not combine to ptest when the extract is not from the first vector result
+define i1 @whilege_x2_second_result(i64 %next, i64 %end) {
+; CHECK-LABEL: whilege_x2_second_result:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilege { p0.s, p1.s }, x0, x1
+; CHECK-NEXT:    mov z0.s, p1/z, #1 // =0x1
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %predpair = call { <vscale x 4 x i1>, <vscale x 4 x i1> } @llvm.aarch64.sve.whilege.x2.nxv4i1.i64(i64 %next, i64 %end)
+  %predicate = extractvalue { <vscale x 4 x i1>, <vscale x 4 x i1> } %predpair, 1
+  %bit = extractelement <vscale x 4 x i1> %predicate, i64 0
+  ret i1 %bit
+}
+
 declare i64 @llvm.vscale.i64()
 declare <vscale x 4 x i1> @llvm.aarch64.sve.whilege.nxv4i1.i64(i64, i64)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.whilegt.nxv4i1.i64(i64, i64)
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilege.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilege.mir
@@ -587,3 +587,54 @@ body:             |
     $w0 = COPY %7
     RET_ReallyLR implicit $w0
 ...
+
+# PTEST is not redundant when it's Pg operand is a subregister copy, but not
+# from the first subregister of ppr2mul2
+---
+name:            whilege_x2_b64_s64_psub1
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: ppr }
+  - { id: 3, class: ppr2mul2 }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+  - { id: 6, class: gpr32 }
+  - { id: 7, class: gpr32 }
+liveins:
+  - { reg: '$x0', virtual-reg: '%0' }
+  - { reg: '$x1', virtual-reg: '%1' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: whilege_x2_b64_s64_psub1
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[PTRUE_D:%[0-9]+]]:ppr = PTRUE_D 31, implicit $vg
+    ; CHECK-NEXT: [[WHILEGE_2PXX_D:%[0-9]+]]:ppr2mul2 = WHILEGE_2PXX_D [[COPY]], [[COPY1]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:ppr = COPY [[WHILEGE_2PXX_D]].psub0
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:ppr = COPY [[WHILEGE_2PXX_D]].psub1
+    ; CHECK-NEXT: PTEST_PP killed [[PTRUE_D]], killed [[COPY3]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:gpr32 = COPY $wzr
+    ; CHECK-NEXT: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr [[COPY4]], $wzr, 0, implicit $nzcv
+    ; CHECK-NEXT: $w0 = COPY [[CSINCWr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %2:ppr = PTRUE_D 31, implicit $vg
+    %3:ppr2mul2 = WHILEGE_2PXX_D %0, %1, implicit-def $nzcv
+    %4:ppr = COPY %3.psub0
+    %5:ppr = COPY %3.psub1
+    PTEST_PP killed %2, killed %5, implicit-def $nzcv
+    %6:gpr32 = COPY $wzr
+    %7:gpr32 = CSINCWr %6, $wzr, 0, implicit $nzcv
+    $w0 = COPY %7
+    RET_ReallyLR implicit $w0
+...
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilegt.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilegt.mir
@@ -627,3 +627,54 @@ body:             |
     $w0 = COPY %7
     RET_ReallyLR implicit $w0
 ...
+
+# PTEST is not redundant when it's Pg operand is a subregister copy, but not
+# from the first subregister of ppr2mul2
+---
+name:            whilegt_x2_b64_s64_psub1
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: ppr }
+  - { id: 3, class: ppr2mul2 }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+  - { id: 6, class: gpr32 }
+  - { id: 7, class: gpr32 }
+liveins:
+  - { reg: '$x0', virtual-reg: '%0' }
+  - { reg: '$x1', virtual-reg: '%1' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: whilegt_x2_b64_s64_psub1
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[PTRUE_D:%[0-9]+]]:ppr = PTRUE_D 31, implicit $vg
+    ; CHECK-NEXT: [[WHILEGT_2PXX_D:%[0-9]+]]:ppr2mul2 = WHILEGT_2PXX_D [[COPY]], [[COPY1]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:ppr = COPY [[WHILEGT_2PXX_D]].psub0
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:ppr = COPY [[WHILEGT_2PXX_D]].psub1
+    ; CHECK-NEXT: PTEST_PP killed [[PTRUE_D]], killed [[COPY3]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:gpr32 = COPY $wzr
+    ; CHECK-NEXT: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr [[COPY4]], $wzr, 0, implicit $nzcv
+    ; CHECK-NEXT: $w0 = COPY [[CSINCWr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %2:ppr = PTRUE_D 31, implicit $vg
+    %3:ppr2mul2 = WHILEGT_2PXX_D %0, %1, implicit-def $nzcv
+    %4:ppr = COPY %3.psub0
+    %5:ppr = COPY %3.psub1
+    PTEST_PP killed %2, killed %5, implicit-def $nzcv
+    %6:gpr32 = COPY $wzr
+    %7:gpr32 = CSINCWr %6, $wzr, 0, implicit $nzcv
+    $w0 = COPY %7
+    RET_ReallyLR implicit $w0
+...
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilehi.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilehi.mir
@@ -587,3 +587,54 @@ body:             |
     $w0 = COPY %7
     RET_ReallyLR implicit $w0
 ...
+
+# PTEST is not redundant when it's Pg operand is a subregister copy, but not
+# from the first subregister of ppr2mul2
+---
+name:            whilehi_x2_b64_s64_psub1
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: ppr }
+  - { id: 3, class: ppr2mul2 }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+  - { id: 6, class: gpr32 }
+  - { id: 7, class: gpr32 }
+liveins:
+  - { reg: '$x0', virtual-reg: '%0' }
+  - { reg: '$x1', virtual-reg: '%1' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: whilehi_x2_b64_s64_psub1
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[PTRUE_D:%[0-9]+]]:ppr = PTRUE_D 31, implicit $vg
+    ; CHECK-NEXT: [[WHILEHI_2PXX_D:%[0-9]+]]:ppr2mul2 = WHILEHI_2PXX_D [[COPY]], [[COPY1]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:ppr = COPY [[WHILEHI_2PXX_D]].psub0
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:ppr = COPY [[WHILEHI_2PXX_D]].psub1
+    ; CHECK-NEXT: PTEST_PP killed [[PTRUE_D]], killed [[COPY3]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:gpr32 = COPY $wzr
+    ; CHECK-NEXT: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr [[COPY4]], $wzr, 0, implicit $nzcv
+    ; CHECK-NEXT: $w0 = COPY [[CSINCWr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %2:ppr = PTRUE_D 31, implicit $vg
+    %3:ppr2mul2 = WHILEHI_2PXX_D %0, %1, implicit-def $nzcv
+    %4:ppr = COPY %3.psub0
+    %5:ppr = COPY %3.psub1
+    PTEST_PP killed %2, killed %5, implicit-def $nzcv
+    %6:gpr32 = COPY $wzr
+    %7:gpr32 = CSINCWr %6, $wzr, 0, implicit $nzcv
+    $w0 = COPY %7
+    RET_ReallyLR implicit $w0
+...
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilehs.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilehs.mir
@@ -587,3 +587,54 @@ body:             |
     $w0 = COPY %7
     RET_ReallyLR implicit $w0
 ...
+
+# PTEST is not redundant when it's Pg operand is a subregister copy, but not
+# from the first subregister of ppr2mul2
+---
+name:            whilehs_x2_b64_s64_psub1
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: ppr }
+  - { id: 3, class: ppr2mul2 }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+  - { id: 6, class: gpr32 }
+  - { id: 7, class: gpr32 }
+liveins:
+  - { reg: '$x0', virtual-reg: '%0' }
+  - { reg: '$x1', virtual-reg: '%1' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: whilehs_x2_b64_s64_psub1
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[PTRUE_D:%[0-9]+]]:ppr = PTRUE_D 31, implicit $vg
+    ; CHECK-NEXT: [[WHILEHS_2PXX_D:%[0-9]+]]:ppr2mul2 = WHILEHS_2PXX_D [[COPY]], [[COPY1]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:ppr = COPY [[WHILEHS_2PXX_D]].psub0
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:ppr = COPY [[WHILEHS_2PXX_D]].psub1
+    ; CHECK-NEXT: PTEST_PP killed [[PTRUE_D]], killed [[COPY3]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:gpr32 = COPY $wzr
+    ; CHECK-NEXT: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr [[COPY4]], $wzr, 0, implicit $nzcv
+    ; CHECK-NEXT: $w0 = COPY [[CSINCWr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %2:ppr = PTRUE_D 31, implicit $vg
+    %3:ppr2mul2 = WHILEHS_2PXX_D %0, %1, implicit-def $nzcv
+    %4:ppr = COPY %3.psub0
+    %5:ppr = COPY %3.psub1
+    PTEST_PP killed %2, killed %5, implicit-def $nzcv
+    %6:gpr32 = COPY $wzr
+    %7:gpr32 = CSINCWr %6, $wzr, 0, implicit $nzcv
+    $w0 = COPY %7
+    RET_ReallyLR implicit $w0
+...
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilele.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilele.mir
@@ -587,3 +587,54 @@ body:             |
     $w0 = COPY %7
     RET_ReallyLR implicit $w0
 ...
+
+# PTEST is not redundant when it's Pg operand is a subregister copy, but not
+# from the first subregister of ppr2mul2
+---
+name:            whilele_x2_b64_s64_psub1
+alignment:       2
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: ppr }
+  - { id: 3, class: ppr2mul2 }
+  - { id: 4, class: ppr }
+  - { id: 5, class: ppr }
+  - { id: 6, class: gpr32 }
+  - { id: 7, class: gpr32 }
+liveins:
+  - { reg: '$x0', virtual-reg: '%0' }
+  - { reg: '$x1', virtual-reg: '%1' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: whilele_x2_b64_s64_psub1
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[PTRUE_D:%[0-9]+]]:ppr = PTRUE_D 31, implicit $vg
+    ; CHECK-NEXT: [[WHILELE_2PXX_D:%[0-9]+]]:ppr2mul2 = WHILELE_2PXX_D [[COPY]], [[COPY1]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:ppr = COPY [[WHILELE_2PXX_D]].psub0
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:ppr = COPY [[WHILELE_2PXX_D]].psub1
+    ; CHECK-NEXT: PTEST_PP killed [[PTRUE_D]], killed [[COPY3]], implicit-def $nzcv
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:gpr32 = COPY $wzr
+    ; CHECK-NEXT: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr [[COPY4]], $wzr, 0, implicit $nzcv
+    ; CHECK-NEXT: $w0 = COPY [[CSINCWr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %2:ppr = PTRUE_D 31, implicit $vg
+    %3:ppr2mul2 = WHILELE_2PXX_D %0, %1, implicit-def $nzcv
+    %4:ppr = COPY %3.psub0
+    %5:ppr = COPY %3.psub1
+    PTEST_PP killed %2, killed %5, implicit-def $nzcv
+    %6:gpr32 = COPY $wzr
+    %7:gpr32 = CSINCWr %6, $wzr, 0, implicit $nzcv
+    $w0 = COPY %7
+    RET_ReallyLR implicit $w0
+...
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilelo.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilelo.mir
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilels.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilels.mir
diff --git a/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilelt.mir b/llvm/test/CodeGen/AArch64/sve-ptest-removal-whilelt.mir
