From 018a9641ff1a488f3ffbe1251666696fa5b29915 Mon Sep 17 00:00:00 2001
From: Amara Emerson <amara@apple.com>
Date: Fri, 28 May 2021 15:33:02 -0700
Subject: [PATCH] [AArch64][GlobalISel] Fix a crash during selection of a
 G_ZEXT(s8 = G_LOAD)

We have special handling for a zext of a load <32b because the load does a zext
for free. In that case, we just select the G_ZEXT as if it were a copy but this
triggered the copy checking code to balk at the mismatched size.

This was being hidden because normally these get combined into G_ZEXTLOAD but
for atomics this doesn't happen. The test case here just uses a normal load
because the particular atomic isn't supported yet anyway.
---
 .../GISel/AArch64InstructionSelector.cpp      |  9 +++++
 .../GlobalISel/select-zext-as-copy.mir        | 35 +++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/select-zext-as-copy.mir

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index 657e2e653c13..6c66f47629b0 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -958,6 +958,15 @@ static bool selectCopy(MachineInstr &I, const TargetInstrInfo &TII,
                       << " operand\n");
     return false;
   }
+
+  // If this a GPR ZEXT that we want to just reduce down into a copy.
+  // The sizes will be mismatched with the source < 32b but that's ok.
+  if (I.getOpcode() == TargetOpcode::G_ZEXT) {
+    I.setDesc(TII.get(AArch64::COPY));
+    assert(SrcRegBank.getID() == AArch64::GPRRegBankID);
+    return selectCopy(I, TII, MRI, TRI, RBI);
+  }
+
   I.setDesc(TII.get(AArch64::COPY));
   return CheckCopy();
 }
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-zext-as-copy.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-zext-as-copy.mir
new file mode 100644
index 000000000000..53d0147f57fe
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-zext-as-copy.mir
@@ -0,0 +1,35 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-- -run-pass=instruction-select -verify-machineinstrs -global-isel-abort=1 %s -o - | FileCheck %s
+---
+name:            zext_of_load_copy
+alignment:       4
+exposesReturnsTwice: false
+legalized:       true
+regBankSelected: true
+liveins:
+  - { reg: '$x0', virtual-reg: '' }
+  - { reg: '$x1', virtual-reg: '' }
+fixedStack:      []
+stack:           []
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: zext_of_load_copy
+    ; CHECK: [[DEF:%[0-9]+]]:gpr64common = IMPLICIT_DEF
+    ; CHECK: [[LDRBBui:%[0-9]+]]:gpr32 = LDRBBui [[DEF]], 0 :: (load 1)
+    ; CHECK: [[SUBREG_TO_REG:%[0-9]+]]:gpr64all = SUBREG_TO_REG 0, [[LDRBBui]], %subreg.sub_32
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64 = COPY [[SUBREG_TO_REG]]
+    ; CHECK: [[ANDXri:%[0-9]+]]:gpr64sp = ANDXri [[COPY]], 4096
+    ; CHECK: $x0 = COPY [[ANDXri]]
+    %3:gpr(p0) = G_IMPLICIT_DEF
+    %2:gpr(s8) = G_LOAD %3(p0) :: (load 1)
+    %4:gpr(s64) = G_ZEXT %2(s8)
+    %5:gpr(s64) = G_CONSTANT i64 1
+    %6:gpr(s64) = G_AND %4, %5
+    $x0 = COPY %6
+...