266 lines
8.7 KiB
C++
266 lines
8.7 KiB
C++
|
//===-- X86PreTileConfig.cpp - Tile Register Configure---------------------===//
|
||
|
|
//
|
||
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
||
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
||
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
//
|
||
|
|
/// \file Pass to pre-config the shape of AMX register
|
||
|
|
/// AMX register need to be configured before use. The shape of AMX register
|
||
|
|
/// is encoded in the 1st and 2nd machine operand of AMX pseudo instructions.
|
||
|
|
/// The pldtilecfg is to config tile registers. It should dominator all AMX
|
||
|
|
/// instructions. The pldtilecfg produce a virtual cfg register and the cfg
|
||
|
|
/// register is used by all AMX instructions.
|
||
|
|
/// This pass is to find the common dominator of all AMX instructions and
|
||
|
|
/// insert the pldtilecfg instruction. Besides the cfg register that pldtilecfg
|
||
|
|
/// produces is inserted as the last operand of each AMX instruction. We use
|
||
|
|
/// this scheme to model the def-use relationship between AMX config instruction
|
||
|
|
/// and other AMX instructions. Below is an example.
|
||
|
|
///
|
||
|
|
/// ----B1----
|
||
|
|
/// / \
|
||
|
|
/// / \
|
||
|
|
/// B2 B3
|
||
|
|
/// %1:tile = PTILELOADDV %2:tile = PTILELOADDV
|
||
|
|
///
|
||
|
|
/// is transformed to
|
||
|
|
///
|
||
|
|
/// B1
|
||
|
|
/// %25:tilecfg = PLDTILECFG
|
||
|
|
/// / \
|
||
|
|
/// / \
|
||
|
|
/// %1:tile = PTILELOADDV %25 %2:tile = PTILELOADDV %25
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
|
||
|
|
#include "X86.h"
|
||
|
|
#include "X86InstrBuilder.h"
|
||
|
|
#include "X86RegisterInfo.h"
|
||
|
|
#include "X86Subtarget.h"
|
||
|
|
#include "llvm/CodeGen/MachineDominators.h"
|
||
|
|
#include "llvm/CodeGen/MachineFunctionPass.h"
|
||
|
|
#include "llvm/CodeGen/MachineInstr.h"
|
||
|
|
#include "llvm/CodeGen/MachineRegisterInfo.h"
|
||
|
|
#include "llvm/CodeGen/Passes.h"
|
||
|
|
#include "llvm/CodeGen/TargetInstrInfo.h"
|
||
|
|
#include "llvm/CodeGen/TargetRegisterInfo.h"
|
||
|
|
#include "llvm/CodeGen/TileShapeInfo.h"
|
||
|
|
#include "llvm/InitializePasses.h"
|
||
|
|
|
||
|
|
using namespace llvm;
|
||
|
|
|
||
|
|
#define DEBUG_TYPE "tile-pre-config"
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
|
||
|
|
class X86PreTileConfig : public MachineFunctionPass {
|
||
|
|
// context
|
||
|
|
MachineFunction *MF = nullptr;
|
||
|
|
const X86Subtarget *ST = nullptr;
|
||
|
|
const TargetRegisterInfo *TRI;
|
||
|
|
const TargetInstrInfo *TII;
|
||
|
|
MachineDominatorTree *DomTree = nullptr;
|
||
|
|
MachineRegisterInfo *MRI = nullptr;
|
||
|
|
|
||
|
|
MachineInstr *getTileConfigPoint();
|
||
|
|
|
||
|
|
public:
|
||
|
|
X86PreTileConfig() : MachineFunctionPass(ID) {}
|
||
|
|
|
||
|
|
/// Return the pass name.
|
||
|
|
StringRef getPassName() const override {
|
||
|
|
return "Tile Register Pre-configure";
|
||
|
|
}
|
||
|
|
|
||
|
|
/// X86PreTileConfig analysis usage.
|
||
|
|
void getAnalysisUsage(AnalysisUsage &AU) const override;
|
||
|
|
|
||
|
|
/// Perform register allocation.
|
||
|
|
bool runOnMachineFunction(MachineFunction &mf) override;
|
||
|
|
|
||
|
|
static char ID;
|
||
|
|
};
|
||
|
|
|
||
|
|
} // end anonymous namespace
|
||
|
|
|
||
|
|
char X86PreTileConfig::ID = 0;
|
||
|
|
|
||
|
|
INITIALIZE_PASS_BEGIN(X86PreTileConfig, "tilepreconfig",
|
||
|
|
"Tile Register Configure", false, false)
|
||
|
|
INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
|
||
|
|
INITIALIZE_PASS_END(X86PreTileConfig, "tilepreconfig",
|
||
|
|
"Tile Register Configure", false, false)
|
||
|
|
|
||
|
|
void X86PreTileConfig::getAnalysisUsage(AnalysisUsage &AU) const {
|
||
|
|
AU.setPreservesAll();
|
||
|
|
AU.addRequired<MachineDominatorTree>();
|
||
|
|
MachineFunctionPass::getAnalysisUsage(AU);
|
||
|
|
}
|
||
|
|
|
||
|
|
static Register buildConfigMI(MachineBasicBlock::iterator MI, int FrameIdx,
|
||
|
|
const TargetInstrInfo *TII,
|
||
|
|
MachineRegisterInfo *MRI,
|
||
|
|
const X86Subtarget *ST) {
|
||
|
|
auto *MBB = MI->getParent();
|
||
|
|
|
||
|
|
// FIXME: AMX should assume AVX512 enabled.
|
||
|
|
if (ST->hasAVX512()) {
|
||
|
|
// Zero stack slot.
|
||
|
|
Register Zmm = MRI->createVirtualRegister(&X86::VR512RegClass);
|
||
|
|
BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VPXORDZrr), Zmm)
|
||
|
|
.addReg(Zmm, RegState::Undef)
|
||
|
|
.addReg(Zmm, RegState::Undef);
|
||
|
|
addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VMOVUPSZmr)),
|
||
|
|
FrameIdx)
|
||
|
|
.addReg(Zmm);
|
||
|
|
}
|
||
|
|
|
||
|
|
// build psuedo ldtilecfg
|
||
|
|
Register VReg = MRI->createVirtualRegister(&X86::TILECFGRegClass);
|
||
|
|
|
||
|
|
addFrameReference(
|
||
|
|
BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::PLDTILECFG), VReg), FrameIdx);
|
||
|
|
|
||
|
|
return VReg;
|
||
|
|
}
|
||
|
|
|
||
|
|
static ShapeT getShape(const MachineInstr &MI, MachineRegisterInfo *MRI) {
|
||
|
|
unsigned Opcode = MI.getOpcode();
|
||
|
|
switch (Opcode) {
|
||
|
|
default:
|
||
|
|
llvm_unreachable("Unexpected machine instruction on tile");
|
||
|
|
case X86::PTILELOADDV:
|
||
|
|
case X86::PTDPBSSDV:
|
||
|
|
case X86::PTILEZEROV:
|
||
|
|
MachineOperand &MO1 = const_cast<MachineOperand &>(MI.getOperand(1));
|
||
|
|
MachineOperand &MO2 = const_cast<MachineOperand &>(MI.getOperand(2));
|
||
|
|
ShapeT Shape(&MO1, &MO2, MRI);
|
||
|
|
return Shape;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
MachineInstr *X86PreTileConfig::getTileConfigPoint() {
|
||
|
|
DenseMap<Register, ShapeT> PhysShapeInfo;
|
||
|
|
MachineBasicBlock *MBB = nullptr;
|
||
|
|
DenseSet<const MachineInstr *> MIs;
|
||
|
|
for (unsigned i = 0, e = MRI->getNumVirtRegs(); i != e; ++i) {
|
||
|
|
Register VirtReg = Register::index2VirtReg(i);
|
||
|
|
if (MRI->reg_nodbg_empty(VirtReg))
|
||
|
|
continue;
|
||
|
|
const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
|
||
|
|
if (RC.getID() != X86::TILERegClassID)
|
||
|
|
continue;
|
||
|
|
|
||
|
|
// Find the common dominator for all MI that define tile register.
|
||
|
|
for (const MachineOperand &MO : MRI->def_operands(VirtReg)) {
|
||
|
|
if (MO.isUndef())
|
||
|
|
continue;
|
||
|
|
const auto *MI = MO.getParent();
|
||
|
|
// PHI or IMPLICIT_DEF instructiion.
|
||
|
|
// There must be a input tile before PHI instruction.
|
||
|
|
if (MI->isTransient())
|
||
|
|
continue;
|
||
|
|
if (!MBB)
|
||
|
|
MBB = const_cast<MachineBasicBlock *>(MI->getParent());
|
||
|
|
MBB = DomTree->findNearestCommonDominator(
|
||
|
|
MBB, const_cast<MachineBasicBlock *>(MI->getParent()));
|
||
|
|
|
||
|
|
// Collect the instructions that define shape.
|
||
|
|
ShapeT Shape = getShape(*MI, MRI);
|
||
|
|
std::array<MachineOperand *, 2> ShapeMOs = {Shape.getRow(),
|
||
|
|
Shape.getCol()};
|
||
|
|
for (auto *ShapeMO : ShapeMOs) {
|
||
|
|
Register ShapeReg = ShapeMO->getReg();
|
||
|
|
for (const MachineOperand &MO : MRI->def_operands(ShapeReg)) {
|
||
|
|
const auto *ShapeMI = MO.getParent();
|
||
|
|
MIs.insert(ShapeMI);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (!MBB)
|
||
|
|
return nullptr;
|
||
|
|
// This pass is before the pass of eliminating PHI node, so it
|
||
|
|
// is in SSA form.
|
||
|
|
assert(MRI->isSSA() && "Not SSA form in pre-tile config");
|
||
|
|
// Shape def should dominate tile config MBB.
|
||
|
|
// def s s1 s2
|
||
|
|
// / \ \ /
|
||
|
|
// / \ \ /
|
||
|
|
// conf s3=phi(s1,s2)
|
||
|
|
// |
|
||
|
|
// c
|
||
|
|
//
|
||
|
|
for (const auto *MI : MIs) {
|
||
|
|
const MachineBasicBlock *ShapeMBB = MI->getParent();
|
||
|
|
if (DomTree->dominates(ShapeMBB, MBB))
|
||
|
|
continue;
|
||
|
|
if (MI->isMoveImmediate())
|
||
|
|
continue;
|
||
|
|
report_fatal_error(MF->getName() + ": Failed to config tile register, "
|
||
|
|
"please define the shape earlier");
|
||
|
|
}
|
||
|
|
|
||
|
|
// ldtilecfg should be inserted after the MI that define the shape.
|
||
|
|
MachineBasicBlock::reverse_instr_iterator I, E;
|
||
|
|
for (I = MBB->instr_rbegin(), E = MBB->instr_rend(); I != E; ++I) {
|
||
|
|
auto *MI = &*I;
|
||
|
|
if (MIs.count(MI) && (!MI->isMoveImmediate()))
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
MachineBasicBlock::iterator MII;
|
||
|
|
if (I == E)
|
||
|
|
MII = MBB->getFirstNonPHI();
|
||
|
|
else {
|
||
|
|
MII = MachineBasicBlock::iterator(&*I);
|
||
|
|
MII++;
|
||
|
|
}
|
||
|
|
return &*MII;
|
||
|
|
}
|
||
|
|
|
||
|
|
static void addTileCFGUse(MachineFunction &MF, Register CFG) {
|
||
|
|
for (MachineBasicBlock &MBB : MF) {
|
||
|
|
|
||
|
|
// Traverse the basic block.
|
||
|
|
for (MachineInstr &MI : MBB) {
|
||
|
|
unsigned Opcode = MI.getOpcode();
|
||
|
|
switch (Opcode) {
|
||
|
|
default:
|
||
|
|
break;
|
||
|
|
case X86::PTILELOADDV:
|
||
|
|
case X86::PTILESTOREDV:
|
||
|
|
case X86::PTDPBSSDV:
|
||
|
|
case X86::PTILEZEROV:
|
||
|
|
unsigned NumOperands = MI.getNumOperands();
|
||
|
|
MI.RemoveOperand(NumOperands - 1);
|
||
|
|
MI.addOperand(MF, MachineOperand::CreateReg(CFG, false));
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
bool X86PreTileConfig::runOnMachineFunction(MachineFunction &mf) {
|
||
|
|
MF = &mf;
|
||
|
|
MRI = &mf.getRegInfo();
|
||
|
|
ST = &mf.getSubtarget<X86Subtarget>();
|
||
|
|
TRI = ST->getRegisterInfo();
|
||
|
|
TII = mf.getSubtarget().getInstrInfo();
|
||
|
|
DomTree = &getAnalysis<MachineDominatorTree>();
|
||
|
|
|
||
|
|
MachineInstr *MI = getTileConfigPoint();
|
||
|
|
if (!MI)
|
||
|
|
return false;
|
||
|
|
unsigned Size = ST->getTileConfigSize();
|
||
|
|
Align Alignment = ST->getTileConfigAlignment();
|
||
|
|
int SS = mf.getFrameInfo().CreateStackObject(Size, Alignment, false);
|
||
|
|
Register CFG = buildConfigMI(MI, SS, TII, MRI, ST);
|
||
|
|
addTileCFGUse(mf, CFG);
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
FunctionPass *llvm::createX86PreTileConfigPass() {
|
||
|
|
return new X86PreTileConfig();
|
||
|
|
}
|