1//===-- RISCVZilsdOptimizer.cpp - RISC-V Zilsd Load/Store Optimizer ------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file contains a pass that performs load/store optimizations for the
10// RISC-V Zilsd extension. It combines pairs of 32-bit load/store instructions
11// into single 64-bit LD/SD instructions when possible.
12//
13// The pass runs in two phases:
14// 1. Pre-allocation: Reschedules loads/stores to bring consecutive memory
15// accesses closer together and forms LD/SD pairs with register hints.
16// 2. Post-allocation: Fixes invalid LD/SD instructions if register allocation
17// didn't provide suitable consecutive registers.
18//
19// Note: second phase is integrated into RISCVLoadStoreOptimizer
20//
21//===----------------------------------------------------------------------===//
22
23#include "RISCV.h"
24#include "RISCVInstrInfo.h"
25#include "RISCVRegisterInfo.h"
26#include "RISCVSubtarget.h"
27#include "llvm/ADT/DenseMap.h"
28#include "llvm/ADT/SmallVector.h"
29#include "llvm/ADT/Statistic.h"
30#include "llvm/Analysis/AliasAnalysis.h"
31#include "llvm/CodeGen/MachineBasicBlock.h"
32#include "llvm/CodeGen/MachineDominators.h"
33#include "llvm/CodeGen/MachineFunction.h"
34#include "llvm/CodeGen/MachineFunctionPass.h"
35#include "llvm/CodeGen/MachineInstr.h"
36#include "llvm/CodeGen/MachineInstrBuilder.h"
37#include "llvm/CodeGen/MachineRegisterInfo.h"
38#include "llvm/InitializePasses.h"
39#include "llvm/Support/CommandLine.h"
40#include "llvm/Support/Debug.h"
41#include <algorithm>
42
43using namespace llvm;
44
45#define DEBUG_TYPE "riscv-zilsd-opt"
46
47STATISTIC(NumLDFormed, "Number of LD instructions formed");
48STATISTIC(NumSDFormed, "Number of SD instructions formed");
49
50static cl::opt<bool>
51 DisableZilsdOpt("disable-riscv-zilsd-opt", cl::Hidden, cl::init(Val: false),
52 cl::desc("Disable Zilsd load/store optimization"));
53
54static cl::opt<unsigned> MaxRescheduleDistance(
55 "riscv-zilsd-max-reschedule-distance", cl::Hidden, cl::init(Val: 10),
56 cl::desc("Maximum distance for rescheduling load/store instructions"));
57
58namespace {
59
60//===----------------------------------------------------------------------===//
61// Pre-allocation Zilsd optimization pass
62//===----------------------------------------------------------------------===//
63class RISCVPreAllocZilsdOpt : public MachineFunctionPass {
64public:
65 static char ID;
66
67 RISCVPreAllocZilsdOpt() : MachineFunctionPass(ID) {}
68
69 bool runOnMachineFunction(MachineFunction &MF) override;
70
71 StringRef getPassName() const override {
72 return "RISC-V pre-allocation Zilsd load/store optimization";
73 }
74
75 MachineFunctionProperties getRequiredProperties() const override {
76 return MachineFunctionProperties().setIsSSA();
77 }
78
79 void getAnalysisUsage(AnalysisUsage &AU) const override {
80 AU.addRequired<AAResultsWrapperPass>();
81 AU.addRequired<MachineDominatorTreeWrapperPass>();
82 AU.setPreservesCFG();
83 MachineFunctionPass::getAnalysisUsage(AU);
84 }
85 enum class MemoryOffsetKind {
86 Imm = 0,
87 Global = 1,
88 CPI = 2,
89 BlockAddr = 3,
90 FrameIdx = 4,
91 Unknown = 5,
92 };
93 using MemOffset = std::pair<MemoryOffsetKind, int>;
94 using BaseRegInfo = std::pair<unsigned, MemoryOffsetKind>;
95
96private:
97 bool isMemoryOp(const MachineInstr &MI);
98 bool rescheduleLoadStoreInstrs(MachineBasicBlock *MBB);
99 bool canFormLdSdPair(MachineInstr *MI0, MachineInstr *MI1);
100 bool rescheduleOps(MachineBasicBlock *MBB,
101 SmallVectorImpl<MachineInstr *> &MIs, BaseRegInfo Base,
102 bool IsLoad,
103 DenseMap<MachineInstr *, unsigned> &MI2LocMap);
104 bool isSafeToMove(MachineInstr *MI, MachineInstr *Target, bool MoveForward);
105 MemOffset getMemoryOpOffset(const MachineInstr &MI);
106
107 const RISCVSubtarget *STI;
108 const RISCVInstrInfo *TII;
109 const RISCVRegisterInfo *TRI;
110 MachineRegisterInfo *MRI;
111 AliasAnalysis *AA;
112 MachineDominatorTree *DT;
113 Align RequiredAlign;
114};
115
116} // end anonymous namespace
117
118char RISCVPreAllocZilsdOpt::ID = 0;
119
120INITIALIZE_PASS_BEGIN(RISCVPreAllocZilsdOpt, "riscv-prera-zilsd-opt",
121 "RISC-V pre-allocation Zilsd optimization", false, false)
122INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
123INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
124INITIALIZE_PASS_END(RISCVPreAllocZilsdOpt, "riscv-prera-zilsd-opt",
125 "RISC-V pre-allocation Zilsd optimization", false, false)
126
127//===----------------------------------------------------------------------===//
128// Pre-allocation pass implementation
129//===----------------------------------------------------------------------===//
130
131bool RISCVPreAllocZilsdOpt::runOnMachineFunction(MachineFunction &MF) {
132
133 if (DisableZilsdOpt || skipFunction(F: MF.getFunction()))
134 return false;
135
136 STI = &MF.getSubtarget<RISCVSubtarget>();
137
138 // Only run on RV32 with Zilsd extension
139 if (STI->is64Bit() || !STI->hasStdExtZilsd())
140 return false;
141
142 TII = STI->getInstrInfo();
143 TRI = STI->getRegisterInfo();
144 MRI = &MF.getRegInfo();
145 AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
146 DT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
147
148 // Check alignment: default is 8-byte, but allow 4-byte with tune feature
149 // If unaligned scalar memory is enabled, allow any alignment
150 RequiredAlign = STI->getZilsdAlign();
151 bool Modified = false;
152 for (auto &MBB : MF) {
153 Modified |= rescheduleLoadStoreInstrs(MBB: &MBB);
154 }
155
156 return Modified;
157}
158
159RISCVPreAllocZilsdOpt::MemOffset
160RISCVPreAllocZilsdOpt::getMemoryOpOffset(const MachineInstr &MI) {
161 switch (MI.getOpcode()) {
162 case RISCV::LW:
163 case RISCV::SW: {
164 // For LW/SW, the base is in operand 1 and offset is in operand 2
165 const MachineOperand &BaseOp = MI.getOperand(i: 1);
166 const MachineOperand &OffsetOp = MI.getOperand(i: 2);
167
168 // Handle immediate offset
169 if (OffsetOp.isImm()) {
170 if (BaseOp.isFI())
171 return std::make_pair(x: MemoryOffsetKind::FrameIdx, y: OffsetOp.getImm());
172 return std::make_pair(x: MemoryOffsetKind::Imm, y: OffsetOp.getImm());
173 }
174
175 // Handle symbolic operands with MO_LO flag (from MergeBaseOffset)
176 if (OffsetOp.getTargetFlags() & RISCVII::MO_LO) {
177 if (OffsetOp.isGlobal())
178 return std::make_pair(x: MemoryOffsetKind::Global, y: OffsetOp.getOffset());
179 if (OffsetOp.isCPI())
180 return std::make_pair(x: MemoryOffsetKind::CPI, y: OffsetOp.getOffset());
181 if (OffsetOp.isBlockAddress())
182 return std::make_pair(x: MemoryOffsetKind::BlockAddr,
183 y: OffsetOp.getOffset());
184 }
185
186 break;
187 }
188 default:
189 break;
190 }
191
192 return std::make_pair(x: MemoryOffsetKind::Unknown, y: 0);
193}
194
195bool RISCVPreAllocZilsdOpt::canFormLdSdPair(MachineInstr *MI0,
196 MachineInstr *MI1) {
197 if (!MI0->hasOneMemOperand() || !MI1->hasOneMemOperand())
198 return false;
199
200 // Get offsets and check they are consecutive
201 int Offset0 = getMemoryOpOffset(MI: *MI0).second;
202 int Offset1 = getMemoryOpOffset(MI: *MI1).second;
203
204 // Offsets must be 4 bytes apart
205 if (Offset1 - Offset0 != 4)
206 return false;
207
208 // We need to guarantee the alignment(base + offset) is legal.
209 const MachineMemOperand *MMO = *MI0->memoperands_begin();
210 if (MMO->getAlign() < RequiredAlign)
211 return false;
212
213 // Check that the two destination/source registers are different for
214 // load/store respectively.
215 if (MI0->getOperand(i: 0).getReg() == MI1->getOperand(i: 0).getReg())
216 return false;
217
218 return true;
219}
220
221bool RISCVPreAllocZilsdOpt::isSafeToMove(MachineInstr *MI, MachineInstr *Target,
222 bool MoveForward) {
223 MachineBasicBlock *MBB = MI->getParent();
224 MachineBasicBlock::iterator Start = MI->getIterator();
225 MachineBasicBlock::iterator End = Target->getIterator();
226
227 if (!MoveForward)
228 std::swap(a&: Start, b&: End);
229
230 // Increment Start to skip the current instruction
231 if (Start != MBB->end())
232 ++Start;
233
234 Register DefReg = MI->getOperand(i: 0).getReg();
235 const MachineOperand &BaseOp = MI->getOperand(i: 1);
236
237 unsigned ScanCount = 0;
238 for (auto It = Start; It != End; ++It, ++ScanCount) {
239 // Don't move across calls or terminators
240 if (It->isCall() || It->isTerminator()) {
241 LLVM_DEBUG(dbgs() << "Cannot move across call/terminator: " << *It);
242 return false;
243 }
244
245 // Don't move across instructions that modify memory barrier
246 if (It->hasUnmodeledSideEffects()) {
247 LLVM_DEBUG(dbgs() << "Cannot move across instruction with side effects: "
248 << *It);
249 return false;
250 }
251
252 // Check if the base register is modified
253 if (BaseOp.isReg() && It->modifiesRegister(Reg: BaseOp.getReg(), TRI)) {
254 LLVM_DEBUG(dbgs() << "Base register " << BaseOp.getReg()
255 << " modified by: " << *It);
256 return false;
257 }
258
259 // For loads, check if the loaded value is used
260 if (MI->mayLoad() &&
261 (It->readsRegister(Reg: DefReg, TRI) || It->modifiesRegister(Reg: DefReg, TRI))) {
262 LLVM_DEBUG(dbgs() << "Destination register " << DefReg
263 << " used by: " << *It);
264 return false;
265 }
266
267 // For stores, check if the stored register is modified
268 if (MI->mayStore() && It->modifiesRegister(Reg: DefReg, TRI)) {
269 LLVM_DEBUG(dbgs() << "Source register " << DefReg
270 << " modified by: " << *It);
271 return false;
272 }
273
274 // Check for memory operation interference
275 if (It->mayLoadOrStore() && It->mayAlias(AA, Other: *MI, /*UseTBAA*/ false)) {
276 LLVM_DEBUG(dbgs() << "Memory operation interference detected\n");
277 return false;
278 }
279 }
280
281 return true;
282}
283
284bool RISCVPreAllocZilsdOpt::rescheduleOps(
285 MachineBasicBlock *MBB, SmallVectorImpl<MachineInstr *> &MIs,
286 BaseRegInfo Base, bool IsLoad,
287 DenseMap<MachineInstr *, unsigned> &MI2LocMap) {
288 // Sort by offset, at this point it ensure base reg and MemoryOffsetKind are
289 // same, so we just need to simply sort by offset value.
290 llvm::sort(Start: MIs.begin(), End: MIs.end(), Comp: [this](MachineInstr *A, MachineInstr *B) {
291 return getMemoryOpOffset(MI: *A).second < getMemoryOpOffset(MI: *B).second;
292 });
293
294 bool Modified = false;
295
296 // Try to pair consecutive operations
297 for (size_t i = 0; i + 1 < MIs.size(); i++) {
298 MachineInstr *MI0 = MIs[i];
299 MachineInstr *MI1 = MIs[i + 1];
300
301 Register FirstReg = MI0->getOperand(i: 0).getReg();
302 Register SecondReg = MI1->getOperand(i: 0).getReg();
303 const MachineOperand &BaseOp = MI0->getOperand(i: 1);
304 const MachineOperand &OffsetOp = MI0->getOperand(i: 2);
305 assert((BaseOp.isReg() || BaseOp.isFI()) &&
306 "Base register should be register or frame index");
307
308 // At this point, MI0 and MI1 are:
309 // 1. both either LW or SW.
310 // 2. guaranteed to have same memory kind.
311 // 3. guaranteed to have same base register.
312 // 4. already be sorted by offset value.
313 // so we don't have to check these in canFormLdSdPair.
314 if (!canFormLdSdPair(MI0, MI1))
315 continue;
316
317 // Use MI2LocMap to determine which instruction appears later in program
318 // order
319 bool MI1IsLater = MI2LocMap[MI1] > MI2LocMap[MI0];
320
321 // For loads: move later instruction up (backwards) to earlier instruction
322 // For stores: move earlier instruction down (forwards) to later instruction
323 MachineInstr *MoveInstr, *TargetInstr;
324 if (IsLoad) {
325 // For loads: move the later instruction to the earlier one
326 MoveInstr = MI1IsLater ? MI1 : MI0;
327 TargetInstr = MI1IsLater ? MI0 : MI1;
328 } else {
329 // For stores: move the earlier instruction to the later one
330 MoveInstr = MI1IsLater ? MI0 : MI1;
331 TargetInstr = MI1IsLater ? MI1 : MI0;
332 }
333
334 unsigned Distance = MI1IsLater ? MI2LocMap[MI1] - MI2LocMap[MI0]
335 : MI2LocMap[MI0] - MI2LocMap[MI1];
336 if (!isSafeToMove(MI: MoveInstr, Target: TargetInstr, MoveForward: !IsLoad) ||
337 Distance > MaxRescheduleDistance)
338 continue;
339
340 // Move the instruction to the target position
341 MachineBasicBlock::iterator InsertPos = TargetInstr->getIterator();
342 ++InsertPos;
343
344 // If we need to move an instruction, do it now
345 if (MoveInstr != TargetInstr)
346 MBB->splice(Where: InsertPos, Other: MBB, From: MoveInstr->getIterator());
347
348 // Create the paired instruction
349 MachineInstrBuilder MIB;
350 DebugLoc DL = MI0->getDebugLoc();
351
352 if (IsLoad) {
353 MIB = BuildMI(BB&: *MBB, I: InsertPos, MIMD: DL, MCID: TII->get(Opcode: RISCV::PseudoLD_RV32_OPT))
354 .addReg(RegNo: FirstReg, Flags: RegState::Define)
355 .addReg(RegNo: SecondReg, Flags: RegState::Define);
356 ++NumLDFormed;
357 LLVM_DEBUG(dbgs() << "Formed LD: " << *MIB << "\n");
358 } else {
359 MIB = BuildMI(BB&: *MBB, I: InsertPos, MIMD: DL, MCID: TII->get(Opcode: RISCV::PseudoSD_RV32_OPT))
360 .addReg(RegNo: FirstReg)
361 .addReg(RegNo: SecondReg);
362 ++NumSDFormed;
363 LLVM_DEBUG(dbgs() << "Formed SD: " << *MIB << "\n");
364 }
365
366 if (BaseOp.isReg())
367 MIB = MIB.addReg(RegNo: BaseOp.getReg());
368 else
369 MIB = MIB.addFrameIndex(Idx: BaseOp.getIndex());
370 MIB = MIB.add(MO: OffsetOp);
371
372 // Copy memory operands
373 MIB.cloneMergedMemRefs(OtherMIs: {MI0, MI1});
374
375 // Add register allocation hints for consecutive registers
376 // RISC-V Zilsd requires even/odd register pairs
377 // Only set hints for virtual registers (physical registers already have
378 // encoding)
379 if (FirstReg.isVirtual() && SecondReg.isVirtual()) {
380 // For virtual registers, we can't determine even/odd yet, but we can hint
381 // that they should be allocated as a consecutive pair
382 MRI->setRegAllocationHint(VReg: FirstReg, Type: RISCVRI::RegPairEven, PrefReg: SecondReg);
383 MRI->setRegAllocationHint(VReg: SecondReg, Type: RISCVRI::RegPairOdd, PrefReg: FirstReg);
384 }
385
386 // Remove the original instructions
387 MI0->eraseFromParent();
388 MI1->eraseFromParent();
389
390 Modified = true;
391
392 // Skip the next instruction since we've already processed it
393 i++;
394 }
395
396 return Modified;
397}
398
399bool RISCVPreAllocZilsdOpt::isMemoryOp(const MachineInstr &MI) {
400 unsigned Opcode = MI.getOpcode();
401 if (Opcode != RISCV::LW && Opcode != RISCV::SW)
402 return false;
403
404 if (!MI.getOperand(i: 1).isReg() && !MI.getOperand(i: 1).isFI())
405 return false;
406
407 // When no memory operands are present, conservatively assume unaligned,
408 // volatile, unfoldable.
409 if (!MI.hasOneMemOperand())
410 return false;
411
412 const MachineMemOperand *MMO = *MI.memoperands_begin();
413
414 if (MMO->isVolatile() || MMO->isAtomic())
415 return false;
416
417 // sw <undef> could probably be eliminated entirely, but for now we just want
418 // to avoid making a mess of it.
419 if (MI.getOperand(i: 0).isReg() && MI.getOperand(i: 0).isUndef())
420 return false;
421
422 // Likewise don't mess with references to undefined addresses.
423 if (MI.getOperand(i: 1).isReg() && MI.getOperand(i: 1).isUndef())
424 return false;
425
426 return true;
427}
428
429bool RISCVPreAllocZilsdOpt::rescheduleLoadStoreInstrs(MachineBasicBlock *MBB) {
430 bool Modified = false;
431
432 // Process the basic block in windows delimited by calls, terminators,
433 // or instructions with duplicate base+offset pairs
434 MachineBasicBlock::iterator MBBI = MBB->begin();
435 MachineBasicBlock::iterator E = MBB->end();
436
437 while (MBBI != E) {
438 // Map from instruction to its location in the current window
439 DenseMap<MachineInstr *, unsigned> MI2LocMap;
440
441 // Map from base register to list of load/store instructions
442 using Base2InstMap = DenseMap<BaseRegInfo, SmallVector<MachineInstr *, 4>>;
443 using BaseVec = SmallVector<BaseRegInfo, 4>;
444 Base2InstMap Base2LdsMap;
445 Base2InstMap Base2StsMap;
446 BaseVec LdBases;
447 BaseVec StBases;
448
449 unsigned Loc = 0;
450
451 // Build the current window of instructions
452 for (; MBBI != E; ++MBBI) {
453 MachineInstr &MI = *MBBI;
454
455 // Stop at barriers (calls and terminators)
456 if (MI.isCall() || MI.isTerminator()) {
457 // Move past the barrier for next iteration
458 ++MBBI;
459 break;
460 }
461
462 // Track instruction location in window
463 if (!MI.isDebugInstr())
464 MI2LocMap[&MI] = ++Loc;
465
466 MemOffset Offset = getMemoryOpOffset(MI);
467 // Skip non-memory operations or it's not a valid memory offset kind.
468 if (!isMemoryOp(MI) || Offset.first == MemoryOffsetKind::Unknown)
469 continue;
470
471 bool IsLd = (MI.getOpcode() == RISCV::LW);
472 const MachineOperand &BaseOp = MI.getOperand(i: 1);
473 unsigned Base;
474 if (BaseOp.isReg())
475 Base = BaseOp.getReg().id();
476 else
477 Base = BaseOp.getIndex();
478 bool StopHere = false;
479
480 // Lambda to find or add base register entries
481 auto FindBases = [&](Base2InstMap &Base2Ops, BaseVec &Bases) {
482 auto [BI, Inserted] = Base2Ops.try_emplace(Key: {Base, Offset.first});
483 if (Inserted) {
484 // First time seeing this base register
485 BI->second.push_back(Elt: &MI);
486 Bases.push_back(Elt: {Base, Offset.first});
487 return;
488 }
489 // Check if we've seen this exact base+offset before
490 if (any_of(Range&: BI->second, P: [&](const MachineInstr *PrevMI) {
491 return Offset == getMemoryOpOffset(MI: *PrevMI);
492 })) {
493 // Found duplicate base+offset - stop here to process current window
494 StopHere = true;
495 } else {
496 BI->second.push_back(Elt: &MI);
497 }
498 };
499
500 if (IsLd)
501 FindBases(Base2LdsMap, LdBases);
502 else
503 FindBases(Base2StsMap, StBases);
504
505 if (StopHere) {
506 // Found a duplicate (a base+offset combination that's seen earlier).
507 // Backtrack to process the current window.
508 --Loc;
509 break;
510 }
511 }
512
513 // Process the current window - reschedule loads
514 for (auto Base : LdBases) {
515 SmallVectorImpl<MachineInstr *> &Lds = Base2LdsMap[Base];
516 if (Lds.size() > 1) {
517 Modified |= rescheduleOps(MBB, MIs&: Lds, Base, IsLoad: true, MI2LocMap);
518 }
519 }
520
521 // Process the current window - reschedule stores
522 for (auto Base : StBases) {
523 SmallVectorImpl<MachineInstr *> &Sts = Base2StsMap[Base];
524 if (Sts.size() > 1) {
525 Modified |= rescheduleOps(MBB, MIs&: Sts, Base, IsLoad: false, MI2LocMap);
526 }
527 }
528 }
529
530 return Modified;
531}
532
533//===----------------------------------------------------------------------===//
534// Pass creation functions
535//===----------------------------------------------------------------------===//
536
537FunctionPass *llvm::createRISCVPreAllocZilsdOptPass() {
538 return new RISCVPreAllocZilsdOpt();
539}
540