[RFC][WIP][RISCV] Fold ADDI of LLA + Load/Store code sequences

[luismarques]

Summary

RFC: I'm not sure this is the best approach for achieving this...?

For medany code models we aren't folding the ADDI offset of the AUIPC+ADDI+Load/Store (i.e. LLA+Load/Store) into the Load/Store offset. We have a pass for that but it runs before register allocation, while the LLA pseudoinstruction expansion runs just before code emission. As an alternative to this patch, I'm not sure either of those could be changed to be able to run in the opposite order and handle this case?

(This is currently a WIP. TODO: implement stores, etc.)

[jrtc27]

This then duplicates the AUIPC though and scales as 2*N rather than N+2. If you want to do this properly you really need to have a peephole pass that runs after the pseudo has been expanded.

[luismarques]

Thanks for the feedback.

llvm/lib/Target/RISCV/RISCVExpandPseudoInsts.cpp

 private:
   bool expandMBB(MachineBasicBlock &MBB);
   bool expandMI(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                 MachineBasicBlock::iterator &NextMBBI);
   bool expandAuipcInstPair(MachineBasicBlock &MBB,
                            MachineBasicBlock::iterator MBBI,
                            MachineBasicBlock::iterator &NextMBBI,
                            unsigned FlagsHi, unsigned SecondOpcode);
+  bool expandLoad(MachineBasicBlock &MBB,
+                 MachineBasicBlock::iterator MBBI,
+                 MachineBasicBlock::iterator &NextMBBI);
+  bool expandStore(MachineBasicBlock &MBB,
+                  MachineBasicBlock::iterator MBBI,
+                  MachineBasicBlock::iterator &NextMBBI);
   bool expandLoadLocalAddress(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator MBBI,
                               MachineBasicBlock::iterator &NextMBBI);
   bool expandLoadAddress(MachineBasicBlock &MBB,
                          MachineBasicBlock::iterator MBBI,
                          MachineBasicBlock::iterator &NextMBBI);
   bool expandLoadTLSIEAddress(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator MBBI,

 
 bool RISCVExpandPseudo::expandMI(MachineBasicBlock &MBB,
                                  MachineBasicBlock::iterator MBBI,
                                  MachineBasicBlock::iterator &NextMBBI) {
   // RISCVInstrInfo::getInstSizeInBytes expects that the total size of the
   // expanded instructions for each pseudo is correct in the Size field of the
   // tablegen definition for the pseudo.
   switch (MBBI->getOpcode()) {
+  case RISCV::PseudoLB:
+  case RISCV::PseudoLBU:
+  case RISCV::PseudoLH:
+  case RISCV::PseudoLHU:
+  case RISCV::PseudoLW:
+  case RISCV::PseudoLWU:
+  case RISCV::PseudoLD:
+    return expandLoad(MBB, MBBI, NextMBBI);
+  case RISCV::PseudoSB:
+  case RISCV::PseudoSH:
+  case RISCV::PseudoSW:
+  case RISCV::PseudoSD:
+    return expandStore(MBB, MBBI, NextMBBI);
   case RISCV::PseudoLLA:
     return expandLoadLocalAddress(MBB, MBBI, NextMBBI);
   case RISCV::PseudoLA:
     return expandLoadAddress(MBB, MBBI, NextMBBI);
   case RISCV::PseudoLA_TLS_IE:
     return expandLoadTLSIEAddress(MBB, MBBI, NextMBBI);
   case RISCV::PseudoLA_TLS_GD:
     return expandLoadTLSGDAddress(MBB, MBBI, NextMBBI);
bool RISCVExpandPseudo::expandAuipcInstPair(
   LivePhysRegs LiveRegs;
   computeAndAddLiveIns(LiveRegs, *NewMBB);
 
   NextMBBI = MBB.end();
   MI.eraseFromParent();
   return true;
 }
 
+bool RISCVExpandPseudo::expandLoad(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+    MachineBasicBlock::iterator &NextMBBI) {
+  MachineInstr &MI = *MBBI;
+  unsigned SecondOpcode;
+  switch (MI.getOpcode()) {
+    case RISCV::PseudoLB:  SecondOpcode = RISCV::LB; break;
+    case RISCV::PseudoLBU: SecondOpcode = RISCV::LBU; break;
+    case RISCV::PseudoLH:  SecondOpcode = RISCV::LH; break;
+    case RISCV::PseudoLHU: SecondOpcode = RISCV::LHU; break;
+    case RISCV::PseudoLW:  SecondOpcode = RISCV::LW; break;
+    case RISCV::PseudoLWU: SecondOpcode = RISCV::LWU; break;
+    case RISCV::PseudoLD:  SecondOpcode = RISCV::LD; break;
+    default:
+      llvm_unreachable("Unexpected opcode");
+  }
+  return expandAuipcInstPair(MBB, MBBI, NextMBBI, RISCVII::MO_PCREL_HI,
+                             SecondOpcode);
+}
+
+bool RISCVExpandPseudo::expandStore(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+    MachineBasicBlock::iterator &NextMBBI) {
+  MachineInstr &MI = *MBBI;
+  unsigned SecondOpcode;
+  switch (MI.getOpcode()) {
+    case RISCV::PseudoSW:
+      SecondOpcode = RISCV::SW;
+      break;
+    default:
+      llvm_unreachable("Unexpected opcode");
+  }
+  MI.eraseFromParent(); // TODO
+  return true;
+}
+
 bool RISCVExpandPseudo::expandLoadLocalAddress(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
     MachineBasicBlock::iterator &NextMBBI) {
   return expandAuipcInstPair(MBB, MBBI, NextMBBI, RISCVII::MO_PCREL_HI,
                              RISCV::ADDI);
 }
 
 bool RISCVExpandPseudo::expandLoadAddress(

llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp

case RISCV::FSD:
     break;
   }
 
   if (!isa<ConstantSDNode>(N->getOperand(OffsetOpIdx)))
     return false;
 
   SDValue Base = N->getOperand(BaseOpIdx);
 
-  // If the base is an ADDI, we can merge it in to the load/store.
-  if (!Base.isMachineOpcode() || Base.getMachineOpcode() != RISCV::ADDI)
+  if (!Base.isMachineOpcode())
     return false;
 
-  SDValue ImmOperand = Base.getOperand(1);
+  // If the base is an ADDI or PseudoLLA, we can merge it in to the load/store.
+  bool BaseIsAddi;
+  switch (Base.getMachineOpcode()) {
+    default:
+      return false;
+    case RISCV::ADDI:
+      BaseIsAddi = true;
+      break;
+    case RISCV::PseudoLLA:
+      BaseIsAddi = false;
+      break;
+  }
+
+  SDValue ImmOperand = Base.getOperand(BaseIsAddi ? 1 : 0);
   uint64_t Offset2 = N->getConstantOperandVal(OffsetOpIdx);
 
   if (auto *Const = dyn_cast<ConstantSDNode>(ImmOperand)) {
     int64_t Offset1 = Const->getSExtValue();
     int64_t CombinedOffset = Offset1 + Offset2;
     if (!isInt<12>(CombinedOffset))
       return false;
     ImmOperand = CurDAG->getTargetConstant(CombinedOffset, SDLoc(ImmOperand),
if (auto *Const = dyn_cast<ConstantSDNode>(ImmOperand)) {
     int64_t CombinedOffset = Offset1 + Offset2;
     ImmOperand = CurDAG->getTargetConstantPool(
         CP->getConstVal(), ImmOperand.getValueType(), CP->getAlign(),
         CombinedOffset, CP->getTargetFlags());
   } else {
     return false;
   }
 
+  if (BaseIsAddi)
     LLVM_DEBUG(dbgs() << "Folding add-immediate into mem-op:\nBase:    ");
+  else
+    LLVM_DEBUG(dbgs() << "Folding load address into mem-op:\nBase:    ");
   LLVM_DEBUG(Base->dump(CurDAG));
   LLVM_DEBUG(dbgs() << "\nN: ");
   LLVM_DEBUG(N->dump(CurDAG));
   LLVM_DEBUG(dbgs() << "\n");
 
+  if (BaseIsAddi) {
     // Modify the offset operand of the load/store.
     if (BaseOpIdx == 0) // Load
       CurDAG->UpdateNodeOperands(N, Base.getOperand(0), ImmOperand,
                                 N->getOperand(2));
     else // Store
       CurDAG->UpdateNodeOperands(N, N->getOperand(0), Base.getOperand(0),
                                 ImmOperand, N->getOperand(3));
-
+  } else {
+    unsigned Opc;
+    switch (N->getMachineOpcode()) {
+    default:
+      return false;
+    case RISCV::LB:  Opc = RISCV::PseudoLB; break;
+    case RISCV::LBU: Opc = RISCV::PseudoLBU; break;
+    case RISCV::LH:  Opc = RISCV::PseudoLH; break;
+    case RISCV::LHU: Opc = RISCV::PseudoLHU; break;
+    case RISCV::LW:  Opc = RISCV::PseudoLW; break;
+    case RISCV::LWU: Opc = RISCV::PseudoLWU; break;
+    case RISCV::LD:  Opc = RISCV::PseudoLD; break;
+    // case RISCV::FLH:
+    // case RISCV::FLW:
+    // case RISCV::FLD:
+    case RISCV::SB: Opc = RISCV::PseudoSB; break;
+    case RISCV::SH: Opc = RISCV::PseudoSH; break;
+    case RISCV::SW: Opc = RISCV::PseudoSW; break;
+    case RISCV::SD: Opc = RISCV::PseudoSD; break;
+    // case RISCV::FSH:
+    // case RISCV::FSW:
+    // case RISCV::FSD:
+    }
+    if (BaseOpIdx == 0) { // Load
+      SDNode *PseudoLoad =
+          CurDAG->getMachineNode(Opc, SDLoc(ImmOperand),
+                                N->getValueType(0), MVT::Other, ImmOperand, N->getOperand(2));
+      ReplaceUses(N, PseudoLoad);
+    } else {
+      MVT XLenVT = Subtarget->getXLenVT();
+      auto TmpOp = CurDAG->getRegister(RISCV::X0, XLenVT);
+      SDNode *PseudoStore =
+          CurDAG->getMachineNode(Opc, SDLoc(ImmOperand),
+                                N->getValueType(0), MVT::Other, TmpOp, ImmOperand, N->getOperand(3));
+      ReplaceUses(N, PseudoStore);
+    }
+  }
   return true;
 }
 
 // Try to remove sext.w if the input is a W instruction or can be made into
 // a W instruction cheaply.
 bool RISCVDAGToDAGISel::doPeepholeSExtW(SDNode *N) {
   // Look for the sext.w pattern, addiw rd, rs1, 0.
   if (N->getMachineOpcode() != RISCV::ADDIW ||
       !isNullConstant(N->getOperand(1)))
     return false;
 
   SDValue N0 = N->getOperand(0);
   if (!N0.isMachineOpcode())
     return false;
 
   switch (N0.getMachineOpcode()) {
   default:
     break;
   case RISCV::ADD:
   case RISCV::ADDI:
   case RISCV::SUB:
   case RISCV::MUL:
   case RISCV::SLLI: {
     // Convert sext.w+add/sub/mul to their W instructions. This will create
     // a new independent instruction. This improves latency.
     unsigned Opc;
     switch (N0.getMachineOpcode()) {
     default:
       llvm_unreachable("Unexpected opcode!");
     case RISCV::ADD:  Opc = RISCV::ADDW;  break;
     case RISCV::ADDI: Opc = RISCV::ADDIW; break;
     case RISCV::SUB:  Opc = RISCV::SUBW;  break;
     case RISCV::MUL:  Opc = RISCV::MULW;  break;
     case RISCV::SLLI: Opc = RISCV::SLLIW; break;
     }
 
     SDValue N00 = N0.getOperand(0);
     SDValue N01 = N0.getOperand(1);
 
     // Shift amount needs to be uimm5.
     if (N0.getMachineOpcode() == RISCV::SLLI &&
         !isUInt<5>(cast<ConstantSDNode>(N01)->getSExtValue()))
       break;
 
     SDNode *Result =
         CurDAG->getMachineNode(Opc, SDLoc(N), N->getValueType(0),
                                N00, N01);
     ReplaceUses(N, Result);
     return true;
   }
   case RISCV::ADDW:
   case RISCV::ADDIW:
   case RISCV::SUBW:
   case RISCV::MULW:
   case RISCV::SLLIW:

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
 
 define void @lower_blockaddress() nounwind {
 ; RV32I-SMALL-LABEL: lower_blockaddress:
 ; RV32I-SMALL:       # %bb.0:
 ; RV32I-SMALL-NEXT:    lui a0, %hi(addr)
 ; RV32I-SMALL-NEXT:    li a1, 1
 ; RV32I-SMALL-NEXT:    sw a1, %lo(addr)(a0)
 ; RV32I-SMALL-NEXT:    ret
 ;
 ; RV32I-MEDIUM-LABEL: lower_blockaddress:
 ; RV32I-MEDIUM:       # %bb.0:
-; RV32I-MEDIUM-NEXT:  .LBB1_1: # Label of block must be emitted
-; RV32I-MEDIUM-NEXT:    auipc a0, %pcrel_hi(addr)
-; RV32I-MEDIUM-NEXT:    addi a0, a0, %pcrel_lo(.LBB1_1)
-; RV32I-MEDIUM-NEXT:    li a1, 1
-; RV32I-MEDIUM-NEXT:    sw a1, 0(a0)
 ; RV32I-MEDIUM-NEXT:    ret
   store volatile i8* blockaddress(@lower_blockaddress, %block), i8** @addr
   ret void
 
 block:
   unreachable
 }