From 6ff09ce061df499b518150f33006d1b0dcfdabd7 Mon Sep 17 00:00:00 2001
From: Amara Emerson <amara@apple.com>
Date: Sun, 21 Feb 2021 21:16:20 -0800
Subject: [PATCH] [AArch64][GlobalISel] Fix <16 x s8> G_DUP regbankselect to
 assign source to gpr.

We can only select this type if the source is on GPR, not FPR.
---
 .../AArch64/GISel/AArch64RegisterBankInfo.cpp      |  7 +++++--
 .../CodeGen/AArch64/GlobalISel/regbank-dup.mir     | 23 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp
index 0e12fff..cbc027c 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp
@@ -666,9 +666,12 @@ AArch64RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
   switch (Opc) {
   case AArch64::G_DUP: {
     Register ScalarReg = MI.getOperand(1).getReg();
+    LLT ScalarTy = MRI.getType(ScalarReg);
     auto ScalarDef = MRI.getVRegDef(ScalarReg);
-    if (getRegBank(ScalarReg, MRI, TRI) == &AArch64::FPRRegBank ||
-        onlyDefinesFP(*ScalarDef, MRI, TRI))
+    // s8 is an exception for G_DUP, which we always want on gpr.
+    if (ScalarTy.getSizeInBits() != 8 &&
+        (getRegBank(ScalarReg, MRI, TRI) == &AArch64::FPRRegBank ||
+         onlyDefinesFP(*ScalarDef, MRI, TRI)))
       OpRegBankIdx = {PMI_FirstFPR, PMI_FirstFPR};
     else
       OpRegBankIdx = {PMI_FirstFPR, PMI_FirstGPR};
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/regbank-dup.mir b/llvm/test/CodeGen/AArch64/GlobalISel/regbank-dup.mir
index 7240705..4cd6eef 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/regbank-dup.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/regbank-dup.mir
@@ -152,3 +152,26 @@ body:             |
     RET_ReallyLR implicit $q0
 
 ...
+---
+name:            v416s8_gpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: v416s8_gpr
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr(s32) = COPY $w0
+    ; CHECK: %trunc:gpr(s8) = G_TRUNC [[COPY]](s32)
+    ; CHECK: [[DUP:%[0-9]+]]:fpr(<16 x s8>) = G_DUP %trunc(s8)
+    ; CHECK: $q0 = COPY [[DUP]](<16 x s8>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(s32) = COPY $w0
+    %trunc:_(s8) = G_TRUNC %0(s32)
+    %1:_(<16 x s8>) = G_DUP %trunc(s8)
+    $q0 = COPY %1(<16 x s8>)
+    RET_ReallyLR implicit $q0
+
+...
-- 
2.7.4