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 MachineFunctionPass::getAnalysisUsage(AU);
117 }
118};
119
120char AArch64PTrueCoalescingLegacy::ID = 0;
121
122} // end anonymous namespace
123
124INITIALIZE_PASS_BEGIN(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
125 "AArch64 PTRUE Coalescing", false, false)
126INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
127INITIALIZE_PASS_END(AArch64PTrueCoalescingLegacy, DEBUG_TYPE,
128 "AArch64 PTRUE Coalescing", false, false)
129
130unsigned
131AArch64PTrueCoalescingImpl::getSmallestElementSizeInUse(Register Reg) const {
132 // SSA form only applies to virtual registers.
133 if (!Reg.isVirtual())
134 return AArch64::ElementSizeNone;
135
136 unsigned SmallestElementSize = AArch64::ElementSizeNone;
137
138 for (MachineOperand &UseMO : MRI->use_nodbg_operands(Reg)) {
139 assert(UseMO.getSubReg() == 0 && "Unexpected SubReg!");
140 MachineInstr *UseMI = UseMO.getParent();
141
142 unsigned ElementSize = TII->getElementSizeForOpcode(Opc: UseMI->getOpcode());
143 if (ElementSize == AArch64::ElementSizeNone)
144 return AArch64::ElementSizeNone;
145
146 if (SmallestElementSize == AArch64::ElementSizeNone ||
147 SmallestElementSize > ElementSize)
148 SmallestElementSize = ElementSize;
149 }
150
151 return SmallestElementSize;
152}
153
154bool AArch64PTrueCoalescingImpl::tryCoalesce(PredicateInfo &DomPI,
155 PredicateInfo &CanPI) const {
156 assert(DomPI.isValid() && CanPI.isValid());
157 MachineInstr *DomMI = DomPI.MI;
158 MachineInstr *CanMI = CanPI.MI;
159
160 if (DomMI == CanMI || !MDT->dominates(A: DomMI, B: CanMI))
161 return false;
162
163 // A predicate's observable shape is the larger of the element size of the
164 // instruction writing the predicate and the one reading it. First check if
165 // DomPI can replace CanPI as-is for CanPI's users. If not, try changing DomPI
166 // to CanPI's element size, but only if DomPI's existing users would observe
167 // the same shape after that change.
168
169 bool MutateDomPTrue = false;
170 if (std::max(a: CanPI.ElementSize, b: CanPI.SmallestUsedElementSize) !=
171 std::max(a: DomPI.ElementSize, b: CanPI.SmallestUsedElementSize)) {
172 if (std::max(a: CanPI.ElementSize, b: DomPI.SmallestUsedElementSize) !=
173 std::max(a: DomPI.ElementSize, b: DomPI.SmallestUsedElementSize))
174 return false;
175
176 MutateDomPTrue = true;
177 }
178
179 Register DomReg = DomMI->getOperand(i: 0).getReg();
180 Register CanReg = CanMI->getOperand(i: 0).getReg();
181 if (!MRI->constrainRegClass(Reg: DomReg, RC: MRI->getRegClass(Reg: CanReg)))
182 return false;
183
184 LLVM_DEBUG(dbgs() << "Coalescing PTRUE: " << CanMI);
185 LLVM_DEBUG(dbgs() << " with: " << DomMI);
186
187 if (MutateDomPTrue) {
188 LLVM_DEBUG(dbgs() << " updated: " << DomMI);
189 DomMI->setDesc(TII->get(Opcode: CanMI->getOpcode()));
190 DomPI.ElementSize = CanPI.ElementSize;
191 LLVM_DEBUG(dbgs() << " to: " << DomMI);
192 }
193
194 MRI->replaceRegWith(FromReg: CanReg, ToReg: DomReg);
195 MRI->clearKillFlags(Reg: DomReg);
196 CanMI->eraseFromParent();
197
198 // Update DomPI based on uses inherited from CanPI.
199 if (CanPI.SmallestUsedElementSize < DomPI.SmallestUsedElementSize)
200 DomPI.SmallestUsedElementSize = CanPI.SmallestUsedElementSize;
201 CanPI.invalidate();
202 return true;
203}
204
205bool AArch64PTrueCoalescingImpl::run(MachineFunction &MF) {
206 if (!EnablePTrueCoalescing ||
207 !MF.getSubtarget<AArch64Subtarget>().isSVEorStreamingSVEAvailable())
208 return false;
209
210 TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
211 MRI = &MF.getRegInfo();
212
213 assert(MRI->isSSA() && "Expected to be run on SSA form!");
214
215 // TODO: Until we prove candidates share the same VG definition, do not
216 // coalesce in functions that define VG.
217 if (!MRI->def_empty(RegNo: AArch64::VG))
218 return false;
219
220 // A list of predicate setting instructions with some usage information.
221 SmallVector<PredicateInfo, 8> PIs;
222
223 // Build a list of predicates whose uses all have a known size.
224 for (MachineBasicBlock &MBB : MF)
225 for (MachineInstr &MI : MBB)
226 if (auto PI = createPredicateInfo(MI))
227 PIs.push_back(Elt: *PI);
228
229 LLVM_DEBUG(dbgs() << "Coalescable PTRUE candidates: " << PIs.size() << "\n");
230 bool Changed = false;
231
232 for (PredicateInfo &DominantPI : PIs) {
233 if (!DominantPI.isValid())
234 continue;
235
236 for (PredicateInfo &CandidatePI : PIs) {
237 if (!CandidatePI.isValid())
238 continue;
239
240 Changed |= tryCoalesce(DomPI&: DominantPI, CanPI&: CandidatePI);
241 }
242 }
243
244 return Changed;
245}
246
247bool AArch64PTrueCoalescingLegacy::runOnMachineFunction(MachineFunction &MF) {
248 MachineDominatorTree &MDT =
249 getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
250 return AArch64PTrueCoalescingImpl(MDT).run(MF);
251}
252
253FunctionPass *llvm::createAArch64PTrueCoalescingLegacyPass() {
254 return new AArch64PTrueCoalescingLegacy();
255}
256
257PreservedAnalyses
258AArch64PTrueCoalescingPass::run(MachineFunction &MF,
259 MachineFunctionAnalysisManager &MFAM) {
260 MachineDominatorTree &MDT = MFAM.getResult<MachineDominatorTreeAnalysis>(IR&: MF);
261 const bool Changed = AArch64PTrueCoalescingImpl(MDT).run(MF);
262 if (!Changed)
263 return PreservedAnalyses::all();
264
265 auto PA = getMachineFunctionPassPreservedAnalyses();
266 PA.preserveSet<CFGAnalyses>();
267 return PA;
268}
269