From 2faffc9daba9b999a91616aee5ebff87ec318dc9 Mon Sep 17 00:00:00 2001
From: Jessica Paquette <jpaquette@apple.com>
Date: Tue, 20 Apr 2021 13:10:30 -0700
Subject: [PATCH] [AArch64][GlobalISel] Fix regbankselect for G_FCMP with
 vector destinations

These should always go to a FPR, since they always use the vector registers.

Differential Revision: https://reviews.llvm.org/D100885
---
 .../AArch64/GISel/AArch64RegisterBankInfo.cpp |  9 +++++--
 .../AArch64/GlobalISel/regbank-fcmp.mir       | 26 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 test/CodeGen/AArch64/GlobalISel/regbank-fcmp.mir

diff --git a/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp b/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp
index 8b488eaa578..84e2a00ce4c 100644
--- a/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp
+++ b/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp
@@ -719,10 +719,15 @@ AArch64RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
       break;
     OpRegBankIdx = {PMI_FirstGPR, PMI_FirstFPR};
     break;
-  case TargetOpcode::G_FCMP:
-    OpRegBankIdx = {PMI_FirstGPR,
+  case TargetOpcode::G_FCMP: {
+    // If the result is a vector, it must use a FPR.
+    AArch64GenRegisterBankInfo::PartialMappingIdx Idx0 =
+        MRI.getType(MI.getOperand(0).getReg()).isVector() ? PMI_FirstFPR
+                                                          : PMI_FirstGPR;
+    OpRegBankIdx = {Idx0,
                     /* Predicate */ PMI_None, PMI_FirstFPR, PMI_FirstFPR};
     break;
+  }
   case TargetOpcode::G_BITCAST:
     // This is going to be a cross register bank copy and this is expensive.
     if (OpRegBankIdx[0] != OpRegBankIdx[1])
diff --git a/test/CodeGen/AArch64/GlobalISel/regbank-fcmp.mir b/test/CodeGen/AArch64/GlobalISel/regbank-fcmp.mir
new file mode 100644
index 00000000000..51c04f6cd39
--- /dev/null
+++ b/test/CodeGen/AArch64/GlobalISel/regbank-fcmp.mir
@@ -0,0 +1,26 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -run-pass=regbankselect -verify-machineinstrs %s -o - | FileCheck %s
+
+...
+---
+name:            vector
+legalized: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+  liveins: $q0, $q1
+    ; Vectors should always end up on a FPR.
+
+    ; CHECK-LABEL: name: vector
+    ; CHECK: liveins: $q0, $q1
+    ; CHECK: %x:fpr(<2 x s64>) = COPY $q0
+    ; CHECK: %y:fpr(<2 x s64>) = COPY $q1
+    ; CHECK: %fcmp:fpr(<2 x s64>) = G_FCMP floatpred(olt), %x(<2 x s64>), %y
+    ; CHECK: $q0 = COPY %fcmp(<2 x s64>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %x:_(<2 x s64>) = COPY $q0
+    %y:_(<2 x s64>) = COPY $q1
+    %fcmp:_(<2 x s64>) = G_FCMP floatpred(olt), %x:_(<2 x s64>), %y:_
+    $q0 = COPY %fcmp
+    RET_ReallyLR implicit $q0
+...