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