From 14c2cf8e187451b51e997c40476b65d5ef9d346e Mon Sep 17 00:00:00 2001
From: Amara Emerson <aemerson@apple.com>
Date: Mon, 27 Jan 2020 15:57:16 -0800
Subject: [PATCH] [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
---
 .../Target/AArch64/AArch64InstructionSelector.cpp  | 13 ++++++--
 .../CodeGen/AArch64/GlobalISel/fold-fp-select.mir  | 37 +++++++++++++++++++++-
 2 files changed, 47 insertions(+), 3 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.
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/fold-fp-select.mir b/llvm/test/CodeGen/AArch64/GlobalISel/fold-fp-select.mir
index 52383f6..7a34b77 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/fold-fp-select.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/fold-fp-select.mir
@@ -4,7 +4,7 @@
 # Verify the following:
 #
 # - We can fold compares into selects.
-# - This only happens when the result of the compare is only used by the select.
+# - This only happens when the result of the compare is only used by selects.
 #
 # Also verify that, for now:
 #
@@ -49,6 +49,41 @@ body:             |
 
 ...
 ---
+name:            fcmp_more_than_one_select
+alignment:       4
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $s0, $s1, $w1
+
+    ; CHECK-LABEL: name: fcmp_more_than_one_select
+    ; CHECK: liveins: $s0, $s1, $w1
+    ; CHECK: [[COPY:%[0-9]+]]:fpr32 = COPY $s0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr32 = COPY $s1
+    ; CHECK: [[FMOVS0_:%[0-9]+]]:fpr32 = FMOVS0
+    ; CHECK: FCMPSri [[COPY]], implicit-def $nzcv
+    ; CHECK: [[FCSELSrrr:%[0-9]+]]:fpr32 = FCSELSrrr [[FMOVS0_]], [[COPY1]], 0, implicit $nzcv
+    ; CHECK: FCMPSri [[COPY]], implicit-def $nzcv
+    ; CHECK: [[FCSELSrrr1:%[0-9]+]]:fpr32 = FCSELSrrr [[COPY1]], [[FMOVS0_]], 0, implicit $nzcv
+    ; CHECK: $s0 = COPY [[FCSELSrrr]]
+    ; CHECK: $s1 = COPY [[FCSELSrrr1]]
+    ; CHECK: RET_ReallyLR implicit $s0
+    %0:fpr(s32) = COPY $s0
+    %1:fpr(s32) = COPY $s1
+    %2:fpr(s32) = G_FCONSTANT float 0.000000e+00
+    %5:gpr(s32) = G_FCMP floatpred(oeq), %0(s32), %2
+    %3:gpr(s1) = G_TRUNC %5(s32)
+    %6:fpr(s1) = COPY %3(s1)
+    %4:fpr(s32) = G_SELECT %6(s1), %2, %1
+    %7:fpr(s32) = G_SELECT %6(s1), %1, %2
+    $s0 = COPY %4(s32)
+    $s1 = COPY %7(s32)
+    RET_ReallyLR implicit $s0
+
+...
+---
 name:            using_icmp
 alignment:       4
 legalized:       true
-- 
2.7.4