1//===- RISCVFoldMemOffset.cpp - Fold ADDI into memory offsets ------------===//
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// Look for ADDIs that can be removed by folding their immediate into later
10// load/store addresses. There may be other arithmetic instructions between the
11// addi and load/store that we need to reassociate through. If the final result
12// of the arithmetic is only used by load/store addresses, we can fold the
13// offset into the all the load/store as long as it doesn't create an offset
14// that is too large.
15//
16//===---------------------------------------------------------------------===//
17
18#include "RISCV.h"
19#include "RISCVSubtarget.h"
20#include "llvm/CodeGen/MachineFunctionPass.h"
21#include "llvm/CodeGen/RegisterClassInfo.h"
22#include <queue>
23
24using namespace llvm;
25
26#define DEBUG_TYPE "riscv-fold-mem-offset"
27#define RISCV_FOLD_MEM_OFFSET_NAME "RISC-V Fold Memory Offset"
28
29namespace {
30
31class RISCVFoldMemOffsetImpl {
32public:
33 bool run(MachineFunction &MF);
34
35private:
36 bool foldOffset(Register OrigReg, int64_t InitialOffset,
37 const MachineRegisterInfo &MRI,
38 DenseMap<MachineInstr *, int64_t> &FoldableInstrs);
39};
40
41class RISCVFoldMemOffsetLegacy : public MachineFunctionPass {
42public:
43 static char ID;
44
45 RISCVFoldMemOffsetLegacy() : MachineFunctionPass(ID) {}
46
47 bool runOnMachineFunction(MachineFunction &MF) override;
48
49 void getAnalysisUsage(AnalysisUsage &AU) const override {
50 AU.setPreservesCFG();
51 AU.addPreserved<MachineRegisterClassInfoWrapperPass>();
52 MachineFunctionPass::getAnalysisUsage(AU);
53 }
54
55 StringRef getPassName() const override { return RISCV_FOLD_MEM_OFFSET_NAME; }
56};
57
58// Wrapper class around a std::optional to allow accumulation.
59class FoldableOffset {
60 std::optional<int64_t> Offset;
61
62public:
63 bool hasValue() const { return Offset.has_value(); }
64 int64_t getValue() const { return *Offset; }
65
66 FoldableOffset &operator=(int64_t RHS) {
67 Offset = RHS;
68 return *this;
69 }
70
71 FoldableOffset &operator+=(int64_t RHS) {
72 if (!Offset)
73 Offset = 0;
74 Offset = (uint64_t)*Offset + (uint64_t)RHS;
75 return *this;
76 }
77
78 int64_t operator*() { return *Offset; }
79};
80
81} // end anonymous namespace
82
83char RISCVFoldMemOffsetLegacy::ID = 0;
84INITIALIZE_PASS(RISCVFoldMemOffsetLegacy, DEBUG_TYPE,
85 RISCV_FOLD_MEM_OFFSET_NAME, false, false)
86
87FunctionPass *llvm::createRISCVFoldMemOffsetLegacyPass() {
88 return new RISCVFoldMemOffsetLegacy();
89}
90
91// Walk forward from the ADDI looking for arithmetic instructions we can
92// analyze or memory instructions that use it as part of their address
93// calculation. For each arithmetic instruction we lookup how the offset
94// contributes to the value in that register use that information to
95// calculate the contribution to the output of this instruction.
96// Only addition and left shift are supported.
97// FIXME: Add multiplication by constant. The constant will be in a register.
98bool RISCVFoldMemOffsetImpl::foldOffset(
99 Register OrigReg, int64_t InitialOffset, const MachineRegisterInfo &MRI,
100 DenseMap<MachineInstr *, int64_t> &FoldableInstrs) {
101 // Map to hold how much the offset contributes to the value of this register.
102 DenseMap<Register, int64_t> RegToOffsetMap;
103
104 // Insert root offset into the map.
105 RegToOffsetMap[OrigReg] = InitialOffset;
106
107 std::queue<Register> Worklist;
108 Worklist.push(x: OrigReg);
109
110 while (!Worklist.empty()) {
111 Register Reg = Worklist.front();
112 Worklist.pop();
113
114 if (!Reg.isVirtual())
115 return false;
116
117 for (auto &User : MRI.use_nodbg_instructions(Reg)) {
118 FoldableOffset Offset;
119
120 switch (User.getOpcode()) {
121 default:
122 return false;
123 case RISCV::ADD:
124 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 1).getReg());
125 I != RegToOffsetMap.end())
126 Offset = I->second;
127 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 2).getReg());
128 I != RegToOffsetMap.end())
129 Offset += I->second;
130 break;
131 case RISCV::SH1ADD:
132 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 1).getReg());
133 I != RegToOffsetMap.end())
134 Offset = (uint64_t)I->second << 1;
135 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 2).getReg());
136 I != RegToOffsetMap.end())
137 Offset += I->second;
138 break;
139 case RISCV::SH2ADD:
140 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 1).getReg());
141 I != RegToOffsetMap.end())
142 Offset = (uint64_t)I->second << 2;
143 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 2).getReg());
144 I != RegToOffsetMap.end())
145 Offset += I->second;
146 break;
147 case RISCV::SH3ADD:
148 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 1).getReg());
149 I != RegToOffsetMap.end())
150 Offset = (uint64_t)I->second << 3;
151 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 2).getReg());
152 I != RegToOffsetMap.end())
153 Offset += I->second;
154 break;
155 case RISCV::ADD_UW:
156 case RISCV::SH1ADD_UW:
157 case RISCV::SH2ADD_UW:
158 case RISCV::SH3ADD_UW:
159 // Don't fold through the zero extended input.
160 if (User.getOperand(i: 1).getReg() == Reg)
161 return false;
162 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 2).getReg());
163 I != RegToOffsetMap.end())
164 Offset = I->second;
165 break;
166 case RISCV::SLLI: {
167 unsigned ShAmt = User.getOperand(i: 2).getImm();
168 if (auto I = RegToOffsetMap.find(Val: User.getOperand(i: 1).getReg());
169 I != RegToOffsetMap.end())
170 Offset = (uint64_t)I->second << ShAmt;
171 break;
172 }
173 case RISCV::LB:
174 case RISCV::LBU:
175 case RISCV::SB:
176 case RISCV::LH:
177 case RISCV::LH_INX:
178 case RISCV::LHU:
179 case RISCV::FLH:
180 case RISCV::SH:
181 case RISCV::SH_INX:
182 case RISCV::FSH:
183 case RISCV::LW:
184 case RISCV::LW_INX:
185 case RISCV::LWU:
186 case RISCV::FLW:
187 case RISCV::SW:
188 case RISCV::SW_INX:
189 case RISCV::FSW:
190 case RISCV::LD:
191 case RISCV::LD_RV32:
192 case RISCV::FLD:
193 case RISCV::SD:
194 case RISCV::SD_RV32:
195 case RISCV::FSD: {
196 // Can't fold into store value.
197 if (User.getOperand(i: 0).getReg() == Reg)
198 return false;
199
200 // Existing offset must be immediate.
201 if (!User.getOperand(i: 2).isImm())
202 return false;
203
204 // Require at least one operation between the ADDI and the load/store.
205 // We have other optimizations that should handle the simple case.
206 if (User.getOperand(i: 1).getReg() == OrigReg)
207 return false;
208
209 auto I = RegToOffsetMap.find(Val: User.getOperand(i: 1).getReg());
210 if (I == RegToOffsetMap.end())
211 return false;
212
213 int64_t LocalOffset = User.getOperand(i: 2).getImm();
214 assert(isInt<12>(LocalOffset));
215 int64_t CombinedOffset = (uint64_t)LocalOffset + (uint64_t)I->second;
216 if (!isInt<12>(x: CombinedOffset))
217 return false;
218
219 FoldableInstrs[&User] = CombinedOffset;
220 continue;
221 }
222 }
223
224 // If we reach here we should have an accumulated offset.
225 assert(Offset.hasValue() && "Expected an offset");
226
227 // If the offset is new or changed, add the destination register to the
228 // work list.
229 int64_t OffsetVal = Offset.getValue();
230 auto P =
231 RegToOffsetMap.try_emplace(Key: User.getOperand(i: 0).getReg(), Args&: OffsetVal);
232 if (P.second) {
233 Worklist.push(x: User.getOperand(i: 0).getReg());
234 } else if (P.first->second != OffsetVal) {
235 P.first->second = OffsetVal;
236 Worklist.push(x: User.getOperand(i: 0).getReg());
237 }
238 }
239 }
240
241 return true;
242}
243
244bool RISCVFoldMemOffsetImpl::run(MachineFunction &MF) {
245 // This optimization may increase size by preventing compression.
246 if (MF.getFunction().hasOptSize())
247 return false;
248
249 MachineRegisterInfo &MRI = MF.getRegInfo();
250
251 bool MadeChange = false;
252 for (MachineBasicBlock &MBB : MF) {
253 for (MachineInstr &MI : llvm::make_early_inc_range(Range&: MBB)) {
254 // FIXME: We can support ADDIW from an LUI+ADDIW pair if the result is
255 // equivalent to LUI+ADDI.
256 if (MI.getOpcode() != RISCV::ADDI)
257 continue;
258
259 // We only want to optimize register ADDIs.
260 if (!MI.getOperand(i: 1).isReg() || !MI.getOperand(i: 2).isImm())
261 continue;
262
263 // Ignore 'li'.
264 if (MI.getOperand(i: 1).getReg() == RISCV::X0)
265 continue;
266
267 int64_t Offset = MI.getOperand(i: 2).getImm();
268 assert(isInt<12>(Offset));
269
270 DenseMap<MachineInstr *, int64_t> FoldableInstrs;
271
272 if (!foldOffset(OrigReg: MI.getOperand(i: 0).getReg(), InitialOffset: Offset, MRI, FoldableInstrs))
273 continue;
274
275 if (FoldableInstrs.empty())
276 continue;
277
278 // We can fold this ADDI.
279 // Rewrite all the instructions.
280 for (auto [MemMI, NewOffset] : FoldableInstrs)
281 MemMI->getOperand(i: 2).setImm(NewOffset);
282
283 MRI.replaceRegWith(FromReg: MI.getOperand(i: 0).getReg(), ToReg: MI.getOperand(i: 1).getReg());
284 MRI.clearKillFlags(Reg: MI.getOperand(i: 1).getReg());
285 MI.eraseFromParent();
286 MadeChange = true;
287 }
288 }
289
290 return MadeChange;
291}
292
293bool RISCVFoldMemOffsetLegacy::runOnMachineFunction(MachineFunction &MF) {
294 if (skipFunction(F: MF.getFunction()))
295 return false;
296 return RISCVFoldMemOffsetImpl().run(MF);
297}
298
299PreservedAnalyses
300RISCVFoldMemOffsetPass::run(MachineFunction &MF,
301 MachineFunctionAnalysisManager &MFAM) {
302 bool Changed = RISCVFoldMemOffsetImpl().run(MF);
303 if (!Changed)
304 return PreservedAnalyses::all();
305
306 PreservedAnalyses PA = getMachineFunctionPassPreservedAnalyses();
307 PA.preserveSet<CFGAnalyses>();
308 PA.preserve<MachineRegisterClassAnalysis>();
309 return PA;
310}
311