[RISCV] Post-RA merge base offset pass

[luismarques]

Summary

This is a stopgap solution until something more along the lines of D123264 and D123265 gets accepted. It conservatively does the folding of the low part of LLA addresses into the offset of load instructions. In my benchmarks it worked better than I was expecting, so maybe this is worth considering if the other patches get stalled?

[luismarques]

Stopgap solution no longer needed. Properly addressed by D123265.

llvm/lib/Target/RISCV/RISCV.h

 void initializeRISCVGatherScatterLoweringPass(PassRegistry &);
 
 FunctionPass *createRISCVSExtWRemovalPass();
 void initializeRISCVSExtWRemovalPass(PassRegistry &);
 
 FunctionPass *createRISCVMergeBaseOffsetOptPass();
 void initializeRISCVMergeBaseOffsetOptPass(PassRegistry &);
 
+FunctionPass *createRISCVPostRAMergeBaseOffsetPass();
+void initializeRISCVPostRAMergeBaseOffsetPass(PassRegistry &);
+
 FunctionPass *createRISCVExpandPseudoPass();
 void initializeRISCVExpandPseudoPass(PassRegistry &);
 
 FunctionPass *createRISCVExpandAtomicPseudoPass();
 void initializeRISCVExpandAtomicPseudoPass(PassRegistry &);
 
 FunctionPass *createRISCVInsertVSETVLIPass();
 void initializeRISCVInsertVSETVLIPass(PassRegistry &);

llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp

for (auto *MI : DeadInstrs)
     MI->eraseFromParent();
   return MadeChange;
 }
 
 /// Returns an instance of the Merge Base Offset Optimization pass.
 FunctionPass *llvm::createRISCVMergeBaseOffsetOptPass() {
   return new RISCVMergeBaseOffsetOpt();
 }
+
+#define RISCV_POSTRA_MERGE_BASE_OFFSET_NAME "RISCV Post-RA Merge Base Offset"
+
+namespace {
+struct RISCVPostRAMergeBaseOffset : public MachineFunctionPass {
+  static char ID;
+  bool runOnMachineFunction(MachineFunction &Fn) override;
+  bool detectAndFoldOffset(MachineInstr &AddrHi, MachineInstr &AddrLo,
+                           MachineInstr &Mem);
+  RISCVPostRAMergeBaseOffset() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override {
+    return RISCV_POSTRA_MERGE_BASE_OFFSET_NAME;
+  }
+
+private:
+  MachineRegisterInfo *MRI;
+  std::set<MachineInstr *> DeadInstrs;
+};
+} // end anonymous namespace
+
+bool RISCVPostRAMergeBaseOffset::runOnMachineFunction(MachineFunction &Fn) {
+  if (skipFunction(Fn.getFunction()))
+    return false;
+
+  bool MadeChange = false;
+  DeadInstrs.clear();
+  MRI = &Fn.getRegInfo();
+  for (MachineBasicBlock &MBB : Fn) {
+    if (MBB.size() < 3)
+      continue;
+    MachineBasicBlock::instr_iterator MII = MBB.instr_begin();
+    MachineBasicBlock::instr_iterator MIE = MBB.instr_end();
+    MachineInstr *AddrHi = &*MII++;
+    MachineInstr *AddrLo = &*MII++;
+    do {
+      MachineInstr *Mem = &*MII++;
+      MadeChange |= detectAndFoldOffset(*AddrHi, *AddrLo, *Mem);
+      AddrHi = AddrLo;
+      AddrLo = Mem;
+    } while (MII != MIE);
+  }
+  // Delete dead instructions.
+  for (auto *MI : DeadInstrs)
+    MI->eraseFromParent();
+  return MadeChange;
+}
+
+bool RISCVPostRAMergeBaseOffset::detectAndFoldOffset(MachineInstr &AddrHi,
+                                                     MachineInstr &AddrLo,
+                                                     MachineInstr &Mem) {
+  if (AddrHi.getOpcode() != RISCV::AUIPC ||
+      AddrHi.getOperand(1).getTargetFlags() != RISCVII::MO_PCREL_HI ||
+      AddrHi.getOperand(1).getType() != MachineOperand::MO_GlobalAddress ||
+      AddrHi.getOperand(1).getOffset() != 0)
+    return false;
+
+  if (AddrLo.getOpcode() != RISCV::ADDI ||
+      AddrLo.getOperand(1).getReg() != AddrHi.getOperand(0).getReg() ||
+      AddrLo.getOperand(2).getTargetFlags() != RISCVII::MO_PCREL_LO ||
+      AddrLo.getOperand(2).getType() != MachineOperand::MO_MachineBasicBlock)
+    return false;
+
+  switch (Mem.getOpcode()) {
+  default:
+    return false;
+  case RISCV::LB:
+  case RISCV::LH:
+  case RISCV::LW:
+  case RISCV::LBU:
+  case RISCV::LHU:
+  case RISCV::LWU:
+  case RISCV::LD:
+  case RISCV::FLH:
+  case RISCV::FLW:
+  case RISCV::FLD:;
+  }
+
+  if (Mem.getOperand(0).getReg() != AddrLo.getOperand(0).getReg() ||
+      Mem.getOperand(1).getReg() != AddrHi.getOperand(0).getReg() ||
+      Mem.getOperand(2).getType() != MachineOperand::MO_Immediate ||
+      Mem.getOperand(2).getImm() != 0)
+    return false;
+
+  Mem.removeOperand(2);
+  Mem.addOperand(AddrLo.getOperand(2));
+  DeadInstrs.insert(&AddrLo);
+  return true;
+}
+
+char RISCVPostRAMergeBaseOffset::ID = 0;
+INITIALIZE_PASS(RISCVPostRAMergeBaseOffset, "riscv-postra-merge-base-offset",
+                RISCV_POSTRA_MERGE_BASE_OFFSET_NAME, false, false)
+
+/// Returns an instance of the Post-RA Merge Base Offset Optimization pass.
+FunctionPass *llvm::createRISCVPostRAMergeBaseOffsetPass() {
+  return new RISCVPostRAMergeBaseOffset();
+}

llvm/lib/Target/RISCV/RISCVTargetMachine.cpp

 }
 
 void RISCVPassConfig::addPreSched2() {}
 
 void RISCVPassConfig::addPreEmitPass() { addPass(&BranchRelaxationPassID); }
 
 void RISCVPassConfig::addPreEmitPass2() {
   addPass(createRISCVExpandPseudoPass());
+  if (TM->getOptLevel() != CodeGenOpt::None)
+    addPass(createRISCVPostRAMergeBaseOffsetPass());
   // Schedule the expansion of AMOs at the last possible moment, avoiding the
   // possibility for other passes to break the requirements for forward
   // progress in the LR/SC block.
   addPass(createRISCVExpandAtomicPseudoPass());
 }
 
 void RISCVPassConfig::addMachineSSAOptimization() {
   TargetPassConfig::addMachineSSAOptimization();

llvm/test/CodeGen/RISCV/O3-pipeline.ll

 ; CHECK-NEXT:       Insert fentry calls
 ; CHECK-NEXT:       Insert XRay ops
 ; CHECK-NEXT:       Implement the 'patchable-function' attribute
 ; CHECK-NEXT:       Branch relaxation pass
 ; CHECK-NEXT:       Contiguously Lay Out Funclets
 ; CHECK-NEXT:       StackMap Liveness Analysis
 ; CHECK-NEXT:       Live DEBUG_VALUE analysis
 ; CHECK-NEXT:       RISCV pseudo instruction expansion pass
+; CHECK-NEXT:       RISCV Post-RA Merge Base Offset
 ; CHECK-NEXT:       RISCV atomic pseudo instruction expansion pass
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       RISCV Assembly Printer
 ; CHECK-NEXT:       Free MachineFunction

llvm/test/CodeGen/RISCV/codemodel-lowering.ll

 ; RV32I-SMALL-NEXT:    lui a0, %hi(G)
 ; RV32I-SMALL-NEXT:    lw a0, %lo(G)(a0)
 ; RV32I-SMALL-NEXT:    ret
 ;
 ; RV32I-MEDIUM-LABEL: lower_global:
 ; RV32I-MEDIUM:       # %bb.0:
 ; RV32I-MEDIUM-NEXT:  .LBB0_1: # Label of block must be emitted
 ; RV32I-MEDIUM-NEXT:    auipc a0, %pcrel_hi(G)
-; RV32I-MEDIUM-NEXT:    addi a0, a0, %pcrel_lo(.LBB0_1)
-; RV32I-MEDIUM-NEXT:    lw a0, 0(a0)
+; RV32I-MEDIUM-NEXT:    lw a0, %pcrel_lo(.LBB0_1)(a0)
 ; RV32I-MEDIUM-NEXT:    ret
   %1 = load volatile i32, i32* @G
   ret i32 %1
 }
 
 ; Check lowering of blockaddresses
 
 @addr = global i8* null