1//===- AArch64PTrueCoalescing.cpp - Coalesce SVE PTRUEs ---------*- C++ -*-===//
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 pass coalesces compatible all-active SVE PTRUE instructions.
10//
11// Consider two all-active PTRUE instructions X and Y with element sizes XSize
12// and YSize. If X dominates Y and XSize <= YSize, then every predicate bit that
13// Y sets is also set by X. In that case, uses of Y can be redirected to X as
14// long as each user of Y only reads predicate bits at YSize granularity or
15// larger.
16//
17// If the dominating PTRUE has a larger element size, we can coalesce the pair
18// by changing the dominating PTRUE to the smaller element size, provided that
19// all of its existing users are also safe with that granularity.
20//
21//===----------------------------------------------------------------------===//
22
23#include "AArch64.h"
24#include "AArch64InstrInfo.h"
25#include "AArch64Subtarget.h"
26#include "llvm/ADT/SmallVector.h"
27#include "llvm/CodeGen/MachineDominators.h"
28#include "llvm/CodeGen/MachineFunctionPass.h"
29#include "llvm/CodeGen/MachineInstr.h"
30#include "llvm/CodeGen/MachineRegisterInfo.h"
31#include "llvm/InitializePasses.h"
32#include "llvm/Pass.h"
33#include "llvm/Support/CommandLine.h"
34#include "llvm/Support/Debug.h"
35
36using namespace llvm;
37
38#define DEBUG_TYPE "aarch64-ptrue-coalesce"
39
40static cl::opt<bool> EnablePTrueCoalescing(
41 "aarch64-enable-ptrue-coalescing", cl::init(Val: false), cl::Hidden,
42 cl::desc("Enable coalescing of compatible AArch64 SVE PTRUE instructions"));
43
44namespace {
45
46class AArch64PTrueCoalescingImpl {
47 const AArch64InstrInfo *TII = nullptr;
48 MachineRegisterInfo *MRI = nullptr;
49 MachineDominatorTree *MDT = nullptr;
50
51public:
52 explicit AArch64PTrueCoalescingImpl(MachineDominatorTree &MDT) : MDT(&MDT) {}
53
54 bool run(MachineFunction &MF);
55
56private:
57 struct PredicateInfo {
58 // Instruction that created the predicate.
59 MachineInstr *MI = nullptr;
60 // Element size of the MI.
61 unsigned ElementSize = AArch64::ElementSizeNone;
62 // Smallest element size of all instructions that use the predicate.
63 unsigned SmallestUsedElementSize = AArch64::ElementSizeNone;
64
65 bool isValid() const {
66 assert(ElementSize != AArch64::ElementSizeNone &&
67 "PTRUE missing element size!");
68 return MI && SmallestUsedElementSize != AArch64::ElementSizeNone;
69 }
70
71 void invalidate() {
72 assert(isValid());
73 MI = nullptr;
74 }
75 };
76
77 std::optional<PredicateInfo> createPredicateInfo(MachineInstr &MI) const {
78 // TODO: Extend support beyond "PTRUE all"?
79 if (!isPTrueOpcode(Opc: MI.getOpcode()) || MI.getOperand(i: 1).getImm() != 31)
80 return std::nullopt;
81
82 Register Pred = MI.getOperand(i: 0).getReg();
83 unsigned SmallestUsedElementSize = getSmallestElementSizeInUse(Reg: Pred);
84 unsigned ElementSize = TII->getElementSizeForOpcode(Opc: MI.getOpcode());
85 assert(ElementSize != AArch64::ElementSizeNone &&
86 "PTRUE missing element size!");
87
88 if (SmallestUsedElementSize == AArch64::ElementSizeNone)
89 return std::nullopt;
90
91 return PredicateInfo{.MI: &MI, .ElementSize: ElementSize, .SmallestUsedElementSize: SmallestUsedElementSize};
92 }
93
94 // Return the smallest element size of all instructions that use Reg, or
95 // AArch64::ElementSizeNone when unknown.
96 unsigned getSmallestElementSizeInUse(Register Reg) const;
97
98 // Try to replace uses of CanPred with DomPred. In some cases that means
99 // modifying DomPred to support smaller element types.
100 bool tryCoalesce(PredicateInfo &DomPred, PredicateInfo &CanPred) const;
101};
102
103class AArch64PTrueCoalescingLegacy : public MachineFunctionPass {
104public:
105 static char ID;
106
107 AArch64PTrueCoalescingLegacy() : MachineFunctionPass(ID) {}
108
109 bool runOnMachineFunction(MachineFunction &MF) override;
110
111 StringRef getPassName() const override { return "AArch64 PTRUE Coalescing"; }
112
113 void getAnalysisUsage(AnalysisUsage &AU) const override {
114 AU.setPreservesCFG();
115 AU.addRequired<MachineDominatorTreeWrapperPass>();
116 AU.addPreserved<MachineDominatorTreeWrapperPass>();
117 MachineFunctionPass::getAnalysisUsage(AU);
118 }
119};
120
121char AArch64PTrueCoalescingLegacy::ID = 0;
122
123} // end anonymous namespace
124
125INITIALIZE_PASS_BEGIN(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
126 "AArch64 PTRUE Coalescing", false, false)
127INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
128INITIALIZE_PASS_END(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
129 "AArch64 PTRUE Coalescing", false, false)
130
131unsigned
132AArch64PTrueCoalescingImpl::getSmallestElementSizeInUse(Register Reg) const {
133 // SSA form only applies to virtual registers.
134 if (!Reg.isVirtual())
135 return AArch64::ElementSizeNone;
136
137 unsigned SmallestElementSize = AArch64::ElementSizeNone;
138
139 for (MachineOperand &UseMO : MRI->use_nodbg_operands(Reg)) {
140 assert(UseMO.getSubReg() == 0 && "Unexpected SubReg!");
141 MachineInstr *UseMI = UseMO.getParent();
142
143 unsigned ElementSize = TII->getElementSizeForOpcode(Opc: UseMI->getOpcode());
144 if (ElementSize == AArch64::ElementSizeNone)
145 return AArch64::ElementSizeNone;
146
147 if (SmallestElementSize == AArch64::ElementSizeNone ||
148 SmallestElementSize > ElementSize)
149 SmallestElementSize = ElementSize;
150 }
151
152 return SmallestElementSize;
153}
154
155bool AArch64PTrueCoalescingImpl::tryCoalesce(PredicateInfo &DomPI,
156 PredicateInfo &CanPI) const {
157 assert(DomPI.isValid() && CanPI.isValid());
158 MachineInstr *DomMI = DomPI.MI;
159 MachineInstr *CanMI = CanPI.MI;
160
161 if (DomMI == CanMI || !MDT->dominates(A: DomMI, B: CanMI))
162 return false;
163
164 // A predicate's observable shape is the larger of the element size of the
165 // instruction writing the predicate and the one reading it. First check if
166 // DomPI can replace CanPI as-is for CanPI's users. If not, try changing DomPI
167 // to CanPI's element size, but only if DomPI's existing users would observe
168 // the same shape after that change.
169
170 bool MutateDomPTrue = false;
171 if (std::max(a: CanPI.ElementSize, b: CanPI.SmallestUsedElementSize) !=
172 std::max(a: DomPI.ElementSize, b: CanPI.SmallestUsedElementSize)) {
173 if (std::max(a: CanPI.ElementSize, b: DomPI.SmallestUsedElementSize) !=
174 std::max(a: DomPI.ElementSize, b: DomPI.SmallestUsedElementSize))
175 return false;
176
177 MutateDomPTrue = true;
178 }
179
180 Register DomReg = DomMI->getOperand(i: 0).getReg();
181 Register CanReg = CanMI->getOperand(i: 0).getReg();
182 if (!MRI->constrainRegClass(Reg: DomReg, RC: MRI->getRegClass(Reg: CanReg)))
183 return false;
184
185 LLVM_DEBUG(dbgs() << "Coalescing PTRUE: " << CanMI);
186 LLVM_DEBUG(dbgs() << " with: " << DomMI);
187
188 if (MutateDomPTrue) {
189 LLVM_DEBUG(dbgs() << " updated: " << DomMI);
190 DomMI->setDesc(TII->get(Opcode: CanMI->getOpcode()));
191 DomPI.ElementSize = CanPI.ElementSize;
192 LLVM_DEBUG(dbgs() << " to: " << DomMI);
193 }
194
195 MRI->replaceRegWith(FromReg: CanReg, ToReg: DomReg);
196 MRI->clearKillFlags(Reg: DomReg);
197 CanMI->eraseFromParent();
198
199 // Update DomPI based on uses inherited from CanPI.
200 if (CanPI.SmallestUsedElementSize < DomPI.SmallestUsedElementSize)
201 DomPI.SmallestUsedElementSize = CanPI.SmallestUsedElementSize;
202 CanPI.invalidate();
203 return true;
204}
205
206bool AArch64PTrueCoalescingImpl::run(MachineFunction &MF) {
207 if (!EnablePTrueCoalescing ||
208 !MF.getSubtarget<AArch64Subtarget>().isSVEorStreamingSVEAvailable())
209 return false;
210
211 TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
212 MRI = &MF.getRegInfo();
213
214 assert(MRI->isSSA() && "Expected to be run on SSA form!");
215
216 // TODO: Until we prove candidates share the same VG definition, do not
217 // coalesce in functions that define VG.
218 if (!MRI->def_empty(RegNo: AArch64::VG))
219 return false;
220
221 // A list of predicate setting instructions with some usage information.
222 SmallVector<PredicateInfo, 8> PIs;
223
224 // Build a list of predicates whose uses all have a known size.
225 for (MachineBasicBlock &MBB : MF)
226 for (MachineInstr &MI : MBB)
227 if (auto PI = createPredicateInfo(MI))
228 PIs.push_back(Elt: *PI);
229
230 LLVM_DEBUG(dbgs() << "Coalescable PTRUE candidates: " << PIs.size() << "\n");
231 bool Changed = false;
232
233 for (PredicateInfo &DominantPI : PIs) {
234 if (!DominantPI.isValid())
235 continue;
236
237 for (PredicateInfo &CandidatePI : PIs) {
238 if (!CandidatePI.isValid())
239 continue;
240
241 Changed |= tryCoalesce(DomPI&: DominantPI, CanPI&: CandidatePI);
242 }
243 }
244
245 return Changed;
246}
247
248bool AArch64PTrueCoalescingLegacy::runOnMachineFunction(MachineFunction &MF) {
249 MachineDominatorTree &MDT =
250 getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
251 return AArch64PTrueCoalescingImpl(MDT).run(MF);
252}
253
254FunctionPass *llvm::createAArch64PTrueCoalescingLegacyPass() {
255 return new AArch64PTrueCoalescingLegacy();
256}
257
258PreservedAnalyses
259AArch64PTrueCoalescingPass::run(MachineFunction &MF,
260 MachineFunctionAnalysisManager &MFAM) {
261 MachineDominatorTree &MDT = MFAM.getResult<MachineDominatorTreeAnalysis>(IR&: MF);
262 const bool Changed = AArch64PTrueCoalescingImpl(MDT).run(MF);
263 if (!Changed)
264 return PreservedAnalyses::all();
265
266 auto PA = getMachineFunctionPassPreservedAnalyses();
267 PA.preserve<MachineDominatorTreeAnalysis>();
268 PA.preserveSet<CFGAnalyses>();
269 return PA;
270}
271