[llvm] [GlobalISel][ARM] Selection of G_CONSTANT using constant pool (PR #96225)

[Serge Pavlov via llvm-commits]

https://github.com/spavloff updated https://github.com/llvm/llvm-project/pull/96225

>From 7c7fe5119052bf22114997dba93e50fed88e307a Mon Sep 17 00:00:00 2001
From: Serge Pavlov <sepavloff at gmail.com>
Date: Fri, 21 Jun 2024 00:04:28 +0700
Subject: [PATCH] [GlobalISel][ARM] Selection of G_CONSTANT using constant pool

Some constants cannot be represented as immediate operands. They can
be loaded from constant pool. This change implement selection of such
constants in the same way as the DAG selector uses.
---
 .../lib/Target/ARM/ARMInstructionSelector.cpp | 58 +++++++++++++++++++
 .../CodeGen/ARM/GlobalISel/select-const.mir   | 19 ++++++
 2 files changed, 77 insertions(+)

diff --git a/llvm/lib/Target/ARM/ARMInstructionSelector.cpp b/llvm/lib/Target/ARM/ARMInstructionSelector.cpp
index 3b3c380e1e1b32..ebbf569e499e2a 100644
--- a/llvm/lib/Target/ARM/ARMInstructionSelector.cpp
+++ b/llvm/lib/Target/ARM/ARMInstructionSelector.cpp
@@ -61,6 +61,10 @@ class ARMInstructionSelector : public InstructionSelector {
   bool selectSelect(MachineInstrBuilder &MIB, MachineRegisterInfo &MRI) const;
   bool selectShift(unsigned ShiftOpc, MachineInstrBuilder &MIB) const;
 
+  bool selectConstantUsingPool(MachineInstrBuilder &MIB) const;
+  void emitLoadFromConstantPool(const Register DefReg, const Constant *CPVal,
+                                MachineInstrBuilder &MIB) const;
+
   // Check if the types match and both operands have the expected size and
   // register bank.
   bool validOpRegPair(MachineRegisterInfo &MRI, unsigned LHS, unsigned RHS,
@@ -812,6 +816,57 @@ bool ARMInstructionSelector::selectShift(unsigned ShiftOpc,
   return constrainSelectedInstRegOperands(*MIB, TII, TRI, RBI);
 }
 
+bool ARMInstructionSelector::selectConstantUsingPool(
+    MachineInstrBuilder &MIB) const {
+  MachineInstr &MI = *MIB;
+  assert(MI.getOpcode() == TargetOpcode::G_CONSTANT && "Expected G_CONSTANT");
+
+  const Register DefReg = MI.getOperand(0).getReg();
+  MachineOperand &ConstOperand = MI.getOperand(1);
+  if (!ConstOperand.isCImm())
+    return false;
+  const Constant *ConstVal = ConstOperand.getCImm();
+  uint64_t ImmVal = ConstOperand.getCImm()->getZExtValue();
+
+  if (ConstantMaterializationCost(ImmVal, Subtarget) > 2 &&
+      !Subtarget->genExecuteOnly()) {
+    emitLoadFromConstantPool(DefReg, ConstVal, MIB);
+    MI.eraseFromParent();
+    return true;
+  }
+
+  return false;
+}
+
+void ARMInstructionSelector::emitLoadFromConstantPool(
+    const Register DefReg, const Constant *ConstVal,
+    MachineInstrBuilder &MIB) const {
+  MachineBasicBlock &MBB = *MIB->getParent();
+  MachineFunction &MF = *MBB.getParent();
+  const DataLayout &DL = MF.getDataLayout();
+
+  auto InsertBefore = std::next(MIB->getIterator());
+  auto &DbgLoc = MIB->getDebugLoc();
+
+  Type *ConstTy = ConstVal->getType();
+  unsigned Size = DL.getTypeStoreSize(ConstTy);
+  Align Alignment = DL.getPrefTypeAlign(ConstTy);
+
+  MachineConstantPool *Pool = MF.getConstantPool();
+  unsigned ConstIndex = Pool->getConstantPoolIndex(ConstVal, Alignment);
+
+  auto Load = BuildMI(MBB, InsertBefore, DbgLoc, TII.get(ARM::LDRcp))
+                  .addDef(DefReg)
+                  .addConstantPoolIndex(ConstIndex)
+                  .addImm(0)
+                  .add(predOps(ARMCC::AL));
+  MachinePointerInfo PtrInfo = MachinePointerInfo::getConstantPool(MF);
+  Load->addMemOperand(MF, MF.getMachineMemOperand(PtrInfo,
+                                                  MachineMemOperand::MOLoad,
+                                                  Size, Alignment));
+  constrainSelectedInstRegOperands(*Load, TII, TRI, RBI);
+}
+
 void ARMInstructionSelector::renderVFPF32Imm(
   MachineInstrBuilder &NewInstBuilder, const MachineInstr &OldInst,
   int OpIdx) const {
@@ -970,6 +1025,9 @@ bool ARMInstructionSelector::select(MachineInstr &I) {
     return selectCopy(I, TII, MRI, TRI, RBI);
   }
   case G_CONSTANT: {
+    if (selectConstantUsingPool(MIB))
+      return true;
+
     if (!MRI.getType(I.getOperand(0).getReg()).isPointer()) {
       // Non-pointer constants should be handled by TableGen.
       LLVM_DEBUG(dbgs() << "Unsupported constant type\n");
diff --git a/llvm/test/CodeGen/ARM/GlobalISel/select-const.mir b/llvm/test/CodeGen/ARM/GlobalISel/select-const.mir
index 5a1ea32a91c9ef..5fd83f8937f969 100644
--- a/llvm/test/CodeGen/ARM/GlobalISel/select-const.mir
+++ b/llvm/test/CodeGen/ARM/GlobalISel/select-const.mir
@@ -23,3 +23,22 @@ body:             |
     MOVPCLR 14 /* CC::al */, $noreg, implicit $r0
 
 ...
+---
+name:            get_from_constpool
+legalized:       true
+regBankSelected: true
+selected:        false
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gprb, preferred-register: '' }
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: get_from_constpool
+    ; CHECK: [[LDRcp:%[0-9]+]]:gpr = LDRcp %const.0, 0, 14 /* CC::al */, $noreg :: (load (s32) from constant-pool)
+    ; CHECK-NEXT: $r0 = COPY [[LDRcp]]
+    ; CHECK-NEXT: BX_RET 14 /* CC::al */, $noreg, implicit $r0
+    %0:gprb(s32) = G_CONSTANT i32 287454020
+    $r0 = COPY %0(s32)
+    BX_RET 14 /* CC::al */, $noreg, implicit $r0
+
+...