1//===- SMEPeepholeOpt.cpp - SME peephole optimization pass-----------------===//
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// This pass tries to remove back-to-back (smstart, smstop) and
9// (smstop, smstart) sequences. The pass is conservative when it cannot
10// determine that it is safe to remove these sequences.
11//===----------------------------------------------------------------------===//
12
13#include "AArch64InstrInfo.h"
14#include "AArch64MachineFunctionInfo.h"
15#include "AArch64Subtarget.h"
16#include "llvm/CodeGen/MachineBasicBlock.h"
17#include "llvm/CodeGen/MachineFunctionPass.h"
18#include "llvm/CodeGen/MachineRegisterInfo.h"
19#include "llvm/CodeGen/TargetRegisterInfo.h"
20
21using namespace llvm;
22
23#define DEBUG_TYPE "aarch64-sme-peephole-opt"
24
25namespace {
26
27struct SMEPeepholeOpt : public MachineFunctionPass {
28 static char ID;
29
30 SMEPeepholeOpt() : MachineFunctionPass(ID) {}
31
32 bool runOnMachineFunction(MachineFunction &MF) override;
33
34 StringRef getPassName() const override {
35 return "SME Peephole Optimization pass";
36 }
37
38 void getAnalysisUsage(AnalysisUsage &AU) const override {
39 AU.setPreservesCFG();
40 MachineFunctionPass::getAnalysisUsage(AU);
41 }
42
43 bool optimizeStartStopPairs(MachineBasicBlock &MBB,
44 bool &HasRemovedAllSMChanges) const;
45 bool visitRegSequence(MachineInstr &MI);
46};
47
48char SMEPeepholeOpt::ID = 0;
49
50} // end anonymous namespace
51
52static bool isConditionalStartStop(const MachineInstr *MI) {
53 return MI->getOpcode() == AArch64::MSRpstatePseudo;
54}
55
56static bool isMatchingStartStopPair(const MachineInstr *MI1,
57 const MachineInstr *MI2) {
58 // We only consider the same type of streaming mode change here, i.e.
59 // start/stop SM, or start/stop ZA pairs.
60 if (MI1->getOperand(i: 0).getImm() != MI2->getOperand(i: 0).getImm())
61 return false;
62
63 // One must be 'start', the other must be 'stop'
64 if (MI1->getOperand(i: 1).getImm() == MI2->getOperand(i: 1).getImm())
65 return false;
66
67 bool IsConditional = isConditionalStartStop(MI: MI2);
68 if (isConditionalStartStop(MI: MI1) != IsConditional)
69 return false;
70
71 if (!IsConditional)
72 return true;
73
74 // Check to make sure the conditional start/stop pairs are identical.
75 if (MI1->getOperand(i: 2).getImm() != MI2->getOperand(i: 2).getImm())
76 return false;
77
78 // Ensure reg masks are identical.
79 if (MI1->getOperand(i: 4).getRegMask() != MI2->getOperand(i: 4).getRegMask())
80 return false;
81
82 // Only consider conditional start/stop pairs which read the same register
83 // holding the original value of pstate.sm. This is somewhat over conservative
84 // as all conditional streaming mode changes only look at the state on entry
85 // to the function.
86 if (MI1->getOperand(i: 3).isReg() && MI2->getOperand(i: 3).isReg()) {
87 Register Reg1 = MI1->getOperand(i: 3).getReg();
88 Register Reg2 = MI2->getOperand(i: 3).getReg();
89 if (Reg1.isPhysical() || Reg2.isPhysical() || Reg1 != Reg2)
90 return false;
91 }
92
93 return true;
94}
95
96static bool ChangesStreamingMode(const MachineInstr *MI) {
97 assert((MI->getOpcode() == AArch64::MSRpstatesvcrImm1 ||
98 MI->getOpcode() == AArch64::MSRpstatePseudo) &&
99 "Expected MI to be a smstart/smstop instruction");
100 return MI->getOperand(i: 0).getImm() == AArch64SVCR::SVCRSM ||
101 MI->getOperand(i: 0).getImm() == AArch64SVCR::SVCRSMZA;
102}
103
104static bool isSVERegOp(const TargetRegisterInfo &TRI,
105 const MachineRegisterInfo &MRI,
106 const MachineOperand &MO) {
107 if (!MO.isReg())
108 return false;
109
110 Register R = MO.getReg();
111 if (R.isPhysical())
112 return llvm::any_of(Range: TRI.subregs_inclusive(Reg: R), P: [](const MCPhysReg &SR) {
113 return AArch64::ZPRRegClass.contains(Reg: SR) ||
114 AArch64::PPRRegClass.contains(Reg: SR);
115 });
116
117 const TargetRegisterClass *RC = MRI.getRegClass(Reg: R);
118 return TRI.getCommonSubClass(A: &AArch64::ZPRRegClass, B: RC) ||
119 TRI.getCommonSubClass(A: &AArch64::PPRRegClass, B: RC);
120}
121
122bool SMEPeepholeOpt::optimizeStartStopPairs(
123 MachineBasicBlock &MBB, bool &HasRemovedAllSMChanges) const {
124 const MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
125 const TargetRegisterInfo &TRI =
126 *MBB.getParent()->getSubtarget().getRegisterInfo();
127
128 bool Changed = false;
129 MachineInstr *Prev = nullptr;
130
131 // Walk through instructions in the block trying to find pairs of smstart
132 // and smstop nodes that cancel each other out. We only permit a limited
133 // set of instructions to appear between them, otherwise we reset our
134 // tracking.
135 unsigned NumSMChanges = 0;
136 unsigned NumSMChangesRemoved = 0;
137 for (MachineInstr &MI : make_early_inc_range(Range&: MBB)) {
138 switch (MI.getOpcode()) {
139 case AArch64::MSRpstatesvcrImm1:
140 case AArch64::MSRpstatePseudo: {
141 if (ChangesStreamingMode(MI: &MI))
142 NumSMChanges++;
143
144 if (!Prev)
145 Prev = &MI;
146 else if (isMatchingStartStopPair(MI1: Prev, MI2: &MI)) {
147 // If they match, we can remove them, and possibly any instructions
148 // that we marked for deletion in between.
149 Prev->eraseFromParent();
150 MI.eraseFromParent();
151 Prev = nullptr;
152 Changed = true;
153 NumSMChangesRemoved += 2;
154 } else {
155 Prev = &MI;
156 }
157 continue;
158 }
159 default:
160 if (!Prev)
161 // Avoid doing expensive checks when Prev is nullptr.
162 continue;
163 break;
164 }
165
166 // Test if the instructions in between the start/stop sequence are agnostic
167 // of streaming mode. If not, the algorithm should reset.
168 switch (MI.getOpcode()) {
169 default:
170 Prev = nullptr;
171 break;
172 case AArch64::COALESCER_BARRIER_FPR16:
173 case AArch64::COALESCER_BARRIER_FPR32:
174 case AArch64::COALESCER_BARRIER_FPR64:
175 case AArch64::COALESCER_BARRIER_FPR128:
176 case AArch64::COPY:
177 // These instructions should be safe when executed on their own, but
178 // the code remains conservative when SVE registers are used. There may
179 // exist subtle cases where executing a COPY in a different mode results
180 // in different behaviour, even if we can't yet come up with any
181 // concrete example/test-case.
182 if (isSVERegOp(TRI, MRI, MO: MI.getOperand(i: 0)) ||
183 isSVERegOp(TRI, MRI, MO: MI.getOperand(i: 1)))
184 Prev = nullptr;
185 break;
186 case AArch64::RestoreZAPseudo:
187 case AArch64::InOutZAUsePseudo:
188 case AArch64::CommitZASavePseudo:
189 case AArch64::SMEStateAllocPseudo:
190 case AArch64::RequiresZASavePseudo:
191 // These instructions only depend on the ZA state, not the streaming mode,
192 // so if the pair of smstart/stop is only changing the streaming mode, we
193 // can permit these instructions.
194 if (Prev->getOperand(i: 0).getImm() != AArch64SVCR::SVCRSM)
195 Prev = nullptr;
196 break;
197 case AArch64::ADJCALLSTACKDOWN:
198 case AArch64::ADJCALLSTACKUP:
199 case AArch64::ANDXri:
200 case AArch64::ADDXri:
201 // We permit these as they don't generate SVE/NEON instructions.
202 break;
203 case AArch64::MSRpstatesvcrImm1:
204 case AArch64::MSRpstatePseudo:
205 llvm_unreachable("Should have been handled");
206 }
207 }
208
209 HasRemovedAllSMChanges =
210 NumSMChanges && (NumSMChanges == NumSMChangesRemoved);
211 return Changed;
212}
213
214// Using the FORM_TRANSPOSED_REG_TUPLE pseudo can improve register allocation
215// of multi-vector intrinsics. However, the pseudo should only be emitted if
216// the input registers of the REG_SEQUENCE are copy nodes where the source
217// register is in a StridedOrContiguous class. For example:
218//
219// %3:zpr2stridedorcontiguous = LD1B_2Z_IMM_PSEUDO ..
220// %4:zpr = COPY %3.zsub1:zpr2stridedorcontiguous
221// %5:zpr = COPY %3.zsub0:zpr2stridedorcontiguous
222// %6:zpr2stridedorcontiguous = LD1B_2Z_PSEUDO ..
223// %7:zpr = COPY %6.zsub1:zpr2stridedorcontiguous
224// %8:zpr = COPY %6.zsub0:zpr2stridedorcontiguous
225// %9:zpr2mul2 = REG_SEQUENCE %5:zpr, %subreg.zsub0, %8:zpr, %subreg.zsub1
226//
227// -> %9:zpr2mul2 = FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO %5:zpr, %8:zpr
228//
229bool SMEPeepholeOpt::visitRegSequence(MachineInstr &MI) {
230 assert(MI.getMF()->getRegInfo().isSSA() && "Expected to be run on SSA form!");
231
232 MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
233 switch (MRI.getRegClass(Reg: MI.getOperand(i: 0).getReg())->getID()) {
234 case AArch64::ZPR2RegClassID:
235 case AArch64::ZPR4RegClassID:
236 case AArch64::ZPR2Mul2RegClassID:
237 case AArch64::ZPR4Mul4RegClassID:
238 break;
239 default:
240 return false;
241 }
242
243 // The first operand is the register class created by the REG_SEQUENCE.
244 // Each operand pair after this consists of a vreg + subreg index, so
245 // for example a sequence of 2 registers will have a total of 5 operands.
246 if (MI.getNumOperands() != 5 && MI.getNumOperands() != 9)
247 return false;
248
249 MCRegister SubReg = MCRegister::NoRegister;
250 for (unsigned I = 1; I < MI.getNumOperands(); I += 2) {
251 MachineOperand &MO = MI.getOperand(i: I);
252
253 MachineOperand *Def = MRI.getOneDef(Reg: MO.getReg());
254 if (!Def || !Def->getParent()->isCopy())
255 return false;
256
257 const MachineOperand &CopySrc = Def->getParent()->getOperand(i: 1);
258 unsigned OpSubReg = CopySrc.getSubReg();
259 if (SubReg == MCRegister::NoRegister)
260 SubReg = OpSubReg;
261
262 MachineOperand *CopySrcOp = MRI.getOneDef(Reg: CopySrc.getReg());
263 if (!CopySrcOp || !CopySrcOp->isReg() || OpSubReg != SubReg ||
264 CopySrcOp->getReg().isPhysical())
265 return false;
266
267 const TargetRegisterClass *CopySrcClass =
268 MRI.getRegClass(Reg: CopySrcOp->getReg());
269 if (CopySrcClass != &AArch64::ZPR2StridedOrContiguousRegClass &&
270 CopySrcClass != &AArch64::ZPR4StridedOrContiguousRegClass)
271 return false;
272 }
273
274 unsigned Opc = MI.getNumOperands() == 5
275 ? AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO
276 : AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO;
277
278 const TargetInstrInfo *TII =
279 MI.getMF()->getSubtarget<AArch64Subtarget>().getInstrInfo();
280 MachineInstrBuilder MIB = BuildMI(BB&: *MI.getParent(), I&: MI, MIMD: MI.getDebugLoc(),
281 MCID: TII->get(Opcode: Opc), DestReg: MI.getOperand(i: 0).getReg());
282 for (unsigned I = 1; I < MI.getNumOperands(); I += 2)
283 MIB.addReg(RegNo: MI.getOperand(i: I).getReg());
284
285 MI.eraseFromParent();
286 return true;
287}
288
289INITIALIZE_PASS(SMEPeepholeOpt, "aarch64-sme-peephole-opt",
290 "SME Peephole Optimization", false, false)
291
292bool SMEPeepholeOpt::runOnMachineFunction(MachineFunction &MF) {
293 if (skipFunction(F: MF.getFunction()))
294 return false;
295
296 AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
297 SMEAttrs SMEFnAttrs = AFI->getSMEFnAttrs();
298
299 if (!MF.getSubtarget<AArch64Subtarget>().hasSME() &&
300 !SMEFnAttrs.hasStreamingCompatibleInterface())
301 return false;
302
303 assert(MF.getRegInfo().isSSA() && "Expected to be run on SSA form!");
304
305 bool Changed = false;
306 bool FunctionHasAllSMChangesRemoved = false;
307
308 // Even if the block lives in a function with no SME attributes attached we
309 // still have to analyze all the blocks because we may call a streaming
310 // function that requires smstart/smstop pairs.
311 for (MachineBasicBlock &MBB : MF) {
312 bool BlockHasAllSMChangesRemoved;
313 Changed |= optimizeStartStopPairs(MBB, HasRemovedAllSMChanges&: BlockHasAllSMChangesRemoved);
314 FunctionHasAllSMChangesRemoved |= BlockHasAllSMChangesRemoved;
315
316 if (MF.getSubtarget<AArch64Subtarget>().isStreaming()) {
317 for (MachineInstr &MI : make_early_inc_range(Range&: MBB))
318 if (MI.getOpcode() == AArch64::REG_SEQUENCE)
319 Changed |= visitRegSequence(MI);
320 }
321 }
322
323 if (FunctionHasAllSMChangesRemoved)
324 AFI->setHasStreamingModeChanges(false);
325
326 return Changed;
327}
328
329FunctionPass *llvm::createSMEPeepholeOptPass() { return new SMEPeepholeOpt(); }
330