207 lines · cpp
1//===-- X86TileConfig.cpp - Tile Register 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 config the shape of AMX physical registers10/// AMX register need to be configured before use. In X86PreTileConfig pass11/// the pldtilecfg instruction is inserted, however at that time we don't12/// know the shape of each physical tile registers, because the register13/// allocation is not done yet. This pass runs after egister allocation14/// pass. It collects the shape information of each physical tile register15/// and store the shape in the stack slot that is allocated for load config16/// to tile config register.17//18//===----------------------------------------------------------------------===//19 20#include "X86.h"21#include "X86InstrBuilder.h"22#include "X86MachineFunctionInfo.h"23#include "X86Subtarget.h"24#include "llvm/CodeGen/LiveIntervals.h"25#include "llvm/CodeGen/MachineFrameInfo.h"26#include "llvm/CodeGen/MachineFunctionPass.h"27#include "llvm/CodeGen/MachineInstr.h"28#include "llvm/CodeGen/MachineRegisterInfo.h"29#include "llvm/CodeGen/Passes.h"30#include "llvm/CodeGen/TargetInstrInfo.h"31#include "llvm/CodeGen/TargetRegisterInfo.h"32#include "llvm/CodeGen/TileShapeInfo.h"33#include "llvm/CodeGen/VirtRegMap.h"34#include "llvm/InitializePasses.h"35 36using namespace llvm;37 38#define DEBUG_TYPE "tileconfig"39 40namespace {41 42struct X86TileConfig : public MachineFunctionPass {43 44 X86TileConfig() : MachineFunctionPass(ID) {}45 46 /// Return the pass name.47 StringRef getPassName() const override { return "Tile Register Configure"; }48 49 /// X86TileConfig analysis usage.50 void getAnalysisUsage(AnalysisUsage &AU) const override {51 AU.setPreservesAll();52 AU.addRequired<VirtRegMapWrapperLegacy>();53 AU.addRequired<LiveIntervalsWrapperPass>();54 MachineFunctionPass::getAnalysisUsage(AU);55 }56 57 /// Perform register allocation.58 bool runOnMachineFunction(MachineFunction &mf) override;59 60 MachineFunctionProperties getRequiredProperties() const override {61 return MachineFunctionProperties().setNoPHIs();62 }63 64 static char ID;65};66 67} // end anonymous namespace68 69char X86TileConfig::ID = 0;70 71INITIALIZE_PASS_BEGIN(X86TileConfig, DEBUG_TYPE, "Tile Register Configure",72 false, false)73INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)74INITIALIZE_PASS_END(X86TileConfig, DEBUG_TYPE, "Tile Register Configure", false,75 false)76 77bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {78 X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();79 // Early exit in the common case of non-AMX code.80 if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)81 return false;82 83 const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>();84 const X86RegisterInfo *TRI = ST.getRegisterInfo();85 const TargetInstrInfo *TII = ST.getInstrInfo();86 MachineRegisterInfo &MRI = MF.getRegInfo();87 LiveIntervals &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();88 VirtRegMap &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();89 90 if (VRM.isShapeMapEmpty())91 return false;92 93 int SS = INT_MAX;94 for (MachineBasicBlock &MBB : MF) {95 for (MachineInstr &MI : MBB) {96 if (MI.getOpcode() == X86::PLDTILECFGV) {97 SS = MI.getOperand(0).getIndex();98 break;99 }100 }101 if (SS != INT_MAX)102 break;103 }104 // Didn't find PLDTILECFGV, just return false;105 if (SS == INT_MAX)106 return false;107 108 // Try to find a point to insert MIs for constant shapes.109 // Here we are leveraging the palette id inserted in PreRA pass.110 unsigned ConstPos = 0;111 MachineInstr *ConstMI = nullptr;112 for (MachineInstr &MI : MF.front()) {113 if (MI.getOpcode() == X86::MOV8mi && SS == MI.getOperand(0).getIndex()) {114 ConstMI = &MI;115 break;116 }117 ++ConstPos;118 }119 assert(ConstMI && "Cannot find an insertion point");120 121 unsigned AMXRegNum = TRI->getRegClass(X86::TILERegClassID)->getNumRegs();122 SmallVector<Register, 8> Phys2Virt(AMXRegNum, 0);123 for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {124 Register VirtReg = Register::index2VirtReg(I);125 if (MRI.reg_nodbg_empty(VirtReg))126 continue;127 if (!TRI->isTileRegisterClass(MRI.getRegClass(VirtReg)))128 continue;129 MCRegister PhysReg = VRM.getPhys(VirtReg);130 if (!PhysReg)131 continue;132 unsigned Index = PhysReg - X86::TMM0;133 if (!Phys2Virt[Index])134 Phys2Virt[Index] = VirtReg;135 }136 137 // Fill in the shape of each tile physical register.138 for (unsigned I = 0; I < AMXRegNum; ++I) {139 if (!Phys2Virt[I])140 continue;141 DebugLoc DL;142 bool IsRow = true;143 MachineInstr *NewMI = nullptr;144 ShapeT Shape = VRM.getShape(Phys2Virt[I]);145 for (auto &R : {Shape.getRow()->getReg(), Shape.getCol()->getReg()}) {146 // Here is the data format for the tile config.147 // 0 palette148 // 1 start_row149 // 2-15 reserved, must be zero150 // 16-17 tile0.colsb Tile 0 bytes per row.151 // 18-19 tile1.colsb Tile 1 bytes per row.152 // 20-21 tile2.colsb Tile 2 bytes per row.153 // ... (sequence continues)154 // 30-31 tile7.colsb Tile 7 bytes per row.155 // 32-47 reserved, must be zero156 // 48 tile0.rows Tile 0 rows.157 // 49 tile1.rows Tile 1 rows.158 // 50 tile2.rows Tile 2 rows.159 // ... (sequence continues)160 // 55 tile7.rows Tile 7 rows.161 // 56-63 reserved, must be zero162 int64_t Imm = INT64_MAX;163 int Offset = IsRow ? 48 + I : 16 + I * 2;164 for (auto &DefMI : MRI.def_instructions(R)) {165 MachineBasicBlock &MBB = *DefMI.getParent();166 if (DefMI.isMoveImmediate()) {167 if (Imm != INT64_MAX) {168 // FIXME: We should handle this case in future.169 assert(Imm == DefMI.getOperand(1).getImm() &&170 "Cannot initialize with different shapes");171 continue;172 }173 Imm = DefMI.getOperand(1).getImm();174 175 NewMI = addFrameReference(176 BuildMI(MF.front(), ++ConstMI->getIterator(), DL,177 TII->get(IsRow ? X86::MOV8mi : X86::MOV16mi)),178 SS, Offset)179 .addImm(Imm);180 ConstMI = NewMI;181 LIS.InsertMachineInstrInMaps(*NewMI);182 } else {183 unsigned SubIdx = IsRow ? X86::sub_8bit : X86::sub_16bit;184 unsigned RegSize = TRI->getRegSizeInBits(*MRI.getRegClass(R));185 if ((IsRow && RegSize == 8) || (!IsRow && RegSize == 16))186 SubIdx = 0;187 auto Iter = DefMI.getIterator();188 if (&MBB == &MF.front() &&189 (unsigned)std::distance(MBB.instr_begin(), Iter) < ConstPos)190 Iter = ConstMI->getIterator();191 NewMI = addFrameReference(192 BuildMI(MBB, ++Iter, DL,193 TII->get(IsRow ? X86::MOV8mr : X86::MOV16mr)),194 SS, Offset)195 .addReg(R, 0, SubIdx);196 SlotIndex SIdx = LIS.InsertMachineInstrInMaps(*NewMI);197 LIS.extendToIndices(LIS.getInterval(R), {SIdx.getRegSlot()});198 }199 }200 IsRow = false;201 }202 }203 return true;204}205 206FunctionPass *llvm::createX86TileConfigPass() { return new X86TileConfig(); }207