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/ADT/SmallVector.h" 17 #include "llvm/CodeGen/MachineBasicBlock.h" 18 #include "llvm/CodeGen/MachineFunctionPass.h" 19 #include "llvm/CodeGen/MachineRegisterInfo.h" 20 #include "llvm/CodeGen/TargetRegisterInfo.h" 21 22 using namespace llvm; 23 24 #define DEBUG_TYPE "aarch64-sme-peephole-opt" 25 26 namespace { 27 28 struct SMEPeepholeOpt : public MachineFunctionPass { 29 static char ID; 30 31 SMEPeepholeOpt() : MachineFunctionPass(ID) { 32 initializeSMEPeepholeOptPass(*PassRegistry::getPassRegistry()); 33 } 34 35 bool runOnMachineFunction(MachineFunction &MF) override; 36 37 StringRef getPassName() const override { 38 return "SME Peephole Optimization pass"; 39 } 40 41 void getAnalysisUsage(AnalysisUsage &AU) const override { 42 AU.setPreservesCFG(); 43 MachineFunctionPass::getAnalysisUsage(AU); 44 } 45 46 bool optimizeStartStopPairs(MachineBasicBlock &MBB, 47 bool &HasRemovedAllSMChanges) const; 48 }; 49 50 char SMEPeepholeOpt::ID = 0; 51 52 } // end anonymous namespace 53 54 static bool isConditionalStartStop(const MachineInstr *MI) { 55 return MI->getOpcode() == AArch64::MSRpstatePseudo; 56 } 57 58 static bool isMatchingStartStopPair(const MachineInstr *MI1, 59 const MachineInstr *MI2) { 60 // We only consider the same type of streaming mode change here, i.e. 61 // start/stop SM, or start/stop ZA pairs. 62 if (MI1->getOperand(0).getImm() != MI2->getOperand(0).getImm()) 63 return false; 64 65 // One must be 'start', the other must be 'stop' 66 if (MI1->getOperand(1).getImm() == MI2->getOperand(1).getImm()) 67 return false; 68 69 bool IsConditional = isConditionalStartStop(MI2); 70 if (isConditionalStartStop(MI1) != IsConditional) 71 return false; 72 73 if (!IsConditional) 74 return true; 75 76 // Check to make sure the conditional start/stop pairs are identical. 77 if (MI1->getOperand(2).getImm() != MI2->getOperand(2).getImm()) 78 return false; 79 80 // Ensure reg masks are identical. 81 if (MI1->getOperand(4).getRegMask() != MI2->getOperand(4).getRegMask()) 82 return false; 83 84 // This optimisation is unlikely to happen in practice for conditional 85 // smstart/smstop pairs as the virtual registers for pstate.sm will always 86 // be different. 87 // TODO: For this optimisation to apply to conditional smstart/smstop, 88 // this pass will need to do more work to remove redundant calls to 89 // __arm_sme_state. 90 91 // Only consider conditional start/stop pairs which read the same register 92 // holding the original value of pstate.sm, as some conditional start/stops 93 // require the state on entry to the function. 94 if (MI1->getOperand(3).isReg() && MI2->getOperand(3).isReg()) { 95 Register Reg1 = MI1->getOperand(3).getReg(); 96 Register Reg2 = MI2->getOperand(3).getReg(); 97 if (Reg1.isPhysical() || Reg2.isPhysical() || Reg1 != Reg2) 98 return false; 99 } 100 101 return true; 102 } 103 104 static bool ChangesStreamingMode(const MachineInstr *MI) { 105 assert((MI->getOpcode() == AArch64::MSRpstatesvcrImm1 || 106 MI->getOpcode() == AArch64::MSRpstatePseudo) && 107 "Expected MI to be a smstart/smstop instruction"); 108 return MI->getOperand(0).getImm() == AArch64SVCR::SVCRSM || 109 MI->getOperand(0).getImm() == AArch64SVCR::SVCRSMZA; 110 } 111 112 static bool isSVERegOp(const TargetRegisterInfo &TRI, 113 const MachineRegisterInfo &MRI, 114 const MachineOperand &MO) { 115 if (!MO.isReg()) 116 return false; 117 118 Register R = MO.getReg(); 119 if (R.isPhysical()) 120 return llvm::any_of(TRI.subregs_inclusive(R), [](const MCPhysReg &SR) { 121 return AArch64::ZPRRegClass.contains(SR) || 122 AArch64::PPRRegClass.contains(SR); 123 }); 124 125 const TargetRegisterClass *RC = MRI.getRegClass(R); 126 return TRI.getCommonSubClass(&AArch64::ZPRRegClass, RC) || 127 TRI.getCommonSubClass(&AArch64::PPRRegClass, RC); 128 } 129 130 bool SMEPeepholeOpt::optimizeStartStopPairs( 131 MachineBasicBlock &MBB, bool &HasRemovedAllSMChanges) const { 132 const MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo(); 133 const TargetRegisterInfo &TRI = 134 *MBB.getParent()->getSubtarget().getRegisterInfo(); 135 136 bool Changed = false; 137 MachineInstr *Prev = nullptr; 138 SmallVector<MachineInstr *, 4> ToBeRemoved; 139 140 // Convenience function to reset the matching of a sequence. 141 auto Reset = [&]() { 142 Prev = nullptr; 143 ToBeRemoved.clear(); 144 }; 145 146 // Walk through instructions in the block trying to find pairs of smstart 147 // and smstop nodes that cancel each other out. We only permit a limited 148 // set of instructions to appear between them, otherwise we reset our 149 // tracking. 150 unsigned NumSMChanges = 0; 151 unsigned NumSMChangesRemoved = 0; 152 for (MachineInstr &MI : make_early_inc_range(MBB)) { 153 switch (MI.getOpcode()) { 154 case AArch64::MSRpstatesvcrImm1: 155 case AArch64::MSRpstatePseudo: { 156 if (ChangesStreamingMode(&MI)) 157 NumSMChanges++; 158 159 if (!Prev) 160 Prev = &MI; 161 else if (isMatchingStartStopPair(Prev, &MI)) { 162 // If they match, we can remove them, and possibly any instructions 163 // that we marked for deletion in between. 164 Prev->eraseFromParent(); 165 MI.eraseFromParent(); 166 for (MachineInstr *TBR : ToBeRemoved) 167 TBR->eraseFromParent(); 168 ToBeRemoved.clear(); 169 Prev = nullptr; 170 Changed = true; 171 NumSMChangesRemoved += 2; 172 } else { 173 Reset(); 174 Prev = &MI; 175 } 176 continue; 177 } 178 default: 179 if (!Prev) 180 // Avoid doing expensive checks when Prev is nullptr. 181 continue; 182 break; 183 } 184 185 // Test if the instructions in between the start/stop sequence are agnostic 186 // of streaming mode. If not, the algorithm should reset. 187 switch (MI.getOpcode()) { 188 default: 189 Reset(); 190 break; 191 case AArch64::COALESCER_BARRIER_FPR16: 192 case AArch64::COALESCER_BARRIER_FPR32: 193 case AArch64::COALESCER_BARRIER_FPR64: 194 case AArch64::COALESCER_BARRIER_FPR128: 195 case AArch64::COPY: 196 // These instructions should be safe when executed on their own, but 197 // the code remains conservative when SVE registers are used. There may 198 // exist subtle cases where executing a COPY in a different mode results 199 // in different behaviour, even if we can't yet come up with any 200 // concrete example/test-case. 201 if (isSVERegOp(TRI, MRI, MI.getOperand(0)) || 202 isSVERegOp(TRI, MRI, MI.getOperand(1))) 203 Reset(); 204 break; 205 case AArch64::ADJCALLSTACKDOWN: 206 case AArch64::ADJCALLSTACKUP: 207 case AArch64::ANDXri: 208 case AArch64::ADDXri: 209 // We permit these as they don't generate SVE/NEON instructions. 210 break; 211 case AArch64::VGRestorePseudo: 212 case AArch64::VGSavePseudo: 213 // When the smstart/smstop are removed, we should also remove 214 // the pseudos that save/restore the VG value for CFI info. 215 ToBeRemoved.push_back(&MI); 216 break; 217 case AArch64::MSRpstatesvcrImm1: 218 case AArch64::MSRpstatePseudo: 219 llvm_unreachable("Should have been handled"); 220 } 221 } 222 223 HasRemovedAllSMChanges = 224 NumSMChanges && (NumSMChanges == NumSMChangesRemoved); 225 return Changed; 226 } 227 228 INITIALIZE_PASS(SMEPeepholeOpt, "aarch64-sme-peephole-opt", 229 "SME Peephole Optimization", false, false) 230 231 bool SMEPeepholeOpt::runOnMachineFunction(MachineFunction &MF) { 232 if (skipFunction(MF.getFunction())) 233 return false; 234 235 if (!MF.getSubtarget<AArch64Subtarget>().hasSME()) 236 return false; 237 238 assert(MF.getRegInfo().isSSA() && "Expected to be run on SSA form!"); 239 240 bool Changed = false; 241 bool FunctionHasAllSMChangesRemoved = false; 242 243 // Even if the block lives in a function with no SME attributes attached we 244 // still have to analyze all the blocks because we may call a streaming 245 // function that requires smstart/smstop pairs. 246 for (MachineBasicBlock &MBB : MF) { 247 bool BlockHasAllSMChangesRemoved; 248 Changed |= optimizeStartStopPairs(MBB, BlockHasAllSMChangesRemoved); 249 FunctionHasAllSMChangesRemoved |= BlockHasAllSMChangesRemoved; 250 } 251 252 AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>(); 253 if (FunctionHasAllSMChangesRemoved) 254 AFI->setHasStreamingModeChanges(false); 255 256 return Changed; 257 } 258 259 FunctionPass *llvm::createSMEPeepholeOptPass() { return new SMEPeepholeOpt(); } 260