1 //===-- X86PreTileConfig.cpp - Tile Register Pre-configure-----------------===// 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 /// \file Pass to pre-config the shapes of AMX registers 10 /// AMX register needs to be configured before use. The shapes of AMX register 11 /// are encoded in the 1st and 2nd machine operand of AMX pseudo instructions. 12 /// 13 /// The instruction ldtilecfg is used to config the shapes. It must be reachable 14 /// for all variable shapes. ldtilecfg will be inserted more than once if we 15 /// cannot find a dominating point for all AMX instructions. 16 /// 17 /// The configure register is caller saved according to ABI. We need to insert 18 /// ldtilecfg again after the call instruction if callee clobbers any AMX 19 /// registers. 20 /// 21 /// This pass calculates all points that ldtilecfg need to be inserted to and 22 /// insert them. It reports error if the reachability conditions aren't met. 23 // 24 //===----------------------------------------------------------------------===// 25 26 #include "X86.h" 27 #include "X86InstrBuilder.h" 28 #include "X86RegisterInfo.h" 29 #include "X86Subtarget.h" 30 #include "llvm/CodeGen/MachineFunctionPass.h" 31 #include "llvm/CodeGen/MachineInstr.h" 32 #include "llvm/CodeGen/MachineLoopInfo.h" 33 #include "llvm/CodeGen/MachineRegisterInfo.h" 34 #include "llvm/CodeGen/Passes.h" 35 #include "llvm/CodeGen/TargetInstrInfo.h" 36 #include "llvm/CodeGen/TargetRegisterInfo.h" 37 #include "llvm/InitializePasses.h" 38 39 using namespace llvm; 40 41 #define DEBUG_TYPE "tile-pre-config" 42 #define REPORT_CONFIG_FAIL \ 43 report_fatal_error( \ 44 MF.getName() + \ 45 ": Failed to config tile register, please define the shape earlier"); 46 47 namespace { 48 49 struct MIRef { 50 MachineInstr *MI = nullptr; 51 MachineBasicBlock *MBB = nullptr; 52 // A virtual position for instruction that will be inserted after MI. 53 size_t Pos = 0; 54 MIRef() = default; 55 MIRef(MachineBasicBlock *MBB) : MBB(MBB) { 56 for (auto I = MBB->begin(), E = MBB->end(); I != E && I->isPHI(); 57 ++I, ++Pos) 58 MI = &*I; 59 } 60 MIRef(MachineInstr *MI) 61 : MI(MI), MBB(MI->getParent()), 62 Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {} 63 MIRef(MachineInstr *MI, MachineBasicBlock *MBB) 64 : MI(MI), MBB(MBB), 65 Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {} 66 MIRef(MachineInstr *MI, MachineBasicBlock *MBB, size_t Pos) 67 : MI(MI), MBB(MBB), Pos(Pos) {} 68 operator bool() const { return MBB != nullptr; } 69 bool operator==(const MIRef &RHS) const { 70 return MI == RHS.MI && MBB == RHS.MBB; 71 } 72 bool operator!=(const MIRef &RHS) const { return !(*this == RHS); } 73 bool operator<(const MIRef &RHS) const { 74 return MBB < RHS.MBB || (MBB == RHS.MBB && Pos < RHS.Pos); 75 } 76 bool operator>(const MIRef &RHS) const { 77 return MBB > RHS.MBB || (MBB == RHS.MBB && Pos > RHS.Pos); 78 } 79 }; 80 81 struct BBInfo { 82 MIRef FirstAMX; 83 MIRef LastCall; 84 bool HasAMXRegLiveIn = false; 85 bool TileCfgForbidden = false; 86 bool NeedTileCfgLiveIn = false; 87 }; 88 89 class X86PreTileConfig : public MachineFunctionPass { 90 MachineRegisterInfo *MRI; 91 const MachineLoopInfo *MLI; 92 SmallSet<MachineInstr *, 8> DefVisited; 93 DenseMap<MachineBasicBlock *, BBInfo> BBVisitedInfo; 94 DenseMap<MachineBasicBlock *, SmallVector<MIRef, 8>> ShapeBBs; 95 96 /// Check if the callee will clobber AMX registers. 97 bool isDestructiveCall(MachineInstr &MI, BitVector UsableRegs) { 98 auto Iter = llvm::find_if( 99 MI.operands(), [](MachineOperand &MO) { return MO.isRegMask(); }); 100 if (Iter == MI.operands_end()) 101 return false; 102 UsableRegs.clearBitsInMask(Iter->getRegMask()); 103 return !UsableRegs.none(); 104 } 105 106 /// Check if MI is AMX pseudo instruction. 107 bool isAMXInstruction(MachineInstr &MI) { 108 if (MI.isPHI() || MI.isDebugInstr() || MI.getNumOperands() < 3) 109 return false; 110 MachineOperand &MO = MI.getOperand(0); 111 // We can simply check if it is AMX instruction by its def. 112 // But we should exclude old API which uses physical registers. 113 if (MO.isReg() && MO.getReg().isVirtual() && 114 MRI->getRegClass(MO.getReg())->getID() == X86::TILERegClassID) { 115 collectShapeInfo(MI); 116 return true; 117 } 118 // PTILESTOREDV is the only exception that doesn't def a AMX register. 119 return MI.getOpcode() == X86::PTILESTOREDV; 120 } 121 122 /// Check if it is an edge from loop bottom to loop head. 123 bool isLoopBackEdge(MachineBasicBlock *Header, MachineBasicBlock *Bottom) { 124 return MLI->isLoopHeader(Header) && 125 MLI->getLoopFor(Header)->getBottomBlock() == Bottom; 126 } 127 128 /// Collect the shape def information for later use. 129 void collectShapeInfo(MachineInstr &MI); 130 131 /// Try to hoist shapes definded below AMX instructions. 132 bool hoistShapesInBB(MachineBasicBlock *MBB, SmallVectorImpl<MIRef> &Shapes) { 133 MIRef &FirstAMX = BBVisitedInfo[MBB].FirstAMX; 134 auto FirstShapeBelowAMX = llvm::lower_bound(Shapes, FirstAMX); 135 auto InsertPoint = FirstAMX.MI->getIterator(); 136 for (auto I = FirstShapeBelowAMX, E = Shapes.end(); I != E; ++I) { 137 // Do not hoist instructions that access memory. 138 if (I->MI->mayLoadOrStore()) 139 return false; 140 for (auto &MO : I->MI->operands()) { 141 if (MO.isDef()) 142 continue; 143 // Do not hoist instructions if the sources' def under AMX instruction. 144 // TODO: We can handle isMoveImmediate MI here. 145 if (MO.isReg() && MIRef(MRI->getVRegDef(MO.getReg())) > FirstAMX) 146 return false; 147 // TODO: Maybe need more checks here. 148 } 149 MBB->insert(InsertPoint, I->MI->removeFromParent()); 150 } 151 // We only need to mark the last shape in the BB now. 152 Shapes.clear(); 153 Shapes.push_back(MIRef(&*--InsertPoint, MBB)); 154 return true; 155 } 156 157 public: 158 X86PreTileConfig() : MachineFunctionPass(ID) {} 159 160 /// Return the pass name. 161 StringRef getPassName() const override { 162 return "Tile Register Pre-configure"; 163 } 164 165 /// X86PreTileConfig analysis usage. 166 void getAnalysisUsage(AnalysisUsage &AU) const override { 167 AU.setPreservesAll(); 168 AU.addRequired<MachineLoopInfo>(); 169 MachineFunctionPass::getAnalysisUsage(AU); 170 } 171 172 /// Clear MF related structures. 173 void releaseMemory() override { 174 ShapeBBs.clear(); 175 DefVisited.clear(); 176 BBVisitedInfo.clear(); 177 } 178 179 /// Perform ldtilecfg instructions inserting. 180 bool runOnMachineFunction(MachineFunction &MF) override; 181 182 static char ID; 183 }; 184 185 } // end anonymous namespace 186 187 char X86PreTileConfig::ID = 0; 188 189 INITIALIZE_PASS_BEGIN(X86PreTileConfig, "tilepreconfig", 190 "Tile Register Pre-configure", false, false) 191 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo) 192 INITIALIZE_PASS_END(X86PreTileConfig, "tilepreconfig", 193 "Tile Register Pre-configure", false, false) 194 195 void X86PreTileConfig::collectShapeInfo(MachineInstr &MI) { 196 auto RecordShape = [&](MachineInstr *MI, MachineBasicBlock *MBB) { 197 MIRef MIR(MI, MBB); 198 auto I = llvm::lower_bound(ShapeBBs[MBB], MIR); 199 if (I == ShapeBBs[MBB].end() || *I != MIR) 200 ShapeBBs[MBB].insert(I, MIR); 201 }; 202 203 SmallVector<Register, 8> WorkList( 204 {MI.getOperand(1).getReg(), MI.getOperand(2).getReg()}); 205 while (!WorkList.empty()) { 206 Register R = WorkList.pop_back_val(); 207 MachineInstr *DefMI = MRI->getVRegDef(R); 208 MachineBasicBlock *DefMBB = DefMI->getParent(); 209 if (!DefMI || DefMI->isMoveImmediate() || !DefVisited.insert(DefMI).second) 210 continue; 211 if (DefMI->isPHI()) { 212 for (unsigned I = 1; I < DefMI->getNumOperands(); I += 2) 213 if (isLoopBackEdge(DefMBB, DefMI->getOperand(I + 1).getMBB())) 214 RecordShape(DefMI, DefMBB); // In this case, PHI is also a shape def. 215 else 216 WorkList.push_back(DefMI->getOperand(I).getReg()); 217 } else { 218 RecordShape(DefMI, DefMBB); 219 } 220 } 221 } 222 223 bool X86PreTileConfig::runOnMachineFunction(MachineFunction &MF) { 224 const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>(); 225 const TargetInstrInfo *TII = ST.getInstrInfo(); 226 const TargetRegisterInfo *TRI = ST.getRegisterInfo(); 227 const TargetRegisterClass *RC = TRI->getRegClass(X86::TILERegClassID); 228 229 BitVector AMXRegs(TRI->getNumRegs()); 230 for (unsigned I = 0; I < RC->getNumRegs(); I++) 231 AMXRegs.set(X86::TMM0 + I); 232 233 // Iterate MF to collect information. 234 MRI = &MF.getRegInfo(); 235 MLI = &getAnalysis<MachineLoopInfo>(); 236 SmallSet<MIRef, 8> CfgNeedInsert; 237 SmallVector<MachineBasicBlock *, 8> CfgLiveInBBs; 238 for (auto &MBB : MF) { 239 size_t Pos = 0; 240 for (auto &MI : MBB) { 241 ++Pos; 242 if (isAMXInstruction(MI)) { 243 // If there's call before the AMX, we need to reload tile config. 244 if (BBVisitedInfo[&MBB].LastCall) 245 CfgNeedInsert.insert(BBVisitedInfo[&MBB].LastCall); 246 else // Otherwise, we need tile config to live in this BB. 247 BBVisitedInfo[&MBB].NeedTileCfgLiveIn = true; 248 // Always record the first AMX in case there's shape def after it. 249 if (!BBVisitedInfo[&MBB].FirstAMX) 250 BBVisitedInfo[&MBB].FirstAMX = MIRef(&MI, &MBB, Pos); 251 } else if (MI.isCall() && isDestructiveCall(MI, AMXRegs)) { 252 // Record the call only if the callee clobbers all AMX registers. 253 BBVisitedInfo[&MBB].LastCall = MIRef(&MI, &MBB, Pos); 254 } 255 } 256 if (BBVisitedInfo[&MBB].NeedTileCfgLiveIn) { 257 if (&MBB == &MF.front()) 258 CfgNeedInsert.insert(MIRef(&MBB)); 259 else 260 CfgLiveInBBs.push_back(&MBB); 261 } 262 if (BBVisitedInfo[&MBB].FirstAMX || BBVisitedInfo[&MBB].HasAMXRegLiveIn) 263 for (auto *Succ : MBB.successors()) 264 if (!isLoopBackEdge(Succ, &MBB)) 265 BBVisitedInfo[Succ].HasAMXRegLiveIn = true; 266 } 267 268 // Update NeedTileCfgLiveIn for predecessors. 269 while (!CfgLiveInBBs.empty()) { 270 MachineBasicBlock *MBB = CfgLiveInBBs.pop_back_val(); 271 for (auto *Pred : MBB->predecessors()) { 272 if (BBVisitedInfo[Pred].LastCall) { 273 CfgNeedInsert.insert(BBVisitedInfo[Pred].LastCall); 274 } else if (!BBVisitedInfo[Pred].NeedTileCfgLiveIn) { 275 BBVisitedInfo[Pred].NeedTileCfgLiveIn = true; 276 if (Pred == &MF.front()) 277 CfgNeedInsert.insert(MIRef(Pred)); 278 else 279 CfgLiveInBBs.push_back(Pred); 280 } 281 } 282 } 283 284 // There's no AMX instruction if we didn't find a tile config live in point. 285 if (CfgNeedInsert.empty()) 286 return false; 287 288 // Avoid to insert ldtilecfg before any shape defs. 289 SmallVector<MachineBasicBlock *, 8> WorkList; 290 for (auto &I : ShapeBBs) { 291 // TODO: We can hoist shapes across BBs here. 292 if (BBVisitedInfo[I.first].HasAMXRegLiveIn) 293 REPORT_CONFIG_FAIL 294 if (BBVisitedInfo[I.first].FirstAMX && 295 BBVisitedInfo[I.first].FirstAMX < I.second.back() && 296 !hoistShapesInBB(I.first, I.second)) 297 REPORT_CONFIG_FAIL 298 WorkList.push_back(I.first); 299 } 300 while (!WorkList.empty()) { 301 MachineBasicBlock *MBB = WorkList.pop_back_val(); 302 for (auto *Pred : MBB->predecessors()) { 303 if (!BBVisitedInfo[Pred].TileCfgForbidden && !isLoopBackEdge(MBB, Pred)) { 304 BBVisitedInfo[Pred].TileCfgForbidden = true; 305 WorkList.push_back(Pred); 306 } 307 } 308 } 309 310 DebugLoc DL; 311 SmallSet<MIRef, 8> VisitedOrInserted; 312 int SS = MF.getFrameInfo().CreateStackObject( 313 ST.getTileConfigSize(), ST.getTileConfigAlignment(), false); 314 315 // Try to insert for the tile config live in points. 316 for (auto I : CfgNeedInsert) { 317 SmallSet<MIRef, 8> InsertPoints; 318 SmallVector<MIRef, 8> WorkList({I}); 319 while (!WorkList.empty()) { 320 MIRef I = WorkList.pop_back_val(); 321 if (!VisitedOrInserted.count(I)) { 322 if (!BBVisitedInfo[I.MBB].TileCfgForbidden) { 323 // If the BB is all shapes reachable, stop sink and try to insert. 324 InsertPoints.insert(I); 325 } else { 326 // Avoid the BB to be multi visited. 327 VisitedOrInserted.insert(I); 328 // Sink the inserting point along the chain with NeedTileCfgLiveIn = 329 // true when MBB isn't all shapes reachable. 330 for (auto *Succ : I.MBB->successors()) 331 if (BBVisitedInfo[Succ].NeedTileCfgLiveIn) 332 WorkList.push_back(MIRef(Succ)); 333 } 334 } 335 } 336 337 // A given point might be forked due to shape conditions are not met. 338 for (MIRef I : InsertPoints) { 339 // Make sure we insert ldtilecfg after the last shape def in MBB. 340 if (ShapeBBs.count(I.MBB) && I < ShapeBBs[I.MBB].back()) 341 I = ShapeBBs[I.MBB].back(); 342 // There're chances the MBB is sunk more than once. Record it to avoid 343 // multi insert. 344 if (VisitedOrInserted.insert(I).second) { 345 auto II = I.MI ? I.MI->getIterator() : I.MBB->instr_begin(); 346 addFrameReference(BuildMI(*I.MBB, ++II, DL, TII->get(X86::LDTILECFG)), 347 SS); 348 } 349 } 350 } 351 352 // Zero stack slot. 353 MachineBasicBlock &MBB = MF.front(); 354 MachineInstr *MI = &*MBB.begin(); 355 if (ST.hasAVX512()) { 356 Register Zmm = MRI->createVirtualRegister(&X86::VR512RegClass); 357 BuildMI(MBB, MI, DL, TII->get(X86::VPXORDZrr), Zmm) 358 .addReg(Zmm, RegState::Undef) 359 .addReg(Zmm, RegState::Undef); 360 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSZmr)), SS) 361 .addReg(Zmm); 362 } else if (ST.hasAVX2()) { 363 Register Ymm = MRI->createVirtualRegister(&X86::VR256RegClass); 364 BuildMI(MBB, MI, DL, TII->get(X86::VPXORYrr), Ymm) 365 .addReg(Ymm, RegState::Undef) 366 .addReg(Ymm, RegState::Undef); 367 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSYmr)), SS) 368 .addReg(Ymm); 369 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSYmr)), SS, 32) 370 .addReg(Ymm); 371 } else { 372 assert(ST.hasSSE2() && "AMX should assume SSE2 enabled"); 373 Register Xmm = MRI->createVirtualRegister(&X86::VR128RegClass); 374 BuildMI(MBB, MI, DL, TII->get(X86::PXORrr), Xmm) 375 .addReg(Xmm, RegState::Undef) 376 .addReg(Xmm, RegState::Undef); 377 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS) 378 .addReg(Xmm); 379 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS, 16) 380 .addReg(Xmm); 381 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS, 32) 382 .addReg(Xmm); 383 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS, 48) 384 .addReg(Xmm); 385 } 386 // Fill in the palette first. 387 addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOV8mi)), SS).addImm(1); 388 389 return true; 390 } 391 392 FunctionPass *llvm::createX86PreTileConfigPass() { 393 return new X86PreTileConfig(); 394 } 395