[AArch64][GlobalISel] Add isel support for G_BLOCK_ADDR

[aemerson]

Summary

Add isel support for G_BLOCK_ADDR.

Also refactors some existing code to materialize addresses for the large code model so it can be shared between G_GLOBAL_VALUE and G_BLOCK_ADDR.

This implements PR36390.

lib/Target/AArch64/AArch64InstructionSelector.cpp

ComplexRendererFns selectAddrModeIndexed(MachineOperand &Root,
                                            unsigned Size) const;
   template <int Width>
   ComplexRendererFns selectAddrModeIndexed(MachineOperand &Root) const {
     return selectAddrModeIndexed(Root, Width / 8);
   }
 
   void renderTruncImm(MachineInstrBuilder &MIB, const MachineInstr &MI) const;
 
+  // Materialize a GlobalValue or BlockAddress using a movz+movk sequence.
+  void materializeLargeCMVal(MachineInstr &I, const Value *V,
+                             unsigned char OpFlags) const;
+
   const AArch64TargetMachine &TM;
   const AArch64Subtarget &STI;
   const AArch64InstrInfo &TII;
   const AArch64RegisterInfo &TRI;
   const AArch64RegisterBankInfo &RBI;
 
 #define GET_GLOBALISEL_PREDICATES_DECL
 #include "AArch64GenGlobalISel.inc"
MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(AArch64::STRXui))
             .addImm(0)
             .addMemOperand(*I.memoperands_begin());
 
   constrainSelectedInstRegOperands(*MIB, TII, TRI, RBI);
   I.eraseFromParent();
   return true;
 }
 
+void AArch64InstructionSelector::materializeLargeCMVal(
+    MachineInstr &I, const Value *V, unsigned char OpFlags) const {
+  MachineBasicBlock &MBB = *I.getParent();
+  MachineFunction &MF = *MBB.getParent();
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  unsigned MovZDstReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
+  auto InsertPt = std::next(I.getIterator());
+  auto MovZ = BuildMI(MBB, InsertPt, I.getDebugLoc(), TII.get(AArch64::MOVZXi))
+                  .addDef(MovZDstReg);
+  MovZ->addOperand(MF, I.getOperand(1));
+  MovZ->getOperand(1).setTargetFlags(OpFlags | AArch64II::MO_G0 |
+                                     AArch64II::MO_NC);
+  MovZ->addOperand(MF, MachineOperand::CreateImm(0));
+  constrainSelectedInstRegOperands(*MovZ, TII, TRI, RBI);
+
+  auto BuildMovK = [&](unsigned SrcReg, unsigned char Flags, unsigned Offset,
+                       unsigned ForceDstReg) {
+    unsigned DstReg = ForceDstReg
+                          ? ForceDstReg
+                          : MRI.createVirtualRegister(&AArch64::GPR64RegClass);
+    auto MovI =
+        BuildMI(MBB, InsertPt, MovZ->getDebugLoc(), TII.get(AArch64::MOVKXi))
+            .addDef(DstReg)
+            .addReg(SrcReg);
+    if (auto *GV = dyn_cast<GlobalValue>(V)) {
+      MovI->addOperand(MF, MachineOperand::CreateGA(
+                               GV, MovZ->getOperand(1).getOffset(), Flags));
+    } else {
+      MovI->addOperand(
+          MF, MachineOperand::CreateBA(cast<BlockAddress>(V),
+                                       MovZ->getOperand(1).getOffset(), Flags));
+    }
+    MovI->addOperand(MF, MachineOperand::CreateImm(Offset));
+    constrainSelectedInstRegOperands(*MovI, TII, TRI, RBI);
+    return DstReg;
+  };
+  unsigned DstReg = BuildMovK(MovZ->getOperand(0).getReg(),
+                              AArch64II::MO_G1 | AArch64II::MO_NC, 16, 0);
+  DstReg = BuildMovK(DstReg, AArch64II::MO_G2 | AArch64II::MO_NC, 32, 0);
+  BuildMovK(DstReg, AArch64II::MO_G3, 48, I.getOperand(0).getReg());
+  return;
+}
+
 bool AArch64InstructionSelector::select(MachineInstr &I,
                                         CodeGenCoverage &CoverageInfo) const {
   assert(I.getParent() && "Instruction should be in a basic block!");
   assert(I.getParent()->getParent() && "Instruction should be in a function!");
 
   MachineBasicBlock &MBB = *I.getParent();
   MachineFunction &MF = *MBB.getParent();
   MachineRegisterInfo &MRI = MF.getRegInfo();
if (GV->isThreadLocal()) {
       return false;
     }
     unsigned char OpFlags = STI.ClassifyGlobalReference(GV, TM);
     if (OpFlags & AArch64II::MO_GOT) {
       I.setDesc(TII.get(AArch64::LOADgot));
       I.getOperand(1).setTargetFlags(OpFlags);
     } else if (TM.getCodeModel() == CodeModel::Large) {
       // Materialize the global using movz/movk instructions.
-      unsigned MovZDstReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
-      auto InsertPt = std::next(I.getIterator());
-      auto MovZ =
-          BuildMI(MBB, InsertPt, I.getDebugLoc(), TII.get(AArch64::MOVZXi))
-              .addDef(MovZDstReg);
-      MovZ->addOperand(MF, I.getOperand(1));
-      MovZ->getOperand(1).setTargetFlags(OpFlags | AArch64II::MO_G0 |
-                                         AArch64II::MO_NC);
-      MovZ->addOperand(MF, MachineOperand::CreateImm(0));
-      constrainSelectedInstRegOperands(*MovZ, TII, TRI, RBI);
-
-      auto BuildMovK = [&](unsigned SrcReg, unsigned char Flags,
-                           unsigned Offset, unsigned ForceDstReg) {
-        unsigned DstReg =
-            ForceDstReg ? ForceDstReg
-                        : MRI.createVirtualRegister(&AArch64::GPR64RegClass);
-        auto MovI = BuildMI(MBB, InsertPt, MovZ->getDebugLoc(),
-                            TII.get(AArch64::MOVKXi))
-                        .addDef(DstReg)
-                        .addReg(SrcReg);
-        MovI->addOperand(MF, MachineOperand::CreateGA(
-                                 GV, MovZ->getOperand(1).getOffset(), Flags));
-        MovI->addOperand(MF, MachineOperand::CreateImm(Offset));
-        constrainSelectedInstRegOperands(*MovI, TII, TRI, RBI);
-        return DstReg;
-      };
-      unsigned DstReg = BuildMovK(MovZ->getOperand(0).getReg(),
-                                  AArch64II::MO_G1 | AArch64II::MO_NC, 16, 0);
-      DstReg = BuildMovK(DstReg, AArch64II::MO_G2 | AArch64II::MO_NC, 32, 0);
-      BuildMovK(DstReg, AArch64II::MO_G3, 48, I.getOperand(0).getReg());
+      materializeLargeCMVal(I, GV, OpFlags);
       I.eraseFromParent();
       return true;
     } else {
       I.setDesc(TII.get(AArch64::MOVaddr));
       I.getOperand(1).setTargetFlags(OpFlags | AArch64II::MO_PAGE);
       MachineInstrBuilder MIB(MF, I);
       MIB.addGlobalAddress(GV, I.getOperand(1).getOffset(),
                            OpFlags | AArch64II::MO_PAGEOFF | AArch64II::MO_NC);
case TargetOpcode::G_INTRINSIC_W_SIDE_EFFECTS:
     if (!I.getOperand(0).isIntrinsicID())
       return false;
     if (I.getOperand(0).getIntrinsicID() != Intrinsic::trap)
       return false;
     BuildMI(MBB, I, I.getDebugLoc(), TII.get(AArch64::BRK))
       .addImm(1);
     I.eraseFromParent();
     return true;
-  case TargetOpcode::G_IMPLICIT_DEF:
+  case TargetOpcode::G_IMPLICIT_DEF: {
     I.setDesc(TII.get(TargetOpcode::IMPLICIT_DEF));
     const LLT DstTy = MRI.getType(I.getOperand(0).getReg());
     const unsigned DstReg = I.getOperand(0).getReg();
     const RegisterBank &DstRB = *RBI.getRegBank(DstReg, MRI, TRI);
     const TargetRegisterClass *DstRC =
         getRegClassForTypeOnBank(DstTy, DstRB, RBI);
     RBI.constrainGenericRegister(DstReg, *DstRC, MRI);
     return true;
   }
+  case TargetOpcode::G_BLOCK_ADDR: {
+    if (TM.getCodeModel() == CodeModel::Large) {
+      materializeLargeCMVal(I, I.getOperand(1).getBlockAddress(), 0);
+      I.eraseFromParent();
+      return true;
+    } else {
+      I.setDesc(TII.get(AArch64::MOVaddrBA));
+      auto MovMI = BuildMI(MBB, I, I.getDebugLoc(), TII.get(AArch64::MOVaddrBA),
+                           I.getOperand(0).getReg())
+                       .addBlockAddress(I.getOperand(1).getBlockAddress(),
+                                        /* Offset */ 0, AArch64II::MO_PAGE)
+                       .addBlockAddress(
+                           I.getOperand(1).getBlockAddress(), /* Offset */ 0,
+                           AArch64II::MO_NC | AArch64II::MO_PAGEOFF);
+      I.eraseFromParent();
+      return constrainSelectedInstRegOperands(*MovMI, TII, TRI, RBI);
+    }
+  }
+  }
 
   return false;
 }
 
 /// SelectArithImmed - Select an immediate value that can be represented as
 /// a 12-bit value shifted left by either 0 or 12.  If so, return true with
 /// Val set to the 12-bit value and Shift set to the shifter operand.
 InstructionSelector::ComplexRendererFns

test/CodeGen/AArch64/GlobalISel/select-blockaddress.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-unknown-unknown -o - -verify-machineinstrs -run-pass=instruction-select %s | FileCheck %s
+# RUN: llc -mtriple=aarch64-unknown-unknown -o - -verify-machineinstrs -run-pass=instruction-select -code-model=large %s | FileCheck %s --check-prefix=LARGE
+--- |
+  ; ModuleID = 'blockaddress.ll'
+  source_filename = "blockaddress.ll"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-none-linux-gnu"
+
+  @addr = global i8* null
+
+  define void @test_blockaddress() {
+    store i8* blockaddress(@test_blockaddress, %block), i8** @addr
+    indirectbr i8* blockaddress(@test_blockaddress, %block), [label %block]
+
+  block:                                            ; preds = %0
+    ret void
+  }
+
+...
+---
+name:            test_blockaddress
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+body:             |
+  ; CHECK-LABEL: name: test_blockaddress
+  ; CHECK: bb.0 (%ir-block.0):
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   [[MOVaddrBA:%[0-9]+]]:gpr64 = MOVaddrBA target-flags(aarch64-page) blockaddress(@test_blockaddress, %ir-block.block), target-flags(aarch64-pageoff, aarch64-nc) blockaddress(@test_blockaddress, %ir-block.block)
+  ; CHECK:   [[MOVaddr:%[0-9]+]]:gpr64common = MOVaddr target-flags(aarch64-page) @addr, target-flags(aarch64-pageoff, aarch64-nc) @addr
+  ; CHECK:   STRXui [[MOVaddrBA]], [[MOVaddr]], 0 :: (store 8 into @addr)
+  ; CHECK:   BR [[MOVaddrBA]]
+  ; CHECK: bb.1.block (address-taken):
+  ; CHECK:   RET_ReallyLR
+  ; LARGE-LABEL: name: test_blockaddress
+  ; LARGE: bb.0 (%ir-block.0):
+  ; LARGE:   successors: %bb.1(0x80000000)
+  ; LARGE:   [[MOVZXi:%[0-9]+]]:gpr64 = MOVZXi target-flags(aarch64-g0, aarch64-nc) blockaddress(@test_blockaddress, %ir-block.block), 0
+  ; LARGE:   [[MOVKXi:%[0-9]+]]:gpr64 = MOVKXi [[MOVZXi]], target-flags(aarch64-g1, aarch64-nc) blockaddress(@test_blockaddress, %ir-block.block), 16
+  ; LARGE:   [[MOVKXi1:%[0-9]+]]:gpr64 = MOVKXi [[MOVKXi]], target-flags(aarch64-g2, aarch64-nc) blockaddress(@test_blockaddress, %ir-block.block), 32
+  ; LARGE:   [[MOVKXi2:%[0-9]+]]:gpr64 = MOVKXi [[MOVKXi1]], target-flags(aarch64-g3) blockaddress(@test_blockaddress, %ir-block.block), 48
+  ; LARGE:   [[MOVZXi1:%[0-9]+]]:gpr64 = MOVZXi target-flags(aarch64-g0, aarch64-nc) @addr, 0
+  ; LARGE:   [[MOVKXi3:%[0-9]+]]:gpr64 = MOVKXi [[MOVZXi1]], target-flags(aarch64-g1, aarch64-nc) @addr, 16
+  ; LARGE:   [[MOVKXi4:%[0-9]+]]:gpr64 = MOVKXi [[MOVKXi3]], target-flags(aarch64-g2, aarch64-nc) @addr, 32
+  ; LARGE:   [[MOVKXi5:%[0-9]+]]:gpr64common = MOVKXi [[MOVKXi4]], target-flags(aarch64-g3) @addr, 48
+  ; LARGE:   STRXui [[MOVKXi2]], [[MOVKXi5]], 0 :: (store 8 into @addr)
+  ; LARGE:   BR [[MOVKXi2]]
+  ; LARGE: bb.1.block (address-taken):
+  ; LARGE:   RET_ReallyLR
+  bb.1 (%ir-block.0):
+    %0:gpr(p0) = G_BLOCK_ADDR blockaddress(@test_blockaddress, %ir-block.block)
+    %1:gpr(p0) = G_GLOBAL_VALUE @addr
+    G_STORE %0(p0), %1(p0) :: (store 8 into @addr)
+    G_BRINDIRECT %0(p0)
+
+  bb.2.block (address-taken):
+    RET_ReallyLR
+
+...