194 lines · cpp
1//===-- X86FastTileConfig.cpp - Fast 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. Before FastRegAllocation pass11/// the ldtilecfg 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 register 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/MachineFrameInfo.h"25#include "llvm/CodeGen/MachineFunctionPass.h"26#include "llvm/CodeGen/MachineInstr.h"27#include "llvm/CodeGen/MachineRegisterInfo.h"28#include "llvm/CodeGen/Passes.h"29#include "llvm/CodeGen/TargetInstrInfo.h"30#include "llvm/CodeGen/TargetRegisterInfo.h"31 32using namespace llvm;33 34#define DEBUG_TYPE "fasttileconfig"35 36namespace {37 38class X86FastTileConfig : public MachineFunctionPass {39 // context40 MachineFunction *MF = nullptr;41 const TargetInstrInfo *TII = nullptr;42 MachineRegisterInfo *MRI = nullptr;43 const TargetRegisterInfo *TRI = nullptr;44 X86MachineFunctionInfo *X86FI = nullptr;45 46 bool configBasicBlock(MachineBasicBlock &MBB);47 48public:49 X86FastTileConfig() : MachineFunctionPass(ID) {}50 51 /// Return the pass name.52 StringRef getPassName() const override {53 return "Fast Tile Register Configure";54 }55 56 void getAnalysisUsage(AnalysisUsage &AU) const override {57 AU.setPreservesAll();58 MachineFunctionPass::getAnalysisUsage(AU);59 }60 61 /// Perform register allocation.62 bool runOnMachineFunction(MachineFunction &MFunc) override;63 64 MachineFunctionProperties getRequiredProperties() const override {65 return MachineFunctionProperties().setNoPHIs();66 }67 68 static char ID;69};70 71} // end anonymous namespace72 73char X86FastTileConfig::ID = 0;74 75INITIALIZE_PASS_BEGIN(X86FastTileConfig, DEBUG_TYPE,76 "Fast Tile Register Configure", false, false)77INITIALIZE_PASS_END(X86FastTileConfig, DEBUG_TYPE,78 "Fast Tile Register Configure", false, false)79 80static bool isTileDef(MachineRegisterInfo *MRI, MachineInstr &MI) {81 // There is no phi instruction after register allocation.82 assert(MI.isPHI() == false);83 // The instruction must have 3 operands: tile def, row, col.84 // It should be AMX pseudo instruction that have shape operand.85 if (MI.isDebugInstr() || MI.isCopy() || MI.getNumOperands() < 3 ||86 !MI.isPseudo())87 return false;88 MachineOperand &MO = MI.getOperand(0);89 90 if (MO.isReg()) {91 Register Reg = MO.getReg();92 // FIXME: It may be used after Greedy RA and the physical93 // register is not rewritten yet.94 if (Reg.isVirtual()) {95 if (MRI->getRegClass(Reg)->getID() == X86::TILERegClassID)96 return true;97 }98 if (Reg >= X86::TMM0 && Reg <= X86::TMM7)99 return true;100 }101 102 return false;103}104 105static unsigned getTMMIndex(Register Reg) {106 if (Reg >= X86::TMM0 && Reg <= X86::TMM7)107 return Reg - X86::TMM0;108 llvm_unreachable("Invalid Tmm Reg!");109}110 111// PreTileConfig should configure the tile registers based on basic112// block.113bool X86FastTileConfig::configBasicBlock(MachineBasicBlock &MBB) {114 bool Change = false;115 SmallVector<std::pair<unsigned, ShapeT>, 6> ShapeInfos;116 for (MachineInstr &MI : reverse(MBB)) {117 if (!isTileDef(MRI, MI) && MI.getOpcode() != X86::PLDTILECFGV)118 continue;119 // AMX instructions that define tile register.120 if (MI.getOpcode() != X86::PLDTILECFGV) {121 MachineOperand &Row = MI.getOperand(1);122 unsigned TMMIdx = getTMMIndex(MI.getOperand(0).getReg());123 MachineOperand &Col = MI.getOperand(2);124 ShapeInfos.push_back({TMMIdx, ShapeT(&Row, &Col)});125 } else { // PLDTILECFGV126 // Rewrite the shape information to memory. Stack slot should have127 // been initialized to zero in pre config.128 int SS = MI.getOperand(0).getIndex(); // tile config stack slot.129 for (auto &ShapeInfo : ShapeInfos) {130 DebugLoc DL;131 unsigned TMMIdx = ShapeInfo.first;132 Register RowReg = ShapeInfo.second.getRow()->getReg();133 Register ColReg = ShapeInfo.second.getCol()->getReg();134 // Here is the data format for the tile config.135 // 0 palette136 // 1 start_row137 // 2-15 reserved, must be zero138 // 16-17 tile0.colsb Tile 0 bytes per row.139 // 18-19 tile1.colsb Tile 1 bytes per row.140 // 20-21 tile2.colsb Tile 2 bytes per row.141 // ... (sequence continues)142 // 30-31 tile7.colsb Tile 7 bytes per row.143 // 32-47 reserved, must be zero144 // 48 tile0.rows Tile 0 rows.145 // 49 tile1.rows Tile 1 rows.146 // 50 tile2.rows Tile 2 rows.147 // ... (sequence continues)148 // 55 tile7.rows Tile 7 rows.149 // 56-63 reserved, must be zero150 int RowOffset = 48 + TMMIdx;151 int ColOffset = 16 + TMMIdx * 2;152 153 Register SubRowReg = TRI->getSubReg(RowReg, X86::sub_8bit);154 BuildMI(MBB, MI, DL, TII->get(X86::IMPLICIT_DEF), SubRowReg);155 MachineInstrBuilder StoreRow =156 BuildMI(MBB, MI, DL, TII->get(X86::MOV8mr));157 addFrameReference(StoreRow, SS, RowOffset).addReg(SubRowReg);158 159 MachineInstrBuilder StoreCol =160 BuildMI(MBB, MI, DL, TII->get(X86::MOV16mr));161 addFrameReference(StoreCol, SS, ColOffset).addReg(ColReg);162 }163 ShapeInfos.clear();164 Change = true;165 }166 }167 168 return Change;169}170 171bool X86FastTileConfig::runOnMachineFunction(MachineFunction &MFunc) {172 X86FI = MFunc.getInfo<X86MachineFunctionInfo>();173 // Early exit in the common case of non-AMX code.174 if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)175 return false;176 177 MF = &MFunc;178 MRI = &MFunc.getRegInfo();179 const TargetSubtargetInfo *ST = &MFunc.getSubtarget<X86Subtarget>();180 TRI = ST->getRegisterInfo();181 TII = MFunc.getSubtarget().getInstrInfo();182 bool Change = false;183 184 // Loop over all of the basic blocks, eliminating virtual register references185 for (MachineBasicBlock &MBB : MFunc)186 Change |= configBasicBlock(MBB);187 188 return Change;189}190 191FunctionPass *llvm::createX86FastTileConfigPass() {192 return new X86FastTileConfig();193}194