[AArch64][GlobalISel] Don't bail out of the select(cmp(a, b)) -> csel optimization with multiple users.

It can still be beneficial to do the optimization if the result of the compare
is used by *another* select.

Differential Revision: https://reviews.llvm.org/D73511

1 files changed, 11 insertions, 2 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
index a582df9..a943d25 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -3508,8 +3508,17 @@ bool AArch64InstructionSelector::tryOptSelect(MachineInstr &I) const {
   MachineInstr *CondDef = MRI.getVRegDef(I.getOperand(1).getReg());
   while (CondDef) {
     // We can only fold if all of the defs have one use.
-    if (!MRI.hasOneUse(CondDef->getOperand(0).getReg()))
-      return false;
+    Register CondDefReg = CondDef->getOperand(0).getReg();
+    if (!MRI.hasOneUse(CondDefReg)) {
+      // Unless it's another select.
+      for (auto UI = MRI.use_instr_begin(CondDefReg), UE = MRI.use_instr_end();
+           UI != UE; ++UI) {
+        if (CondDef == &*UI)
+          continue;
+        if (UI->getOpcode() != TargetOpcode::G_SELECT)
+          return false;
+      }
+    }
 
     // We can skip over G_TRUNC since the condition is 1-bit.
     // Truncating/extending can have no impact on the value.