[llvm][RISCV] Implement Zilsd load/store pair optimization

llvm/lib/Target/RISCV/CMakeLists.txt

@@ -71,6 +71,7 @@ add_llvm_target(RISCVCodeGen
   RISCVVLOptimizer.cpp
   RISCVVMV0Elimination.cpp
   RISCVZacasABIFix.cpp
+  RISCVZilsdOptimizer.cpp
   GISel/RISCVCallLowering.cpp
   GISel/RISCVInstructionSelector.cpp
   GISel/RISCVLegalizerInfo.cpp

llvm/lib/Target/RISCV/RISCV.h

@@ -93,6 +93,9 @@ void initializeRISCVPushPopOptPass(PassRegistry &);
 FunctionPass *createRISCVLoadStoreOptPass();
 void initializeRISCVLoadStoreOptPass(PassRegistry &);
 
+FunctionPass *createRISCVPreAllocZilsdOptPass();
+void initializeRISCVPreAllocZilsdOptPass(PassRegistry &);
+
 FunctionPass *createRISCVZacasABIFixPass();
 void initializeRISCVZacasABIFixPass(PassRegistry &);
 

llvm/lib/Target/RISCV/RISCVFeatures.td

@@ -194,6 +194,10 @@ def HasStdExtZilsd : Predicate<"Subtarget->hasStdExtZilsd()">,
                        AssemblerPredicate<(all_of FeatureStdExtZilsd),
                                           "'Zilsd' (Load/Store pair instructions)">;
 
+def FeatureZilsd4ByteAlign
+    : SubtargetFeature<"zilsd-4byte-align", "AllowZilsd4ByteAlign", "true",
+                       "Allow 4-byte alignment for Zilsd LD/SD instructions">;
+
 // Multiply Extensions
 
 def FeatureStdExtZmmul

llvm/lib/Target/RISCV/RISCVInstrInfoZilsd.td

@@ -47,6 +47,23 @@ let Predicates = [HasStdExtZilsd, IsRV32] in {
 def PseudoLD_RV32 : PseudoLoad<"ld", GPRPairRV32>;
 def PseudoSD_RV32 : PseudoStore<"sd", GPRPairRV32>;
 
+// Pseudo instructions for load/store optimization with 2 separate registers
+def PseudoLD_RV32_OPT :
+    Pseudo<(outs GPR:$rd1, GPR:$rd2),
+           (ins GPR:$rs1, simm12_lo:$imm12), [], "", ""> {
+  let hasSideEffects = 0;
+  let mayLoad = 1;
+  let mayStore = 0;
+}
+
+def PseudoSD_RV32_OPT :
+    Pseudo<(outs),
+           (ins GPR:$rs1, GPR:$rs2, GPR:$rs3, simm12_lo:$imm12), [], "", ""> {
+  let hasSideEffects = 0;
+  let mayLoad = 0;
+  let mayStore = 1;
+}
+
 def : InstAlias<"ld $rd, (${rs1})", (LD_RV32 GPRPairRV32:$rd, GPR:$rs1, 0), 0>;
 def : InstAlias<"sd $rs2, (${rs1})", (SD_RV32 GPRPairRV32:$rs2, GPR:$rs1, 0), 0>;
 }

llvm/lib/Target/RISCV/RISCVLoadStoreOptimizer.cpp

@@ -11,6 +11,9 @@
 // paired instruction, leveraging hardware support for paired memory accesses.
 // Much of the pairing logic is adapted from the AArch64LoadStoreOpt pass.
 //
+// Post-allocation Zilsd decomposition: Fixes invalid LD/SD instructions if
+// register allocation didn't provide suitable consecutive registers.
+//
 // NOTE: The AArch64LoadStoreOpt pass performs additional optimizations such as
 // merging zero store instructions, promoting loads that read directly from a
 // preceding store, and merging base register updates with load/store
@@ -23,6 +26,7 @@
 
 #include "RISCV.h"
 #include "RISCVTargetMachine.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/MC/TargetRegistry.h"
@@ -38,6 +42,8 @@ using namespace llvm;
 // pairs.
 static cl::opt<unsigned> LdStLimit("riscv-load-store-scan-limit", cl::init(128),
                                    cl::Hidden);
+STATISTIC(NumLD2LW, "Number of LD instructions split back to LW");
+STATISTIC(NumSD2SW, "Number of SD instructions split back to SW");
 
 namespace {
 
@@ -75,6 +81,13 @@ struct RISCVLoadStoreOpt : public MachineFunctionPass {
   mergePairedInsns(MachineBasicBlock::iterator I,
                    MachineBasicBlock::iterator Paired, bool MergeForward);
 
+  // Post reg-alloc zilsd part
+  bool fixInvalidRegPairOp(MachineBasicBlock &MBB,
+                           MachineBasicBlock::iterator &MBBI);
+  bool isConsecutiveRegPair(Register First, Register Second);
+  void splitLdSdIntoTwo(MachineBasicBlock &MBB,
+                        MachineBasicBlock::iterator &MBBI, bool IsLoad);
+
 private:
   AliasAnalysis *AA;
   MachineRegisterInfo *MRI;
@@ -91,9 +104,8 @@ INITIALIZE_PASS(RISCVLoadStoreOpt, DEBUG_TYPE, RISCV_LOAD_STORE_OPT_NAME, false,
 bool RISCVLoadStoreOpt::runOnMachineFunction(MachineFunction &Fn) {
   if (skipFunction(Fn.getFunction()))
     return false;
+
   const RISCVSubtarget &Subtarget = Fn.getSubtarget<RISCVSubtarget>();
-  if (!Subtarget.useLoadStorePairs())
-    return false;
 
   bool MadeChange = false;
   TII = Subtarget.getInstrInfo();
@@ -103,18 +115,34 @@ bool RISCVLoadStoreOpt::runOnMachineFunction(MachineFunction &Fn) {
   ModifiedRegUnits.init(*TRI);
   UsedRegUnits.init(*TRI);
 
-  for (MachineBasicBlock &MBB : Fn) {
-    LLVM_DEBUG(dbgs() << "MBB: " << MBB.getName() << "\n");
+  if (Subtarget.useLoadStorePairs()) {
+    for (MachineBasicBlock &MBB : Fn) {
+      LLVM_DEBUG(dbgs() << "MBB: " << MBB.getName() << "\n");
+
+      for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
+           MBBI != E;) {
+        if (TII->isPairableLdStInstOpc(MBBI->getOpcode()) &&
+            tryToPairLdStInst(MBBI))
+          MadeChange = true;
+        else
+          ++MBBI;
+      }
+    }
+  }
 
-    for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
-         MBBI != E;) {
-      if (TII->isPairableLdStInstOpc(MBBI->getOpcode()) &&
-          tryToPairLdStInst(MBBI))
-        MadeChange = true;
-      else
-        ++MBBI;
+  if (!Subtarget.is64Bit() && Subtarget.hasStdExtZilsd()) {
+    for (auto &MBB : Fn) {
+      for (auto MBBI = MBB.begin(), E = MBB.end(); MBBI != E;) {
+        if (fixInvalidRegPairOp(MBB, MBBI)) {
+          MadeChange = true;
+          // Iterator was updated by fixInvalidRegPairOp
+        } else {
+          ++MBBI;
+        }
+      }
     }
   }
+
   return MadeChange;
 }
 
@@ -395,6 +423,143 @@ RISCVLoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
   return NextI;
 }
 
+//===----------------------------------------------------------------------===//
+// Post reg-alloc zilsd pass implementation
+//===----------------------------------------------------------------------===//
+
+bool RISCVLoadStoreOpt::isConsecutiveRegPair(Register First, Register Second) {
+  // Special case: both registers are zero register - this is valid for storing
+  // zeros
+  if (First == RISCV::X0 && Second == RISCV::X0)
+    return true;
+
+  // Check if registers form a valid even/odd pair for Zilsd
+  unsigned FirstNum = TRI->getEncodingValue(First);
+  unsigned SecondNum = TRI->getEncodingValue(Second);
+
+  // Must be consecutive and first must be even
+  return (FirstNum % 2 == 0) && (SecondNum == FirstNum + 1);
+}
+
+void RISCVLoadStoreOpt::splitLdSdIntoTwo(MachineBasicBlock &MBB,
+                                         MachineBasicBlock::iterator &MBBI,
+                                         bool IsLoad) {
+  MachineInstr *MI = &*MBBI;
+  DebugLoc DL = MI->getDebugLoc();
+
+  Register FirstReg = MI->getOperand(0).getReg();
+  Register SecondReg = MI->getOperand(1).getReg();
+  Register BaseReg = MI->getOperand(2).getReg();
+  int Offset = MI->getOperand(3).getImm();
+
+  unsigned Opc = IsLoad ? RISCV::LW : RISCV::SW;
+
+  // Create two separate instructions
+  if (IsLoad) {
+    auto MIB1 = BuildMI(MBB, MBBI, DL, TII->get(Opc))
+                    .addReg(FirstReg, RegState::Define)
+                    .addReg(BaseReg)
+                    .addImm(Offset);
+
+    auto MIB2 = BuildMI(MBB, MBBI, DL, TII->get(Opc))
+                    .addReg(SecondReg, RegState::Define)
+                    .addReg(BaseReg)
+                    .addImm(Offset + 4);
+
+    // Copy memory operands if the original instruction had them
+    // FIXME: This is overly conservative; the new instruction accesses 4 bytes,
+    // not 8.
+    if (MI->memoperands_begin() != MI->memoperands_end()) {
+      MIB1.cloneMemRefs(*MI);
+      MIB2.cloneMemRefs(*MI);
+    }
+
+    ++NumLD2LW;
+    LLVM_DEBUG(dbgs() << "Split LD back to two LW instructions\n");
+  } else {
+    auto MIB1 = BuildMI(MBB, MBBI, DL, TII->get(Opc))
+                    .addReg(FirstReg)
+                    .addReg(BaseReg)
+                    .addImm(Offset);
+
+    auto MIB2 = BuildMI(MBB, MBBI, DL, TII->get(Opc))
+                    .addReg(SecondReg)
+                    .addReg(BaseReg)
+                    .addImm(Offset + 4);
+
+    // Copy memory operands if the original instruction had them
+    // FIXME: This is overly conservative; the new instruction accesses 4 bytes,
+    // not 8.
+    if (MI->memoperands_begin() != MI->memoperands_end()) {
+      MIB1.cloneMemRefs(*MI);
+      MIB2.cloneMemRefs(*MI);
+    }
+
+    ++NumSD2SW;
+    LLVM_DEBUG(dbgs() << "Split SD back to two SW instructions\n");
+  }
+
+  // Remove the original paired instruction and update iterator
+  MBBI = MBB.erase(MBBI);
+}
+
+bool RISCVLoadStoreOpt::fixInvalidRegPairOp(MachineBasicBlock &MBB,
+                                            MachineBasicBlock::iterator &MBBI) {
+  MachineInstr *MI = &*MBBI;
+  unsigned Opcode = MI->getOpcode();
+
+  // Check if this is a Zilsd pseudo that needs fixing
+  if (Opcode != RISCV::PseudoLD_RV32_OPT && Opcode != RISCV::PseudoSD_RV32_OPT)
+    return false;
+
+  bool IsLoad = (Opcode == RISCV::PseudoLD_RV32_OPT);
+
+  Register FirstReg = MI->getOperand(0).getReg();
+  Register SecondReg = MI->getOperand(1).getReg();
+
+  // Check if we have valid consecutive registers
+  if (!isConsecutiveRegPair(FirstReg, SecondReg)) {
+    // Need to split back into two instructions
+    splitLdSdIntoTwo(MBB, MBBI, IsLoad);
+    return true;
+  }
+
+  // Registers are valid, convert to real LD/SD instruction
+  Register BaseReg = MI->getOperand(2).getReg();
+  int Offset = MI->getOperand(3).getImm();
+  DebugLoc DL = MI->getDebugLoc();
+
+  unsigned RealOpc = IsLoad ? RISCV::LD_RV32 : RISCV::SD_RV32;
+
+  // Create register pair from the two individual registers
+  unsigned RegPair = TRI->getMatchingSuperReg(FirstReg, RISCV::sub_gpr_even,
+                                              &RISCV::GPRPairRegClass);
+  // Create the real LD/SD instruction with register pair
+  MachineInstrBuilder MIB = BuildMI(MBB, MBBI, DL, TII->get(RealOpc));
+
+  if (IsLoad) {
+    // For LD, the register pair is the destination
+    MIB.addReg(RegPair, RegState::Define);
+  } else {
+    // For SD, the register pair is the source
+    MIB.addReg(RegPair);
+  }
+
+  MIB.addReg(BaseReg).addImm(Offset);
+
+  // Copy memory operands if the original instruction had them
+  if (MI->memoperands_begin() != MI->memoperands_end())
+    MIB.cloneMemRefs(*MI);
+
+  LLVM_DEBUG(dbgs() << "Converted pseudo to real instruction: " << *MIB
+                    << "\n");
+
+  // Remove the pseudo instruction and update iterator
+  MBBI = MBB.erase(MBBI);
+
+  return true;
+}
+
 // Returns an instance of the Load / Store Optimization pass.
 FunctionPass *llvm::createRISCVLoadStoreOptPass() {
   return new RISCVLoadStoreOpt();

llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp

@@ -853,6 +853,44 @@ bool RISCVRegisterInfo::getRegAllocationHints(
   const MachineRegisterInfo *MRI = &MF.getRegInfo();
   auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
 
+  // Handle RegPairEven/RegPairOdd hints for Zilsd register pairs
+  std::pair<unsigned, Register> Hint = MRI->getRegAllocationHint(VirtReg);
+  unsigned HintType = Hint.first;
+  Register Partner = Hint.second;
+
+  if (HintType == RISCVRI::RegPairEven || HintType == RISCVRI::RegPairOdd) {
+    // Check if we want the even or odd register of a consecutive pair
+    bool WantOdd = (HintType == RISCVRI::RegPairOdd);
+
+    // First priority: Check if partner is already allocated
+    if (Partner.isVirtual() && VRM && VRM->hasPhys(Partner)) {
+      MCPhysReg PartnerPhys = VRM->getPhys(Partner);
+      // Calculate the exact register we need for consecutive pairing
+      MCPhysReg TargetReg = PartnerPhys + (WantOdd ? 1 : -1);
+
+      // Verify it's valid and available
+      if (RISCV::GPRRegClass.contains(TargetReg) &&
+          is_contained(Order, TargetReg)) {
+        Hints.push_back(TargetReg);
+      }
+    }
+
+    // Second priority: Try to find consecutive register pairs in the allocation
+    // order
+    for (MCPhysReg PhysReg : Order) {
+      if (!PhysReg)
+        continue;
+
+      unsigned RegNum = getEncodingValue(PhysReg);
+      // Check if this register matches the even/odd requirement
+      bool IsOdd = (RegNum % 2 != 0);
+
+      // Verify the pair register exists and is in the same register class
+      if ((WantOdd && IsOdd) || (!WantOdd && !IsOdd))
+        Hints.push_back(PhysReg);
+    }
+  }
+
   bool BaseImplRetVal = TargetRegisterInfo::getRegAllocationHints(
       VirtReg, Order, Hints, MF, VRM, Matrix);
 
@@ -994,6 +1032,35 @@ bool RISCVRegisterInfo::getRegAllocationHints(
   return BaseImplRetVal;
 }
 
+void RISCVRegisterInfo::updateRegAllocHint(Register Reg, Register NewReg,
+                                           MachineFunction &MF) const {
+  MachineRegisterInfo *MRI = &MF.getRegInfo();
+  std::pair<unsigned, Register> Hint = MRI->getRegAllocationHint(Reg);
+
+  // Handle RegPairEven/RegPairOdd hints for Zilsd register pairs
+  if ((Hint.first == RISCVRI::RegPairOdd ||
+       Hint.first == RISCVRI::RegPairEven) &&
+      Hint.second.isVirtual()) {
+    // If 'Reg' is one of the even/odd register pair and it's now changed
+    // (e.g. coalesced) into a different register, the other register of the
+    // pair allocation hint must be updated to reflect the relationship change.
+    Register Partner = Hint.second;
+    std::pair<unsigned, Register> PartnerHint =
+        MRI->getRegAllocationHint(Partner);
+
+    // Make sure partner still points to us
+    if (PartnerHint.second == Reg) {
+      // Update partner to point to NewReg instead of Reg
+      MRI->setRegAllocationHint(Partner, PartnerHint.first, NewReg);
+
+      // If NewReg is virtual, set up the reciprocal hint
+      // NewReg takes over Reg's role, so it gets the SAME hint type as Reg
+      if (NewReg.isVirtual())
+        MRI->setRegAllocationHint(NewReg, Hint.first, Partner);
+    }
+  }
+}
+
 Register
 RISCVRegisterInfo::findVRegWithEncoding(const TargetRegisterClass &RegClass,
                                         uint16_t Encoding) const {

llvm/lib/Target/RISCV/RISCVRegisterInfo.h

@@ -37,6 +37,13 @@ enum : uint8_t {
   NFShiftMask = 0b111 << NFShift,
 };
 
+/// Register allocation hints for Zilsd register pairs
+enum {
+  // Used for Zilsd LD/SD register pairs
+  RegPairOdd = 1,
+  RegPairEven = 2,
+};
+
 /// \returns the IsVRegClass for the register class.
 static inline bool isVRegClass(uint8_t TSFlags) {
   return (TSFlags & IsVRegClassShiftMask) >> IsVRegClassShift;
@@ -143,6 +150,9 @@ struct RISCVRegisterInfo : public RISCVGenRegisterInfo {
                              const MachineFunction &MF, const VirtRegMap *VRM,
                              const LiveRegMatrix *Matrix) const override;
 
+  void updateRegAllocHint(Register Reg, Register NewReg,
+                          MachineFunction &MF) const override;
+
   Register findVRegWithEncoding(const TargetRegisterClass &RegClass,
                                 uint16_t Encoding) const;