676 lines · cpp
1//===- MIR2Vec.cpp - Implementation of MIR2Vec ---------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM4// Exceptions. See the LICENSE file for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8///9/// \file10/// This file implements the MIR2Vec algorithm for Machine IR embeddings.11///12//===----------------------------------------------------------------------===//13 14#include "llvm/CodeGen/MIR2Vec.h"15#include "llvm/ADT/DepthFirstIterator.h"16#include "llvm/ADT/Statistic.h"17#include "llvm/CodeGen/TargetInstrInfo.h"18#include "llvm/IR/Module.h"19#include "llvm/InitializePasses.h"20#include "llvm/Pass.h"21#include "llvm/Support/Errc.h"22#include "llvm/Support/MemoryBuffer.h"23#include "llvm/Support/Regex.h"24 25using namespace llvm;26using namespace mir2vec;27 28#define DEBUG_TYPE "mir2vec"29 30STATISTIC(MIRVocabMissCounter,31 "Number of lookups to MIR entities not present in the vocabulary");32 33namespace llvm {34namespace mir2vec {35cl::OptionCategory MIR2VecCategory("MIR2Vec Options");36 37// FIXME: Use a default vocab when not specified38static cl::opt<std::string>39 VocabFile("mir2vec-vocab-path", cl::Optional,40 cl::desc("Path to the vocabulary file for MIR2Vec"), cl::init(""),41 cl::cat(MIR2VecCategory));42cl::opt<float> OpcWeight("mir2vec-opc-weight", cl::Optional, cl::init(1.0),43 cl::desc("Weight for machine opcode embeddings"),44 cl::cat(MIR2VecCategory));45cl::opt<float> CommonOperandWeight(46 "mir2vec-common-operand-weight", cl::Optional, cl::init(1.0),47 cl::desc("Weight for common operand embeddings"), cl::cat(MIR2VecCategory));48cl::opt<float>49 RegOperandWeight("mir2vec-reg-operand-weight", cl::Optional, cl::init(1.0),50 cl::desc("Weight for register operand embeddings"),51 cl::cat(MIR2VecCategory));52cl::opt<MIR2VecKind> MIR2VecEmbeddingKind(53 "mir2vec-kind", cl::Optional,54 cl::values(clEnumValN(MIR2VecKind::Symbolic, "symbolic",55 "Generate symbolic embeddings for MIR")),56 cl::init(MIR2VecKind::Symbolic), cl::desc("MIR2Vec embedding kind"),57 cl::cat(MIR2VecCategory));58 59} // namespace mir2vec60} // namespace llvm61 62//===----------------------------------------------------------------------===//63// Vocabulary64//===----------------------------------------------------------------------===//65 66MIRVocabulary::MIRVocabulary(VocabMap &&OpcodeMap, VocabMap &&CommonOperandMap,67 VocabMap &&PhysicalRegisterMap,68 VocabMap &&VirtualRegisterMap,69 const TargetInstrInfo &TII,70 const TargetRegisterInfo &TRI,71 const MachineRegisterInfo &MRI)72 : TII(TII), TRI(TRI), MRI(MRI) {73 buildCanonicalOpcodeMapping();74 unsigned CanonicalOpcodeCount = UniqueBaseOpcodeNames.size();75 assert(CanonicalOpcodeCount > 0 &&76 "No canonical opcodes found for target - invalid vocabulary");77 78 buildRegisterOperandMapping();79 80 // Define layout of vocabulary sections81 Layout.OpcodeBase = 0;82 Layout.CommonOperandBase = CanonicalOpcodeCount;83 // We expect same classes for physical and virtual registers84 Layout.PhyRegBase = Layout.CommonOperandBase + std::size(CommonOperandNames);85 Layout.VirtRegBase = Layout.PhyRegBase + RegisterOperandNames.size();86 87 generateStorage(OpcodeMap, CommonOperandMap, PhysicalRegisterMap,88 VirtualRegisterMap);89 Layout.TotalEntries = Storage.size();90}91 92Expected<MIRVocabulary>93MIRVocabulary::create(VocabMap &&OpcodeMap, VocabMap &&CommonOperandMap,94 VocabMap &&PhyRegMap, VocabMap &&VirtRegMap,95 const TargetInstrInfo &TII, const TargetRegisterInfo &TRI,96 const MachineRegisterInfo &MRI) {97 if (OpcodeMap.empty() || CommonOperandMap.empty() || PhyRegMap.empty() ||98 VirtRegMap.empty())99 return createStringError(errc::invalid_argument,100 "Empty vocabulary entries provided");101 102 MIRVocabulary Vocab(std::move(OpcodeMap), std::move(CommonOperandMap),103 std::move(PhyRegMap), std::move(VirtRegMap), TII, TRI,104 MRI);105 106 // Validate Storage after construction107 if (!Vocab.Storage.isValid())108 return createStringError(errc::invalid_argument,109 "Failed to create valid vocabulary storage");110 Vocab.ZeroEmbedding = Embedding(Vocab.Storage.getDimension(), 0.0);111 return std::move(Vocab);112}113 114std::string MIRVocabulary::extractBaseOpcodeName(StringRef InstrName) {115 // Extract base instruction name using regex to capture letters and116 // underscores Examples: "ADD32rr" -> "ADD", "ARITH_FENCE" -> "ARITH_FENCE"117 //118 // TODO: Consider more sophisticated extraction:119 // - Handle complex prefixes like "AVX1_SETALLONES" correctly (Currently, it120 // would naively map to "AVX")121 // - Extract width suffixes (8,16,32,64) as separate features122 // - Capture addressing mode suffixes (r,i,m,ri,etc.) for better analysis123 // (Currently, instances like "MOV32mi" map to "MOV", but "ADDPDrr" would map124 // to "ADDPDrr")125 126 assert(!InstrName.empty() && "Instruction name should not be empty");127 128 // Use regex to extract initial sequence of letters and underscores129 static const Regex BaseOpcodeRegex("([a-zA-Z_]+)");130 SmallVector<StringRef, 2> Matches;131 132 if (BaseOpcodeRegex.match(InstrName, &Matches) && Matches.size() > 1) {133 StringRef Match = Matches[1];134 // Trim trailing underscores135 while (!Match.empty() && Match.back() == '_')136 Match = Match.drop_back();137 return Match.str();138 }139 140 // Fallback to original name if no pattern matches141 return InstrName.str();142}143 144unsigned MIRVocabulary::getCanonicalIndexForBaseName(StringRef BaseName) const {145 assert(!UniqueBaseOpcodeNames.empty() && "Canonical mapping not built");146 auto It = std::find(UniqueBaseOpcodeNames.begin(),147 UniqueBaseOpcodeNames.end(), BaseName.str());148 assert(It != UniqueBaseOpcodeNames.end() &&149 "Base name not found in unique opcodes");150 return std::distance(UniqueBaseOpcodeNames.begin(), It);151}152 153unsigned MIRVocabulary::getCanonicalOpcodeIndex(unsigned Opcode) const {154 auto BaseOpcode = extractBaseOpcodeName(TII.getName(Opcode));155 return getCanonicalIndexForBaseName(BaseOpcode);156}157 158unsigned159MIRVocabulary::getCanonicalIndexForOperandName(StringRef OperandName) const {160 auto It = std::find(std::begin(CommonOperandNames),161 std::end(CommonOperandNames), OperandName);162 assert(It != std::end(CommonOperandNames) &&163 "Operand name not found in common operands");164 return Layout.CommonOperandBase +165 std::distance(std::begin(CommonOperandNames), It);166}167 168unsigned169MIRVocabulary::getCanonicalIndexForRegisterClass(StringRef RegName,170 bool IsPhysical) const {171 auto It = std::find(RegisterOperandNames.begin(), RegisterOperandNames.end(),172 RegName);173 assert(It != RegisterOperandNames.end() &&174 "Register name not found in register operands");175 unsigned LocalIndex = std::distance(RegisterOperandNames.begin(), It);176 return (IsPhysical ? Layout.PhyRegBase : Layout.VirtRegBase) + LocalIndex;177}178 179std::string MIRVocabulary::getStringKey(unsigned Pos) const {180 assert(Pos < Layout.TotalEntries && "Position out of bounds in vocabulary");181 182 // Handle opcodes section183 if (Pos < Layout.CommonOperandBase) {184 // Convert canonical index back to base opcode name185 auto It = UniqueBaseOpcodeNames.begin();186 std::advance(It, Pos);187 assert(It != UniqueBaseOpcodeNames.end() &&188 "Canonical index out of bounds in opcode section");189 return *It;190 }191 192 auto getLocalIndex = [](unsigned Pos, size_t BaseOffset, size_t Bound,193 const char *Msg) {194 unsigned LocalIndex = Pos - BaseOffset;195 assert(LocalIndex < Bound && Msg);196 return LocalIndex;197 };198 199 // Handle common operands section200 if (Pos < Layout.PhyRegBase) {201 unsigned LocalIndex = getLocalIndex(202 Pos, Layout.CommonOperandBase, std::size(CommonOperandNames),203 "Local index out of bounds in common operands");204 return CommonOperandNames[LocalIndex].str();205 }206 207 // Handle physical registers section208 if (Pos < Layout.VirtRegBase) {209 unsigned LocalIndex =210 getLocalIndex(Pos, Layout.PhyRegBase, RegisterOperandNames.size(),211 "Local index out of bounds in physical registers");212 return "PhyReg_" + RegisterOperandNames[LocalIndex];213 }214 215 // Handle virtual registers section216 unsigned LocalIndex =217 getLocalIndex(Pos, Layout.VirtRegBase, RegisterOperandNames.size(),218 "Local index out of bounds in virtual registers");219 return "VirtReg_" + RegisterOperandNames[LocalIndex];220}221 222void MIRVocabulary::generateStorage(const VocabMap &OpcodeMap,223 const VocabMap &CommonOperandsMap,224 const VocabMap &PhyRegMap,225 const VocabMap &VirtRegMap) {226 227 // Helper for handling missing entities in the vocabulary.228 // Currently, we use a zero vector. In the future, we will throw an error to229 // ensure that *all* known entities are present in the vocabulary.230 auto handleMissingEntity = [](StringRef Key) {231 LLVM_DEBUG(errs() << "MIR2Vec: Missing vocabulary entry for " << Key232 << "; using zero vector. This will result in an error "233 "in the future.\n");234 ++MIRVocabMissCounter;235 };236 237 // Initialize opcode embeddings section238 unsigned EmbeddingDim = OpcodeMap.begin()->second.size();239 std::vector<Embedding> OpcodeEmbeddings(Layout.CommonOperandBase,240 Embedding(EmbeddingDim));241 242 // Populate opcode embeddings using canonical mapping243 for (auto COpcodeName : UniqueBaseOpcodeNames) {244 if (auto It = OpcodeMap.find(COpcodeName); It != OpcodeMap.end()) {245 auto COpcodeIndex = getCanonicalIndexForBaseName(COpcodeName);246 assert(COpcodeIndex < Layout.CommonOperandBase &&247 "Canonical index out of bounds");248 OpcodeEmbeddings[COpcodeIndex] = It->second;249 } else {250 handleMissingEntity(COpcodeName);251 }252 }253 254 // Initialize common operand embeddings section255 std::vector<Embedding> CommonOperandEmbeddings(std::size(CommonOperandNames),256 Embedding(EmbeddingDim));257 unsigned OperandIndex = 0;258 for (const auto &CommonOperandName : CommonOperandNames) {259 if (auto It = CommonOperandsMap.find(CommonOperandName.str());260 It != CommonOperandsMap.end()) {261 CommonOperandEmbeddings[OperandIndex] = It->second;262 } else {263 handleMissingEntity(CommonOperandName);264 }265 ++OperandIndex;266 }267 268 // Helper lambda for creating register operand embeddings269 auto createRegisterEmbeddings = [&](const VocabMap &RegMap) {270 std::vector<Embedding> RegEmbeddings(TRI.getNumRegClasses(),271 Embedding(EmbeddingDim));272 unsigned RegOperandIndex = 0;273 for (const auto &RegOperandName : RegisterOperandNames) {274 if (auto It = RegMap.find(RegOperandName); It != RegMap.end())275 RegEmbeddings[RegOperandIndex] = It->second;276 else277 handleMissingEntity(RegOperandName);278 ++RegOperandIndex;279 }280 return RegEmbeddings;281 };282 283 // Initialize register operand embeddings sections284 std::vector<Embedding> PhyRegEmbeddings = createRegisterEmbeddings(PhyRegMap);285 std::vector<Embedding> VirtRegEmbeddings =286 createRegisterEmbeddings(VirtRegMap);287 288 // Scale the vocabulary sections based on the provided weights289 auto scaleVocabSection = [](std::vector<Embedding> &Embeddings,290 double Weight) {291 for (auto &Embedding : Embeddings)292 Embedding *= Weight;293 };294 scaleVocabSection(OpcodeEmbeddings, OpcWeight);295 scaleVocabSection(CommonOperandEmbeddings, CommonOperandWeight);296 scaleVocabSection(PhyRegEmbeddings, RegOperandWeight);297 scaleVocabSection(VirtRegEmbeddings, RegOperandWeight);298 299 std::vector<std::vector<Embedding>> Sections(300 static_cast<unsigned>(Section::MaxSections));301 Sections[static_cast<unsigned>(Section::Opcodes)] =302 std::move(OpcodeEmbeddings);303 Sections[static_cast<unsigned>(Section::CommonOperands)] =304 std::move(CommonOperandEmbeddings);305 Sections[static_cast<unsigned>(Section::PhyRegisters)] =306 std::move(PhyRegEmbeddings);307 Sections[static_cast<unsigned>(Section::VirtRegisters)] =308 std::move(VirtRegEmbeddings);309 310 Storage = ir2vec::VocabStorage(std::move(Sections));311}312 313void MIRVocabulary::buildCanonicalOpcodeMapping() {314 // Check if already built315 if (!UniqueBaseOpcodeNames.empty())316 return;317 318 // Build mapping from opcodes to canonical base opcode indices319 for (unsigned Opcode = 0; Opcode < TII.getNumOpcodes(); ++Opcode) {320 std::string BaseOpcode = extractBaseOpcodeName(TII.getName(Opcode));321 UniqueBaseOpcodeNames.insert(BaseOpcode);322 }323 324 LLVM_DEBUG(dbgs() << "MIR2Vec: Built canonical mapping for target with "325 << UniqueBaseOpcodeNames.size()326 << " unique base opcodes\n");327}328 329void MIRVocabulary::buildRegisterOperandMapping() {330 // Check if already built331 if (!RegisterOperandNames.empty())332 return;333 334 for (unsigned RC = 0; RC < TRI.getNumRegClasses(); ++RC) {335 const TargetRegisterClass *RegClass = TRI.getRegClass(RC);336 if (!RegClass)337 continue;338 339 // Get the register class name340 StringRef ClassName = TRI.getRegClassName(RegClass);341 RegisterOperandNames.push_back(ClassName.str());342 }343}344 345unsigned MIRVocabulary::getCommonOperandIndex(346 MachineOperand::MachineOperandType OperandType) const {347 assert(OperandType != MachineOperand::MO_Register &&348 "Expected non-register operand type");349 assert(OperandType > MachineOperand::MO_Register &&350 OperandType < MachineOperand::MO_Last && "Operand type out of bounds");351 return static_cast<unsigned>(OperandType) - 1;352}353 354unsigned MIRVocabulary::getRegisterOperandIndex(Register Reg) const {355 assert(!RegisterOperandNames.empty() && "Register operand mapping not built");356 assert(Reg.isValid() && "Invalid register; not expected here");357 assert((Reg.isPhysical() || Reg.isVirtual()) &&358 "Expected a physical or virtual register");359 360 const TargetRegisterClass *RegClass = nullptr;361 362 // For physical registers, use TRI to get minimal register class as a363 // physical register can belong to multiple classes. For virtual364 // registers, use MRI to uniquely identify the assigned register class.365 if (Reg.isPhysical())366 RegClass = TRI.getMinimalPhysRegClass(Reg);367 else368 RegClass = MRI.getRegClass(Reg);369 370 if (RegClass)371 return RegClass->getID();372 // Fallback for registers without a class (shouldn't happen)373 llvm_unreachable("Register operand without a valid register class");374 return 0;375}376 377Expected<MIRVocabulary> MIRVocabulary::createDummyVocabForTest(378 const TargetInstrInfo &TII, const TargetRegisterInfo &TRI,379 const MachineRegisterInfo &MRI, unsigned Dim) {380 assert(Dim > 0 && "Dimension must be greater than zero");381 382 float DummyVal = 0.1f;383 384 VocabMap DummyOpcMap, DummyOperandMap, DummyPhyRegMap, DummyVirtRegMap;385 386 // Process opcodes directly without creating temporary vocabulary387 for (unsigned Opcode = 0; Opcode < TII.getNumOpcodes(); ++Opcode) {388 std::string BaseOpcode = extractBaseOpcodeName(TII.getName(Opcode));389 if (DummyOpcMap.count(BaseOpcode) == 0) { // Only add if not already present390 DummyOpcMap[BaseOpcode] = Embedding(Dim, DummyVal);391 DummyVal += 0.1f;392 }393 }394 395 // Add common operands396 for (const auto &CommonOperandName : CommonOperandNames) {397 DummyOperandMap[CommonOperandName.str()] = Embedding(Dim, DummyVal);398 DummyVal += 0.1f;399 }400 401 // Process register classes directly402 for (unsigned RC = 0; RC < TRI.getNumRegClasses(); ++RC) {403 const TargetRegisterClass *RegClass = TRI.getRegClass(RC);404 if (!RegClass)405 continue;406 407 std::string ClassName = TRI.getRegClassName(RegClass);408 DummyPhyRegMap[ClassName] = Embedding(Dim, DummyVal);409 DummyVirtRegMap[ClassName] = Embedding(Dim, DummyVal);410 DummyVal += 0.1f;411 }412 413 // Create vocabulary directly without temporary instance414 return MIRVocabulary::create(415 std::move(DummyOpcMap), std::move(DummyOperandMap),416 std::move(DummyPhyRegMap), std::move(DummyVirtRegMap), TII, TRI, MRI);417}418 419//===----------------------------------------------------------------------===//420// MIR2VecVocabProvider and MIR2VecVocabLegacyAnalysis421//===----------------------------------------------------------------------===//422 423Expected<mir2vec::MIRVocabulary>424MIR2VecVocabProvider::getVocabulary(const Module &M) {425 VocabMap OpcVocab, CommonOperandVocab, PhyRegVocabMap, VirtRegVocabMap;426 427 if (Error Err = readVocabulary(OpcVocab, CommonOperandVocab, PhyRegVocabMap,428 VirtRegVocabMap))429 return std::move(Err);430 431 for (const auto &F : M) {432 if (F.isDeclaration())433 continue;434 435 if (auto *MF = MMI.getMachineFunction(F)) {436 auto &Subtarget = MF->getSubtarget();437 if (const auto *TII = Subtarget.getInstrInfo())438 if (const auto *TRI = Subtarget.getRegisterInfo())439 return mir2vec::MIRVocabulary::create(440 std::move(OpcVocab), std::move(CommonOperandVocab),441 std::move(PhyRegVocabMap), std::move(VirtRegVocabMap), *TII, *TRI,442 MF->getRegInfo());443 }444 }445 return createStringError(errc::invalid_argument,446 "No machine functions found in module");447}448 449Error MIR2VecVocabProvider::readVocabulary(VocabMap &OpcodeVocab,450 VocabMap &CommonOperandVocab,451 VocabMap &PhyRegVocabMap,452 VocabMap &VirtRegVocabMap) {453 if (VocabFile.empty())454 return createStringError(455 errc::invalid_argument,456 "MIR2Vec vocabulary file path not specified; set it "457 "using --mir2vec-vocab-path");458 459 auto BufOrError = MemoryBuffer::getFileOrSTDIN(VocabFile, /*IsText=*/true);460 if (!BufOrError)461 return createFileError(VocabFile, BufOrError.getError());462 463 auto Content = BufOrError.get()->getBuffer();464 465 Expected<json::Value> ParsedVocabValue = json::parse(Content);466 if (!ParsedVocabValue)467 return ParsedVocabValue.takeError();468 469 unsigned OpcodeDim = 0, CommonOperandDim = 0, PhyRegOperandDim = 0,470 VirtRegOperandDim = 0;471 if (auto Err = ir2vec::VocabStorage::parseVocabSection(472 "Opcodes", *ParsedVocabValue, OpcodeVocab, OpcodeDim))473 return Err;474 475 if (auto Err = ir2vec::VocabStorage::parseVocabSection(476 "CommonOperands", *ParsedVocabValue, CommonOperandVocab,477 CommonOperandDim))478 return Err;479 480 if (auto Err = ir2vec::VocabStorage::parseVocabSection(481 "PhysicalRegisters", *ParsedVocabValue, PhyRegVocabMap,482 PhyRegOperandDim))483 return Err;484 485 if (auto Err = ir2vec::VocabStorage::parseVocabSection(486 "VirtualRegisters", *ParsedVocabValue, VirtRegVocabMap,487 VirtRegOperandDim))488 return Err;489 490 // All sections must have the same embedding dimension491 if (!(OpcodeDim == CommonOperandDim && CommonOperandDim == PhyRegOperandDim &&492 PhyRegOperandDim == VirtRegOperandDim)) {493 return createStringError(494 errc::illegal_byte_sequence,495 "MIR2Vec vocabulary sections have different dimensions");496 }497 498 return Error::success();499}500 501char MIR2VecVocabLegacyAnalysis::ID = 0;502INITIALIZE_PASS_BEGIN(MIR2VecVocabLegacyAnalysis, "mir2vec-vocab-analysis",503 "MIR2Vec Vocabulary Analysis", false, true)504INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)505INITIALIZE_PASS_END(MIR2VecVocabLegacyAnalysis, "mir2vec-vocab-analysis",506 "MIR2Vec Vocabulary Analysis", false, true)507 508StringRef MIR2VecVocabLegacyAnalysis::getPassName() const {509 return "MIR2Vec Vocabulary Analysis";510}511 512//===----------------------------------------------------------------------===//513// MIREmbedder and its subclasses514//===----------------------------------------------------------------------===//515 516std::unique_ptr<MIREmbedder> MIREmbedder::create(MIR2VecKind Mode,517 const MachineFunction &MF,518 const MIRVocabulary &Vocab) {519 switch (Mode) {520 case MIR2VecKind::Symbolic:521 return std::make_unique<SymbolicMIREmbedder>(MF, Vocab);522 }523 return nullptr;524}525 526Embedding MIREmbedder::computeEmbeddings(const MachineBasicBlock &MBB) const {527 Embedding MBBVector(Dimension, 0);528 529 // Get instruction info for opcode name resolution530 const auto &Subtarget = MF.getSubtarget();531 const auto *TII = Subtarget.getInstrInfo();532 if (!TII) {533 MF.getFunction().getContext().emitError(534 "MIR2Vec: No TargetInstrInfo available; cannot compute embeddings");535 return MBBVector;536 }537 538 // Process each machine instruction in the basic block539 for (const auto &MI : MBB) {540 // Skip debug instructions and other metadata541 if (MI.isDebugInstr())542 continue;543 MBBVector += computeEmbeddings(MI);544 }545 546 return MBBVector;547}548 549Embedding MIREmbedder::computeEmbeddings() const {550 Embedding MFuncVector(Dimension, 0);551 552 // Consider all reachable machine basic blocks in the function553 for (const auto *MBB : depth_first(&MF))554 MFuncVector += computeEmbeddings(*MBB);555 return MFuncVector;556}557 558SymbolicMIREmbedder::SymbolicMIREmbedder(const MachineFunction &MF,559 const MIRVocabulary &Vocab)560 : MIREmbedder(MF, Vocab) {}561 562std::unique_ptr<SymbolicMIREmbedder>563SymbolicMIREmbedder::create(const MachineFunction &MF,564 const MIRVocabulary &Vocab) {565 return std::make_unique<SymbolicMIREmbedder>(MF, Vocab);566}567 568Embedding SymbolicMIREmbedder::computeEmbeddings(const MachineInstr &MI) const {569 // Skip debug instructions and other metadata570 if (MI.isDebugInstr())571 return Embedding(Dimension, 0);572 573 // Opcode embedding574 Embedding InstructionEmbedding = Vocab[MI.getOpcode()];575 576 // Add operand contributions577 for (const MachineOperand &MO : MI.operands())578 InstructionEmbedding += Vocab[MO];579 580 return InstructionEmbedding;581}582 583//===----------------------------------------------------------------------===//584// Printer Passes585//===----------------------------------------------------------------------===//586 587char MIR2VecVocabPrinterLegacyPass::ID = 0;588INITIALIZE_PASS_BEGIN(MIR2VecVocabPrinterLegacyPass, "print-mir2vec-vocab",589 "MIR2Vec Vocabulary Printer Pass", false, true)590INITIALIZE_PASS_DEPENDENCY(MIR2VecVocabLegacyAnalysis)591INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)592INITIALIZE_PASS_END(MIR2VecVocabPrinterLegacyPass, "print-mir2vec-vocab",593 "MIR2Vec Vocabulary Printer Pass", false, true)594 595bool MIR2VecVocabPrinterLegacyPass::runOnMachineFunction(MachineFunction &MF) {596 return false;597}598 599bool MIR2VecVocabPrinterLegacyPass::doFinalization(Module &M) {600 auto &Analysis = getAnalysis<MIR2VecVocabLegacyAnalysis>();601 auto MIR2VecVocabOrErr = Analysis.getMIR2VecVocabulary(M);602 603 if (!MIR2VecVocabOrErr) {604 OS << "MIR2Vec Vocabulary Printer: Failed to get vocabulary - "605 << toString(MIR2VecVocabOrErr.takeError()) << "\n";606 return false;607 }608 609 auto &MIR2VecVocab = *MIR2VecVocabOrErr;610 unsigned Pos = 0;611 for (const auto &Entry : MIR2VecVocab) {612 OS << "Key: " << MIR2VecVocab.getStringKey(Pos++) << ": ";613 Entry.print(OS);614 }615 616 return false;617}618 619MachineFunctionPass *620llvm::createMIR2VecVocabPrinterLegacyPass(raw_ostream &OS) {621 return new MIR2VecVocabPrinterLegacyPass(OS);622}623 624char MIR2VecPrinterLegacyPass::ID = 0;625INITIALIZE_PASS_BEGIN(MIR2VecPrinterLegacyPass, "print-mir2vec",626 "MIR2Vec Embedder Printer Pass", false, true)627INITIALIZE_PASS_DEPENDENCY(MIR2VecVocabLegacyAnalysis)628INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)629INITIALIZE_PASS_END(MIR2VecPrinterLegacyPass, "print-mir2vec",630 "MIR2Vec Embedder Printer Pass", false, true)631 632bool MIR2VecPrinterLegacyPass::runOnMachineFunction(MachineFunction &MF) {633 auto &Analysis = getAnalysis<MIR2VecVocabLegacyAnalysis>();634 auto VocabOrErr =635 Analysis.getMIR2VecVocabulary(*MF.getFunction().getParent());636 assert(VocabOrErr && "Failed to get MIR2Vec vocabulary");637 auto &MIRVocab = *VocabOrErr;638 639 auto Emb = mir2vec::MIREmbedder::create(MIR2VecEmbeddingKind, MF, MIRVocab);640 if (!Emb) {641 OS << "Error creating MIR2Vec embeddings for function " << MF.getName()642 << "\n";643 return false;644 }645 646 OS << "MIR2Vec embeddings for machine function " << MF.getName() << ":\n";647 OS << "Machine Function vector: ";648 Emb->getMFunctionVector().print(OS);649 650 OS << "Machine basic block vectors:\n";651 for (const MachineBasicBlock &MBB : MF) {652 OS << "Machine basic block: " << MBB.getFullName() << ":\n";653 Emb->getMBBVector(MBB).print(OS);654 }655 656 OS << "Machine instruction vectors:\n";657 for (const MachineBasicBlock &MBB : MF) {658 for (const MachineInstr &MI : MBB) {659 // Skip debug instructions as they are not660 // embedded661 if (MI.isDebugInstr())662 continue;663 664 OS << "Machine instruction: ";665 MI.print(OS);666 Emb->getMInstVector(MI).print(OS);667 }668 }669 670 return false;671}672 673MachineFunctionPass *llvm::createMIR2VecPrinterLegacyPass(raw_ostream &OS) {674 return new MIR2VecPrinterLegacyPass(OS);675}676