[llvm] [AArch64] InsertSelect: Fold both diamond arms into chained CSEL family instructions (PR #212033)

Mugundan S via llvm-commits llvm-commits at lists.llvm.org
Sat Jul 25 09:32:29 PDT 2026


https://github.com/MGN-GIT updated https://github.com/llvm/llvm-project/pull/212033

>From a51a9cdb8d85faa4f1dd91aeb76c2ea78457730f Mon Sep 17 00:00:00 2001
From: Greenie0701 <smugundan12a at gmail.com>
Date: Sat, 25 Jul 2026 21:38:02 +0530
Subject: [PATCH] [AArch64] InsertSelect: fold both diamond arms into chained
 CSEL family instructions

---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  |  52 +++--
 llvm/lib/Transforms/Vectorize/VPlan.cpp       |   9 +-
 llvm/lib/Transforms/Vectorize/VPlan.h         |  14 ++
 llvm/lib/Transforms/Vectorize/VPlanCFG.h      |  11 +
 .../early-ifcvt-insert-select-fold.mir        | 202 ++++++++++++++++++
 5 files changed, 268 insertions(+), 20 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/early-ifcvt-insert-select-fold.mir

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 6b640727c82c9..ff2db343c6fdf 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1039,7 +1039,6 @@ bool AArch64InstrInfo::canInsertSelect(const MachineBasicBlock &MBB,
   unsigned ExtraCondLat = Cond.size() != 1;
 
   // GPRs are handled by csel.
-  // FIXME: Fold in x+1, -x, and ~x when applicable.
   if (AArch64::GPR64allRegClass.hasSubClassEq(RC) ||
       AArch64::GPR32allRegClass.hasSubClassEq(RC)) {
     // Single-cycle csel, csinc, csinv, and csneg.
@@ -1047,7 +1046,7 @@ bool AArch64InstrInfo::canInsertSelect(const MachineBasicBlock &MBB,
     TrueCycles = FalseCycles = 1;
     if (canFoldIntoCSel(MRI, TrueReg))
       TrueCycles = 0;
-    else if (canFoldIntoCSel(MRI, FalseReg))
+    if (canFoldIntoCSel(MRI, FalseReg))
       FalseCycles = 0;
     return true;
   }
@@ -1291,22 +1290,39 @@ void AArch64InstrInfo::insertSelect(MachineBasicBlock &MBB,
 
   // Try folding simple instructions into the csel.
   if (TryFold) {
-    unsigned NewReg = 0;
-    unsigned FoldedOpc = canFoldIntoCSel(MRI, TrueReg, &NewReg);
-    if (FoldedOpc) {
-      // The folded opcodes csinc, csinc and csneg apply the operation to
-      // FalseReg, so we need to invert the condition.
-      CC = AArch64CC::getInvertedCondCode(CC);
-      TrueReg = FalseReg;
-    } else
-      FoldedOpc = canFoldIntoCSel(MRI, FalseReg, &NewReg);
-
-    // Fold the operation. Leave any dead instructions for DCE to clean up.
-    if (FoldedOpc) {
-      FalseReg = NewReg;
-      Opc = FoldedOpc;
-      // Extend the live range of NewReg.
-      MRI.clearKillFlags(NewReg);
+    unsigned NewTrueReg = 0, NewFalseReg = 0;
+    unsigned TrueFoldOpc  = canFoldIntoCSel(MRI, TrueReg,  &NewTrueReg);
+    unsigned FalseFoldOpc = canFoldIntoCSel(MRI, FalseReg, &NewFalseReg);
+
+    if (TrueFoldOpc && FalseFoldOpc) {
+      // Both sides are foldable --> Emit two chained CSEL-family instructions:
+      //   TmpReg = TrueFoldOpc(FalseReg, NewTrueReg, InvCC) --> folds TrueReg
+      //   DstReg = FalseFoldOpc(TmpReg,  NewFalseReg, CC)   --> folds FalseReg
+      Register TmpReg = MRI.createVirtualRegister(RC);
+      AArch64CC::CondCode InvCC = AArch64CC::getInvertedCondCode(CC);
+      BuildMI(MBB, I, DL, get(TrueFoldOpc), TmpReg)
+          .addReg(FalseReg)
+          .addReg(NewTrueReg)
+          .addImm(InvCC);
+      MRI.clearKillFlags(NewTrueReg);
+      // Set up the second instruction to fold FalseReg.
+      TrueReg  = TmpReg;
+      FalseReg = NewFalseReg;
+      Opc      = FalseFoldOpc;
+      MRI.clearKillFlags(NewFalseReg);
+    } else if (TrueFoldOpc) {
+      // Only TrueReg is foldable. The fold operation applies to Rm (slot 2),
+      // so invert the condition and swap Rn/Rm to place it correctly.
+      CC       = AArch64CC::getInvertedCondCode(CC);
+      TrueReg  = FalseReg;
+      FalseReg = NewTrueReg;
+      Opc      = TrueFoldOpc;
+      MRI.clearKillFlags(NewTrueReg);
+    } else if (FalseFoldOpc) {
+      // Only FalseReg is foldable. Direct fold into Rm, no swap needed.
+      FalseReg = NewFalseReg;
+      Opc      = FalseFoldOpc;
+      MRI.clearKillFlags(NewFalseReg);
     }
   }
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index adb2df8baa167..5f51e7638fa6b 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -902,8 +902,10 @@ VPlan::~VPlan() {
       for (unsigned I = 0, E = R.getNumOperands(); I != E; I++)
         R.setOperand(I, &DummyValue);
 
-  for (auto *VPB : CreatedBlocks)
+  for (auto [Idx, VPB] : enumerate(CreatedBlocks)) {
+    assert(VPB->getNumber() == Idx && "block with mismatched number");
     delete VPB;
+  }
   for (VPValue *VPV : getLiveIns())
     delete VPV;
   delete BackedgeTakenCount;
@@ -1293,8 +1295,10 @@ VPlan *VPlan::duplicate() {
   // current to new VPlan.
   unsigned NumBlocksAfterCloning = CreatedBlocks.size();
   for (unsigned I :
-       seq<unsigned>(NumBlocksBeforeCloning, NumBlocksAfterCloning))
+       seq<unsigned>(NumBlocksBeforeCloning, NumBlocksAfterCloning)) {
+    this->CreatedBlocks[I]->setNumber(NewPlan->CreatedBlocks.size());
     NewPlan->CreatedBlocks.push_back(this->CreatedBlocks[I]);
+  }
   CreatedBlocks.truncate(NumBlocksBeforeCloning);
 
   // Update ExitBlocks of the new plan.
@@ -1309,6 +1313,7 @@ VPlan *VPlan::duplicate() {
 
 VPIRBasicBlock *VPlan::createEmptyVPIRBasicBlock(BasicBlock *IRBB) {
   auto *VPIRBB = new VPIRBasicBlock(IRBB);
+  VPIRBB->setNumber(CreatedBlocks.size());
   CreatedBlocks.push_back(VPIRBB);
   return VPIRBB;
 }
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index ff4af5e0ebc9b..b0d012e80b403 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -126,6 +126,9 @@ class LLVM_ABI_FOR_TEST VPBlockBase {
   /// Subclass identifier (for isa/dyn_cast).
   const VPBlockTy SubclassID;
 
+  /// Unique number, used as node number in the dominator tree.
+  unsigned Number;
+
   /// Add \p Successor as the last successor to this block.
   void appendSuccessor(VPBlockBase *Successor) {
     assert(Successor && "Cannot add nullptr successor!");
@@ -350,6 +353,12 @@ class LLVM_ABI_FOR_TEST VPBlockBase {
     return std::distance(Successors.begin(), find(Successors, Succ));
   }
 
+  /// Return the unique number of the block.
+  unsigned getNumber() const { return Number; }
+
+  /// Set the unique number of the block, used for dominator tree.
+  void setNumber(unsigned N) { Number = N; }
+
   /// The method which generates the output IR that correspond to this
   /// VPBlockBase, thereby "executing" the VPlan.
   virtual void execute(VPTransformState *State) = 0;
@@ -5129,6 +5138,7 @@ class VPlan {
   VPBasicBlock *createVPBasicBlock(const Twine &Name,
                                    VPRecipeBase *Recipe = nullptr) {
     auto *VPB = new VPBasicBlock(Name, Recipe);
+    VPB->setNumber(CreatedBlocks.size());
     CreatedBlocks.push_back(VPB);
     return VPB;
   }
@@ -5142,6 +5152,7 @@ class VPlan {
                                   VPBlockBase *Entry = nullptr,
                                   VPBlockBase *Exiting = nullptr) {
     auto *VPB = new VPRegionBlock(CanIVTy, DL, Entry, Exiting, Name);
+    VPB->setNumber(CreatedBlocks.size());
     CreatedBlocks.push_back(VPB);
     return VPB;
   }
@@ -5152,6 +5163,7 @@ class VPlan {
   VPRegionBlock *createReplicateRegion(VPBlockBase *Entry, VPBlockBase *Exiting,
                                        const std::string &Name = "") {
     auto *VPB = new VPRegionBlock(Entry, Exiting, Name);
+    VPB->setNumber(CreatedBlocks.size());
     CreatedBlocks.push_back(VPB);
     return VPB;
   }
@@ -5167,6 +5179,8 @@ class VPlan {
   /// and deleted once the VPlan is destroyed.
   LLVM_ABI_FOR_TEST VPIRBasicBlock *createVPIRBasicBlock(BasicBlock *IRBB);
 
+  unsigned getMaxBlockNumber() const { return CreatedBlocks.size(); }
+
   /// Returns true if the VPlan is based on a loop with an early exit.
   bool hasEarlyExit() const {
     unsigned NumExitPredecessors =
diff --git a/llvm/lib/Transforms/Vectorize/VPlanCFG.h b/llvm/lib/Transforms/Vectorize/VPlanCFG.h
index 3796d12702bb0..6cfca361001e5 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanCFG.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanCFG.h
@@ -310,6 +310,8 @@ template <> struct GraphTraits<VPBlockBase *> {
   static inline ChildIteratorType child_end(NodeRef N) {
     return ChildIteratorType::end(N);
   }
+
+  static unsigned getNumber(NodeRef N) { return N->getNumber(); }
 };
 
 template <> struct GraphTraits<const VPBlockBase *> {
@@ -325,6 +327,8 @@ template <> struct GraphTraits<const VPBlockBase *> {
   static inline ChildIteratorType child_end(NodeRef N) {
     return ChildIteratorType::end(N);
   }
+
+  static unsigned getNumber(NodeRef N) { return N->getNumber(); }
 };
 
 template <> struct GraphTraits<Inverse<VPBlockBase *>> {
@@ -341,6 +345,8 @@ template <> struct GraphTraits<Inverse<VPBlockBase *>> {
   static inline ChildIteratorType child_end(NodeRef N) {
     return ChildIteratorType::end(N);
   }
+
+  static unsigned getNumber(NodeRef N) { return N->getNumber(); }
 };
 
 template <> struct GraphTraits<VPlan *> {
@@ -359,6 +365,11 @@ template <> struct GraphTraits<VPlan *> {
     // matter.
     return nodes_iterator::end(N->getEntry());
   }
+
+  static unsigned getMaxNumber(GraphRef N) { return N->getMaxBlockNumber(); }
+
+  // Nodes are never renumbered.
+  static unsigned getNumberEpoch(GraphRef) { return 0; }
 };
 
 } // namespace llvm
diff --git a/llvm/test/CodeGen/AArch64/early-ifcvt-insert-select-fold.mir b/llvm/test/CodeGen/AArch64/early-ifcvt-insert-select-fold.mir
new file mode 100644
index 0000000000000..9f7f4e6289fe6
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/early-ifcvt-insert-select-fold.mir
@@ -0,0 +1,202 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass=early-ifcvt -verify-machineinstrs %s -o - | FileCheck %s
+
+# Test insertSelect folding in EarlyIfConversion for the both-foldable case
+# (x+1 -> CSINC, -x -> CSNEG, ~x -> CSINV) and single-side-only cases.
+#
+# Positive tests: both sides foldable -> two chained CSELs, no PHI.
+# Negative tests: only one side foldable -> single CSEL-family, no PHI.
+
+--- |
+  define i32 @test_inc_neg(i32 %cond, i32 %a, i32 %b) {
+  entry:
+    br i1 undef, label %btrue, label %bfalse
+  btrue:
+    br label %tail
+  bfalse:
+    br label %tail
+  tail:
+    ret i32 undef
+  }
+  define i32 @test_inc_inv(i32 %cond, i32 %a, i32 %b) {
+  entry:
+    br i1 undef, label %btrue, label %bfalse
+  btrue:
+    br label %tail
+  bfalse:
+    br label %tail
+  tail:
+    ret i32 undef
+  }
+  define i32 @test_only_true_inc(i32 %cond, i32 %a, i32 %b) {
+  entry:
+    br i1 undef, label %btrue, label %bfalse
+  btrue:
+    br label %tail
+  bfalse:
+    br label %tail
+  tail:
+    ret i32 undef
+  }
+  define i32 @test_only_false_neg(i32 %cond, i32 %a, i32 %b) {
+  entry:
+    br i1 undef, label %btrue, label %bfalse
+  btrue:
+    br label %tail
+  bfalse:
+    br label %tail
+  tail:
+    ret i32 undef
+  }
+...
+
+---
+# Positive: both sides foldable.
+# true = a+1 (CSINC), false = -b (CSNEG).
+# Expected: CSINCWr + CSNEGWr chained, no PHI, no plain CSEL.
+name:            test_inc_neg
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: test_inc_neg
+  ; CHECK:      bb.0.entry:
+  ; CHECK-NOT:  PHI
+  ; CHECK-NOT:  CSELWr
+  ; CHECK:      CSINCWr
+  ; CHECK:      CSNEGWr
+  bb.0.entry:
+    successors: %bb.1(0x50000000), %bb.2(0x30000000)
+    liveins: $w0, $w1, $w2
+    %0:gpr32 = COPY $w0
+    %1:gpr32common = COPY $w1
+    %2:gpr32 = COPY $w2
+    CBNZW %0, %bb.1
+    B %bb.2
+
+  bb.1.btrue:
+    successors: %bb.3(0x80000000)
+    %3:gpr32sp = nsw ADDWri %1, 1, 0
+    %6:gpr32 = COPY %3
+    B %bb.3
+
+  bb.2.bfalse:
+    successors: %bb.3(0x80000000)
+    %4:gpr32 = SUBWrr $wzr, %2
+    B %bb.3
+
+  bb.3.tail:
+    %5:gpr32 = PHI %6, %bb.1, %4, %bb.2
+    $w0 = COPY %5
+    RET_ReallyLR implicit $w0
+
+---
+# Positive: both sides foldable.
+# true = a+1 (CSINC), false = ~b (CSINV).
+# Expected: CSINCWr + CSINVWr chained, no PHI, no plain CSEL.
+name:            test_inc_inv
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: test_inc_inv
+  ; CHECK:      bb.0.entry:
+  ; CHECK-NOT:  PHI
+  ; CHECK-NOT:  CSELWr
+  ; CHECK:      CSINCWr
+  ; CHECK:      CSINVWr
+  bb.0.entry:
+    successors: %bb.1(0x50000000), %bb.2(0x30000000)
+    liveins: $w0, $w1, $w2
+    %0:gpr32 = COPY $w0
+    %1:gpr32common = COPY $w1
+    %2:gpr32 = COPY $w2
+    CBNZW %0, %bb.1
+    B %bb.2
+
+  bb.1.btrue:
+    successors: %bb.3(0x80000000)
+    %3:gpr32sp = nsw ADDWri %1, 1, 0
+    %6:gpr32 = COPY %3
+    B %bb.3
+
+  bb.2.bfalse:
+    successors: %bb.3(0x80000000)
+    %4:gpr32 = ORNWrr $wzr, %2
+    B %bb.3
+
+  bb.3.tail:
+    %5:gpr32 = PHI %6, %bb.1, %4, %bb.2
+    $w0 = COPY %5
+    RET_ReallyLR implicit $w0
+
+---
+# Negative: only TrueReg foldable.
+# true = a+1 (CSINC), false = plain reg (not foldable).
+# Expected: single CSINCWr with inverted CC, no PHI, no CSNEGWr/CSINVWr.
+name:            test_only_true_inc
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: test_only_true_inc
+  ; CHECK:      bb.0.entry:
+  ; CHECK-NOT:  PHI
+  ; CHECK-NOT:  CSNEGWr
+  ; CHECK-NOT:  CSINVWr
+  ; CHECK:      CSINCWr
+  bb.0.entry:
+    successors: %bb.1(0x50000000), %bb.2(0x30000000)
+    liveins: $w0, $w1, $w2
+    %0:gpr32 = COPY $w0
+    %1:gpr32common = COPY $w1
+    %2:gpr32 = COPY $w2
+    CBNZW %0, %bb.1
+    B %bb.2
+
+  bb.1.btrue:
+    successors: %bb.3(0x80000000)
+    %3:gpr32sp = nsw ADDWri %1, 1, 0
+    %6:gpr32 = COPY %3
+    B %bb.3
+
+  bb.2.bfalse:
+    successors: %bb.3(0x80000000)
+    %4:gpr32 = COPY %2
+    B %bb.3
+
+  bb.3.tail:
+    %5:gpr32 = PHI %6, %bb.1, %4, %bb.2
+    $w0 = COPY %5
+    RET_ReallyLR implicit $w0
+
+---
+# Negative: only FalseReg foldable.
+# true = plain reg (not foldable), false = -b (CSNEG).
+# Expected: single CSNEGWr, no PHI, no CSINCWr/CSINVWr.
+name:            test_only_false_neg
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: test_only_false_neg
+  ; CHECK:      bb.0.entry:
+  ; CHECK-NOT:  PHI
+  ; CHECK-NOT:  CSINCWr
+  ; CHECK-NOT:  CSINVWr
+  ; CHECK:      CSNEGWr
+  bb.0.entry:
+    successors: %bb.1(0x50000000), %bb.2(0x30000000)
+    liveins: $w0, $w1, $w2
+    %0:gpr32 = COPY $w0
+    %1:gpr32 = COPY $w1
+    %2:gpr32 = COPY $w2
+    CBNZW %0, %bb.1
+    B %bb.2
+
+  bb.1.btrue:
+    successors: %bb.3(0x80000000)
+    %6:gpr32 = COPY %1
+    B %bb.3
+
+  bb.2.bfalse:
+    successors: %bb.3(0x80000000)
+    %4:gpr32 = SUBWrr $wzr, %2
+    B %bb.3
+
+  bb.3.tail:
+    %5:gpr32 = PHI %6, %bb.1, %4, %bb.2
+    $w0 = COPY %5
+    RET_ReallyLR implicit $w0



More information about the llvm-commits mailing list