From 1cb5f9c5d7944a781f94e44b6094ba609f7a680f Mon Sep 17 00:00:00 2001
From: David Green <david.green@arm.com>
Date: Mon, 26 Jul 2021 14:26:38 +0100
Subject: [PATCH] [ARM] Ensure correct regclass in distributing postinc

The register class required for some MVE loads/stores is more
constrained than the register we use when creating postinc. Make sure we
constrain the register class to keep the code correct.
---
 lib/Target/ARM/ARMLoadStoreOptimizer.cpp      | 13 +++++--
 .../CodeGen/Thumb2/mve-postinc-distribute.mir | 35 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/lib/Target/ARM/ARMLoadStoreOptimizer.cpp b/lib/Target/ARM/ARMLoadStoreOptimizer.cpp
index 28fe01c4a57..fd06bfdf352 100644
--- a/lib/Target/ARM/ARMLoadStoreOptimizer.cpp
+++ b/lib/Target/ARM/ARMLoadStoreOptimizer.cpp
@@ -2726,9 +2726,18 @@ static bool isLegalOrConvertableAddressImm(unsigned Opcode, int Imm,
 // by -Offset. This can either happen in-place or be a replacement as MI is
 // converted to another instruction type.
 static void AdjustBaseAndOffset(MachineInstr *MI, Register NewBaseReg,
-                                int Offset, const TargetInstrInfo *TII) {
+                                int Offset, const TargetInstrInfo *TII,
+                                const TargetRegisterInfo *TRI) {
+  // Set the Base reg
   unsigned BaseOp = getBaseOperandIndex(*MI);
   MI->getOperand(BaseOp).setReg(NewBaseReg);
+  // and constrain the reg class to that required by the instruction.
+  MachineFunction *MF = MI->getMF();
+  MachineRegisterInfo &MRI = MF->getRegInfo();
+  const MCInstrDesc &MCID = TII->get(MI->getOpcode());
+  const TargetRegisterClass *TRC = TII->getRegClass(MCID, BaseOp, TRI, *MF);
+  MRI.constrainRegClass(NewBaseReg, TRC);
+
   int OldOffset = MI->getOperand(BaseOp + 1).getImm();
   if (isLegalAddressImm(MI->getOpcode(), OldOffset - Offset, TII))
     MI->getOperand(BaseOp + 1).setImm(OldOffset - Offset);
@@ -2971,7 +2980,7 @@ bool ARMPreAllocLoadStoreOpt::DistributeIncrements(Register Base) {
 
   for (auto *Use : SuccessorAccesses) {
     LLVM_DEBUG(dbgs() << "Changing: "; Use->dump());
-    AdjustBaseAndOffset(Use, NewBaseReg, IncrementOffset, TII);
+    AdjustBaseAndOffset(Use, NewBaseReg, IncrementOffset, TII, TRI);
     LLVM_DEBUG(dbgs() << "  To    : "; Use->dump());
   }
 
diff --git a/test/CodeGen/Thumb2/mve-postinc-distribute.mir b/test/CodeGen/Thumb2/mve-postinc-distribute.mir
index ae488c9f3fb..f2b0510957e 100644
--- a/test/CodeGen/Thumb2/mve-postinc-distribute.mir
+++ b/test/CodeGen/Thumb2/mve-postinc-distribute.mir
@@ -1,5 +1,5 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+mve.fp -run-pass arm-prera-ldst-opt %s -o - | FileCheck %s
+# RUN: llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+mve.fp -run-pass arm-prera-ldst-opt %s -o - -verify-machineinstrs | FileCheck %s
 
 --- |
   define i32* @MVE_VLDRWU32(i32* %x) { unreachable }
@@ -69,6 +69,7 @@
   define i32* @multiple4(i32* %x) { unreachable }
   define i32* @badScale2(i32* %x) { unreachable }
   define i32* @badRange2(i32* %x) { unreachable }
+  define i32* @regtype(i32* %x) { unreachable }
 
 ...
 ---
@@ -1925,3 +1926,35 @@ body:             |
     tBX_RET 14, $noreg, implicit $r0
 
 ...
+---
+name:            regtype
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: tgpr, preferred-register: '' }
+  - { id: 1, class: mqpr, preferred-register: '' }
+  - { id: 2, class: rgpr, preferred-register: '' }
+  - { id: 3, class: rgpr, preferred-register: '' }
+liveins:
+  - { reg: '$r0', virtual-reg: '%0' }
+  - { reg: '$q0', virtual-reg: '%1' }
+body:             |
+  bb.0:
+    liveins: $r0, $q0
+
+    ; CHECK-LABEL: name: regtype
+    ; CHECK: liveins: $r0, $q0
+    ; CHECK: [[COPY:%[0-9]+]]:mqpr = COPY $q0
+    ; CHECK: [[COPY1:%[0-9]+]]:tgpr = COPY $r0
+    ; CHECK: [[t2LDRB_POST:%[0-9]+]]:rgpr, [[t2LDRB_POST1:%[0-9]+]]:tgpr = t2LDRB_POST [[COPY1]], 32, 14 /* CC::al */, $noreg :: (load (s8), align 2)
+    ; CHECK: MVE_VSTRB16 [[COPY]], [[t2LDRB_POST1]], -22, 0, $noreg :: (store (s128), align 8)
+    ; CHECK: $r0 = COPY [[t2LDRB_POST1]]
+    ; CHECK: tBX_RET 14 /* CC::al */, $noreg, implicit $r0
+    %1:mqpr = COPY $q0
+    %0:tgpr = COPY $r0
+    %2:rgpr = t2LDRBi12 %0:tgpr, 0, 14, $noreg :: (load (s8), align 2)
+    MVE_VSTRB16 %1, %0, 10, 0, $noreg :: (store (s128), align 8)
+    %3:rgpr = nuw t2ADDri %0, 32, 14, $noreg, $noreg
+    $r0 = COPY %3
+    tBX_RET 14, $noreg, implicit $r0
+
+...