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