AArch64/GlobalISel: Preserve memory types

[paquette]

LGTM

[arsenm]

30fa074c0a6e38bc605cab892f6a0474e7e6499f

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

Register getStackAddress(uint64_t Size, int64_t Offset,
     const bool IsImmutable = !Flags.isByVal();
 
     int FI = MFI.CreateFixedObject(Size, Offset, IsImmutable);
     MPO = MachinePointerInfo::getFixedStack(MIRBuilder.getMF(), FI);
     auto AddrReg = MIRBuilder.buildFrameIndex(LLT::pointer(0, 64), FI);
     return AddrReg.getReg(0);
   }
 
-  LLT getStackValueStoreType(const DataLayout &, const CCValAssign &VA,
+  LLT getStackValueStoreType(const DataLayout &DL, const CCValAssign &VA,
                              ISD::ArgFlagsTy Flags) const override {
+    // For pointers, we just need to fixup the integer types reported in the
+    // CCValAssign.
+    if (Flags.isPointer())
+      return CallLowering::ValueHandler::getStackValueStoreType(DL, VA, Flags);
     return getStackValueStoreTypeHack(VA);
   }
 
   void assignValueToReg(Register ValVReg, Register PhysReg,
                         CCValAssign &VA) override {
     markPhysRegUsed(PhysReg);
     IncomingValueHandler::assignValueToReg(ValVReg, PhysReg, VA);
   }
 
   void assignValueToAddress(Register ValVReg, Register Addr, LLT MemTy,
                             MachinePointerInfo &MPO, CCValAssign &VA) override {
     MachineFunction &MF = MIRBuilder.getMF();
 
     LLT ValTy(VA.getValVT());
     LLT LocTy(VA.getLocVT());
 
     // Fixup the types for the DAG compatibility hack.
     if (VA.getValVT() == MVT::i8 || VA.getValVT() == MVT::i16)
       std::swap(ValTy, LocTy);
+    else {
+      // The calling code knows if this is a pointer or not, we're only touching
+      // the LocTy for the i8/i16 hack.
+      assert(LocTy.getSizeInBits() == MemTy.getSizeInBits());
+      LocTy = MemTy;
+    }
 
     auto MMO = MF.getMachineMemOperand(
         MPO, MachineMemOperand::MOLoad | MachineMemOperand::MOInvariant, LocTy,
         inferAlignFromPtrInfo(MF, MPO));
     MIRBuilder.buildLoad(ValVReg, Addr, *MMO);
   }
 
   /// How the physical register gets marked varies between formal
Register getStackAddress(uint64_t Size, int64_t Offset,
     MPO = MachinePointerInfo::getStack(MF, Offset);
     return AddrReg.getReg(0);
   }
 
   /// We need to fixup the reported store size for certain value types because
   /// we invert the interpretation of ValVT and LocVT in certain cases. This is
   /// for compatability with the DAG call lowering implementation, which we're
   /// currently building on top of.
-  LLT getStackValueStoreType(const DataLayout &, const CCValAssign &VA,
+  LLT getStackValueStoreType(const DataLayout &DL, const CCValAssign &VA,
                              ISD::ArgFlagsTy Flags) const override {
+    if (Flags.isPointer())
+      return CallLowering::ValueHandler::getStackValueStoreType(DL, VA, Flags);
     return getStackValueStoreTypeHack(VA);
   }
 
   void assignValueToReg(Register ValVReg, Register PhysReg,
                         CCValAssign &VA) override {
     MIB.addUse(PhysReg, RegState::Implicit);
     Register ExtReg = extendRegister(ValVReg, VA);
     MIRBuilder.buildCopy(PhysReg, ExtReg);

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

if (Amount > 31)
     return true; // This will have to remain a register variant.
   auto ExtCst = MIRBuilder.buildConstant(LLT::scalar(64), Amount);
   Observer.changingInstr(MI);
   MI.getOperand(2).setReg(ExtCst.getReg(0));
   Observer.changedInstr(MI);
   return true;
 }
 
+// FIXME: This should be removed and replaced with the generic bitcast legalize
+// action.
 bool AArch64LegalizerInfo::legalizeLoadStore(
     MachineInstr &MI, MachineRegisterInfo &MRI, MachineIRBuilder &MIRBuilder,
     GISelChangeObserver &Observer) const {
   assert(MI.getOpcode() == TargetOpcode::G_STORE ||
          MI.getOpcode() == TargetOpcode::G_LOAD);
   // Here we just try to handle vector loads/stores where our value type might
   // have pointer elements, which the SelectionDAG importer can't handle. To
   // allow the existing patterns for s64 to fire for p0, we just try to bitcast
if (!ValTy.isVector() || !ValTy.getElementType().isPointer() ||
       ValTy.getElementType().getAddressSpace() != 0) {
     LLVM_DEBUG(dbgs() << "Tried to do custom legalization on wrong load/store");
     return false;
   }
 
   unsigned PtrSize = ValTy.getElementType().getSizeInBits();
   const LLT NewTy = LLT::vector(ValTy.getElementCount(), PtrSize);
   auto &MMO = **MI.memoperands_begin();
+  MMO.setType(NewTy);
+
   if (MI.getOpcode() == TargetOpcode::G_STORE) {
     auto Bitcast = MIRBuilder.buildBitcast(NewTy, ValReg);
     MIRBuilder.buildStore(Bitcast.getReg(0), MI.getOperand(1), MMO);
   } else {
     auto NewLoad = MIRBuilder.buildLoad(NewTy, MI.getOperand(1), MMO);
     MIRBuilder.buildBitcast(ValReg, NewLoad);
   }
   MI.eraseFromParent();

llvm/test/CodeGen/AArch64/GlobalISel/call-translator.ll

 ; CHECK-DAG:  - { id: [[STACK0:[0-9]+]], type: default, offset: 0, size: 8,
 ; CHECK-DAG:  - { id: [[STACK8:[0-9]+]], type: default, offset: 8, size: 8,
 ; CHECK-DAG:  - { id: [[STACK16:[0-9]+]], type: default, offset: 16, size: 8,
 ; CHECK: [[LHS_ADDR:%[0-9]+]]:_(p0) = G_FRAME_INDEX %fixed-stack.[[STACK0]]
 ; CHECK: [[LHS:%[0-9]+]]:_(s64) = G_LOAD [[LHS_ADDR]](p0) :: (invariant load (s64) from %fixed-stack.[[STACK0]], align 16)
 ; CHECK: [[RHS_ADDR:%[0-9]+]]:_(p0) = G_FRAME_INDEX %fixed-stack.[[STACK8]]
 ; CHECK: [[RHS:%[0-9]+]]:_(s64) = G_LOAD [[RHS_ADDR]](p0) :: (invariant load (s64) from %fixed-stack.[[STACK8]])
 ; CHECK: [[ADDR_ADDR:%[0-9]+]]:_(p0) = G_FRAME_INDEX %fixed-stack.[[STACK16]]
-; CHECK: [[ADDR:%[0-9]+]]:_(p0) = G_LOAD [[ADDR_ADDR]](p0) :: (invariant load (s64) from %fixed-stack.[[STACK16]], align 16)
+; CHECK: [[ADDR:%[0-9]+]]:_(p0) = G_LOAD [[ADDR_ADDR]](p0) :: (invariant load (p0) from %fixed-stack.[[STACK16]], align 16)
 ; CHECK: [[SUM:%[0-9]+]]:_(s64) = G_ADD [[LHS]], [[RHS]]
 ; CHECK: G_STORE [[SUM]](s64), [[ADDR]](p0)
 define void @test_stack_slots([8 x i64], i64 %lhs, i64 %rhs, i64* %addr) {
   %sum = add i64 %lhs, %rhs
   store i64 %sum, i64* %addr
   ret void
 }
 
 ; CHECK-LABEL: name: test_call_stack
 ; CHECK: [[C42:%[0-9]+]]:_(s64) = G_CONSTANT i64 42
 ; CHECK: [[C12:%[0-9]+]]:_(s64) = G_CONSTANT i64 12
 ; CHECK: [[PTR:%[0-9]+]]:_(p0) = G_CONSTANT i64 0
 ; CHECK: ADJCALLSTACKDOWN 24, 0, implicit-def $sp, implicit $sp
 ; CHECK: [[SP:%[0-9]+]]:_(p0) = COPY $sp
 ; CHECK: [[C42_OFFS:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
 ; CHECK: [[C42_LOC:%[0-9]+]]:_(p0) = G_PTR_ADD [[SP]], [[C42_OFFS]](s64)
 ; CHECK: G_STORE [[C42]](s64), [[C42_LOC]](p0) :: (store (s64) into stack, align 1)
 ; CHECK: [[C12_OFFS:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
 ; CHECK: [[C12_LOC:%[0-9]+]]:_(p0) = G_PTR_ADD [[SP]], [[C12_OFFS]](s64)
 ; CHECK: G_STORE [[C12]](s64), [[C12_LOC]](p0) :: (store (s64) into stack + 8, align 1)
 ; CHECK: [[PTR_OFFS:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
 ; CHECK: [[PTR_LOC:%[0-9]+]]:_(p0) = G_PTR_ADD [[SP]], [[PTR_OFFS]](s64)
-; CHECK: G_STORE [[PTR]](p0), [[PTR_LOC]](p0) :: (store (s64) into stack + 16, align 1)
+; CHECK: G_STORE [[PTR]](p0), [[PTR_LOC]](p0) :: (store (p0) into stack + 16, align 1)
 ; CHECK: BL @test_stack_slots
 ; CHECK: ADJCALLSTACKUP 24, 0, implicit-def $sp, implicit $sp
 define void @test_call_stack() {
   call void @test_stack_slots([8 x i64] undef, i64 42, i64 12, i64* null)
   ret void
 }
 
 ; CHECK-LABEL: name: test_mem_i1

llvm/test/CodeGen/AArch64/GlobalISel/legalize-load-store-vector-of-ptr.mir

body:             |
   bb.1 (%ir-block.0):
     liveins: $q0, $x0
 
     ; CHECK-LABEL: name: store_v2p0
     ; CHECK: liveins: $q0, $x0
     ; CHECK: [[COPY:%[0-9]+]]:_(<2 x p0>) = COPY $q0
     ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x0
     ; CHECK: [[BITCAST:%[0-9]+]]:_(<2 x s64>) = G_BITCAST [[COPY]](<2 x p0>)
-    ; CHECK: G_STORE [[BITCAST]](<2 x s64>), [[COPY1]](p0) :: (store (<2 x p0>) into %ir.ptr)
+    ; CHECK: G_STORE [[BITCAST]](<2 x s64>), [[COPY1]](p0) :: (store (<2 x s64>) into %ir.ptr)
     ; CHECK: RET_ReallyLR
     %0:_(<2 x p0>) = COPY $q0
     %1:_(p0) = COPY $x0
     G_STORE %0(<2 x p0>), %1(p0) :: (store (<2 x p0>) into %ir.ptr)
     RET_ReallyLR
 
 ...
 ---
 name:            load_v2p0
 alignment:       4
 tracksRegLiveness: true
 machineFunctionInfo: {}
 body:             |
   bb.1 (%ir-block.0):
     liveins: $x0
 
     ; CHECK-LABEL: name: load_v2p0
     ; CHECK: liveins: $x0
     ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
-    ; CHECK: [[LOAD:%[0-9]+]]:_(<2 x s64>) = G_LOAD [[COPY]](p0) :: (load (<2 x p0>) from %ir.ptr)
+    ; CHECK: [[LOAD:%[0-9]+]]:_(<2 x s64>) = G_LOAD [[COPY]](p0) :: (load (<2 x s64>) from %ir.ptr)
     ; CHECK: [[BITCAST:%[0-9]+]]:_(<2 x p0>) = G_BITCAST [[LOAD]](<2 x s64>)
     ; CHECK: $q0 = COPY [[BITCAST]](<2 x p0>)
     ; CHECK: RET_ReallyLR implicit $q0
     %0:_(p0) = COPY $x0
     %1:_(<2 x p0>) = G_LOAD %0(p0) :: (load (<2 x p0>) from %ir.ptr)
     $q0 = COPY %1(<2 x p0>)
     RET_ReallyLR implicit $q0