[AARCH64][RegisterCoalescer] clang miscompiles zero-extension to long long

[simonwallis2]

Summary

Implement AArch64 variant of shouldCoalesce() to detect a known failing case
and prevent the coalescing of a 32-bit copy into a 64-bit sign-extending load.

Do not coalesce in the following case:
COPY where source is bottom 32 bits of a 64-register,
and destination is a 32-bit subregister of a 64-bit register,
ie it causes the rest of the register to be implicitly set to zero.

A mir test has been added.

In the test case, the 32-bit copy implements a 32 to 64 bit zero extension
and relies on the upper 32 bits being zeroed.

Coalescing to the result of the 64-bit load meant overwriting
the upper 32 bits incorrectly when the loaded byte was negative.

[arsenm]

llvm/test/CodeGen/AArch64/zext-reg-coalesce.mir

4

Add a comment explaining the test?

You don't need -O1

5–12

Don't need the IR section

24–37

I think you can trim most of these instructions

[simonwallis2]

Thanks for the feedback arsenm.
I've made all of the changes you suggested,
with 1 exception:
I have left in the IR for the declaration of @c.

llvm/test/CodeGen/AArch64/zext-reg-coalesce.mir

5–12

I left the IR section in place because of @c

[john.brawn]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

735

We should have a comment here explaining why coalescing is incorrect in this case. Also I'm not sure if the SrcRC == DstRC is correct here, as the implicit zeroing on subregister write doesn't depend on what the source register is.

[simonwallis2]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

735

We should have a comment here explaining why coalescing is incorrect in this case.

I will reword and tailor the comment that is in the test case.

735

Also I'm not sure if the SrcRC == DstRC is correct here, as the implicit zeroing on subregister write doesn't depend on what the source register is.

Thanks. I will modify the check. Currently it matches:

undef %4.sub_32:gpr64 = COPY %2.sub_32:gpr64

The 2 register classes do not have to be the same,
But they both have to be a subregister of a 64bit int register type so that this will also match:

undef %40.sub_32:gpr64 = COPY %41.sub_32:gpr64common

So I should include GPR64common.

[simonwallis2]

Added comments to implementation of AArch64RegisterInfo::shouldCoalesce()

Removed test for exact match of both TargetRegisterClass;
Replaced it with check that both classes are a subregister of a 64bit int register class.

[john.brawn]

LGTM.

llvm/lib/Target/AArch64/AArch64RegisterInfo.h

public:
   // Debug information queries.
   Register getFrameRegister(const MachineFunction &MF) const override;
 
   unsigned getRegPressureLimit(const TargetRegisterClass *RC,
                                MachineFunction &MF) const override;
 
   unsigned getLocalAddressRegister(const MachineFunction &MF) const;
   bool regNeedsCFI(unsigned Reg, unsigned &RegToUseForCFI) const;
+
+  /// SrcRC and DstRC will be morphed into NewRC if this returns true
+  bool shouldCoalesce(MachineInstr *MI, const TargetRegisterClass *SrcRC,
+                      unsigned SubReg, const TargetRegisterClass *DstRC,
+                      unsigned DstSubReg, const TargetRegisterClass *NewRC,
+                      LiveIntervals &LIS) const override;
 };
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

unsigned AArch64RegisterInfo::getLocalAddressRegister(
   const MachineFunction &MF) const {
   const auto &MFI = MF.getFrameInfo();
   if (!MF.hasEHFunclets() && !MFI.hasVarSizedObjects())
     return AArch64::SP;
   else if (needsStackRealignment(MF))
     return getBaseRegister();
   return getFrameRegister(MF);
 }
+
+/// SrcRC and DstRC will be morphed into NewRC if this returns true
+bool AArch64RegisterInfo::shouldCoalesce(
+    MachineInstr *MI, const TargetRegisterClass *SrcRC, unsigned SubReg,
+    const TargetRegisterClass *DstRC, unsigned DstSubReg,
+    const TargetRegisterClass *NewRC, LiveIntervals &LIS) const {
+  if (MI->isCopy() &&
+      ((DstRC->getID() == AArch64::GPR64RegClassID) ||
+       (DstRC->getID() == AArch64::GPR64commonRegClassID)) &&
+      MI->getOperand(0).getSubReg() && MI->getOperand(1).getSubReg())
+    // Do not coalesce in the case of a 32-bit subregister copy
+    // which implements a 32 to 64 bit zero extension
+    // which relies on the upper 32 bits being zeroed.
+    return false;
+  return true;
+}

llvm/test/CodeGen/AArch64/zext-reg-coalesce.mir

+# RUN: llc -mtriple=aarch64-arm-none-eabi -o - %s \
+# RUN: -run-pass simple-register-coalescing | FileCheck %s
+
+# In this test case, the 32-bit copy implements a 32 to 64 bit zero extension
+# and relies on the upper 32 bits being zeroed.
+# Coalescing to the result of the 64-bit load meant overwriting
+# the upper 32 bits incorrectly when the loaded byte was negative.
+
+--- |
+  @c = local_unnamed_addr global i8 -1, align 4
+
+  define i64 @bug_e(i32 %i32) local_unnamed_addr {
+  ret i64 0
+  }
+...
+---
+name:            bug_e
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    %1:gpr32 = COPY $w0
+    %2:gpr64common = ADRP target-flags(aarch64-page) @c
+    %3:gpr64 = LDRSBXui %2, target-flags(aarch64-pageoff, aarch64-nc) @c :: (dereferenceable load 1 from @c, align 4)
+    %0:gpr32 = COPY %3.sub_32
+  ; CHECK: {{.*}}.sub_32:gpr64 = COPY {{.*}}.sub_32
+    STRBBui %1, %2, target-flags(aarch64-pageoff, aarch64-nc) @c :: (store 1 into @c, align 4)
+    %8:gpr64all = SUBREG_TO_REG 0, %0, %subreg.sub_32
+    $x0 = COPY %8
+  ; CHECK: $x0 = COPY
+    RET_ReallyLR implicit $x0
+...