698 lines · cpp
1//===- IR2Vec.cpp - Implementation of IR2Vec -----------------------------===//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 IR2Vec algorithm.11///12//===----------------------------------------------------------------------===//13 14#include "llvm/Analysis/IR2Vec.h"15 16#include "llvm/ADT/DepthFirstIterator.h"17#include "llvm/ADT/Sequence.h"18#include "llvm/ADT/SmallVector.h"19#include "llvm/ADT/Statistic.h"20#include "llvm/IR/CFG.h"21#include "llvm/IR/Module.h"22#include "llvm/IR/PassManager.h"23#include "llvm/Support/Debug.h"24#include "llvm/Support/Errc.h"25#include "llvm/Support/Error.h"26#include "llvm/Support/ErrorHandling.h"27#include "llvm/Support/Format.h"28#include "llvm/Support/MemoryBuffer.h"29 30using namespace llvm;31using namespace ir2vec;32 33#define DEBUG_TYPE "ir2vec"34 35STATISTIC(VocabMissCounter,36 "Number of lookups to entities not present in the vocabulary");37 38namespace llvm {39namespace ir2vec {40cl::OptionCategory IR2VecCategory("IR2Vec Options");41 42// FIXME: Use a default vocab when not specified43static cl::opt<std::string>44 VocabFile("ir2vec-vocab-path", cl::Optional,45 cl::desc("Path to the vocabulary file for IR2Vec"), cl::init(""),46 cl::cat(IR2VecCategory));47cl::opt<float> OpcWeight("ir2vec-opc-weight", cl::Optional, cl::init(1.0),48 cl::desc("Weight for opcode embeddings"),49 cl::cat(IR2VecCategory));50cl::opt<float> TypeWeight("ir2vec-type-weight", cl::Optional, cl::init(0.5),51 cl::desc("Weight for type embeddings"),52 cl::cat(IR2VecCategory));53cl::opt<float> ArgWeight("ir2vec-arg-weight", cl::Optional, cl::init(0.2),54 cl::desc("Weight for argument embeddings"),55 cl::cat(IR2VecCategory));56cl::opt<IR2VecKind> IR2VecEmbeddingKind(57 "ir2vec-kind", cl::Optional,58 cl::values(clEnumValN(IR2VecKind::Symbolic, "symbolic",59 "Generate symbolic embeddings"),60 clEnumValN(IR2VecKind::FlowAware, "flow-aware",61 "Generate flow-aware embeddings")),62 cl::init(IR2VecKind::Symbolic), cl::desc("IR2Vec embedding kind"),63 cl::cat(IR2VecCategory));64 65} // namespace ir2vec66} // namespace llvm67 68AnalysisKey IR2VecVocabAnalysis::Key;69 70// ==----------------------------------------------------------------------===//71// Local helper functions72//===----------------------------------------------------------------------===//73namespace llvm::json {74inline bool fromJSON(const llvm::json::Value &E, Embedding &Out,75 llvm::json::Path P) {76 std::vector<double> TempOut;77 if (!llvm::json::fromJSON(E, TempOut, P))78 return false;79 Out = Embedding(std::move(TempOut));80 return true;81}82} // namespace llvm::json83 84// ==----------------------------------------------------------------------===//85// Embedding86//===----------------------------------------------------------------------===//87Embedding &Embedding::operator+=(const Embedding &RHS) {88 assert(this->size() == RHS.size() && "Vectors must have the same dimension");89 std::transform(this->begin(), this->end(), RHS.begin(), this->begin(),90 std::plus<double>());91 return *this;92}93 94Embedding Embedding::operator+(const Embedding &RHS) const {95 Embedding Result(*this);96 Result += RHS;97 return Result;98}99 100Embedding &Embedding::operator-=(const Embedding &RHS) {101 assert(this->size() == RHS.size() && "Vectors must have the same dimension");102 std::transform(this->begin(), this->end(), RHS.begin(), this->begin(),103 std::minus<double>());104 return *this;105}106 107Embedding Embedding::operator-(const Embedding &RHS) const {108 Embedding Result(*this);109 Result -= RHS;110 return Result;111}112 113Embedding &Embedding::operator*=(double Factor) {114 std::transform(this->begin(), this->end(), this->begin(),115 [Factor](double Elem) { return Elem * Factor; });116 return *this;117}118 119Embedding Embedding::operator*(double Factor) const {120 Embedding Result(*this);121 Result *= Factor;122 return Result;123}124 125Embedding &Embedding::scaleAndAdd(const Embedding &Src, float Factor) {126 assert(this->size() == Src.size() && "Vectors must have the same dimension");127 for (size_t Itr = 0; Itr < this->size(); ++Itr)128 (*this)[Itr] += Src[Itr] * Factor;129 return *this;130}131 132bool Embedding::approximatelyEquals(const Embedding &RHS,133 double Tolerance) const {134 assert(this->size() == RHS.size() && "Vectors must have the same dimension");135 for (size_t Itr = 0; Itr < this->size(); ++Itr)136 if (std::abs((*this)[Itr] - RHS[Itr]) > Tolerance) {137 LLVM_DEBUG(errs() << "Embedding mismatch at index " << Itr << ": "138 << (*this)[Itr] << " vs " << RHS[Itr]139 << "; Tolerance: " << Tolerance << "\n");140 return false;141 }142 return true;143}144 145void Embedding::print(raw_ostream &OS) const {146 OS << " [";147 for (const auto &Elem : Data)148 OS << " " << format("%.2f", Elem) << " ";149 OS << "]\n";150}151 152// ==----------------------------------------------------------------------===//153// Embedder and its subclasses154//===----------------------------------------------------------------------===//155 156std::unique_ptr<Embedder> Embedder::create(IR2VecKind Mode, const Function &F,157 const Vocabulary &Vocab) {158 switch (Mode) {159 case IR2VecKind::Symbolic:160 return std::make_unique<SymbolicEmbedder>(F, Vocab);161 case IR2VecKind::FlowAware:162 return std::make_unique<FlowAwareEmbedder>(F, Vocab);163 }164 return nullptr;165}166 167Embedding Embedder::computeEmbeddings() const {168 Embedding FuncVector(Dimension, 0.0);169 170 if (F.isDeclaration())171 return FuncVector;172 173 // Consider only the basic blocks that are reachable from entry174 for (const BasicBlock *BB : depth_first(&F))175 FuncVector += computeEmbeddings(*BB);176 return FuncVector;177}178 179Embedding Embedder::computeEmbeddings(const BasicBlock &BB) const {180 Embedding BBVector(Dimension, 0);181 182 // We consider only the non-debug and non-pseudo instructions183 for (const auto &I : BB.instructionsWithoutDebug())184 BBVector += computeEmbeddings(I);185 return BBVector;186}187 188Embedding SymbolicEmbedder::computeEmbeddings(const Instruction &I) const {189 // Currently, we always (re)compute the embeddings for symbolic embedder.190 // This is cheaper than caching the vectors.191 Embedding ArgEmb(Dimension, 0);192 for (const auto &Op : I.operands())193 ArgEmb += Vocab[*Op];194 auto InstVector =195 Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;196 if (const auto *IC = dyn_cast<CmpInst>(&I))197 InstVector += Vocab[IC->getPredicate()];198 return InstVector;199}200 201Embedding FlowAwareEmbedder::computeEmbeddings(const Instruction &I) const {202 // If we have already computed the embedding for this instruction, return it203 auto It = InstVecMap.find(&I);204 if (It != InstVecMap.end())205 return It->second;206 207 // TODO: Handle call instructions differently.208 // For now, we treat them like other instructions209 Embedding ArgEmb(Dimension, 0);210 for (const auto &Op : I.operands()) {211 // If the operand is defined elsewhere, we use its embedding212 if (const auto *DefInst = dyn_cast<Instruction>(Op)) {213 auto DefIt = InstVecMap.find(DefInst);214 // Fixme (#159171): Ideally we should never miss an instruction215 // embedding here.216 // But when we have cyclic dependencies (e.g., phi217 // nodes), we might miss the embedding. In such cases, we fall back to218 // using the vocabulary embedding. This can be fixed by iterating to a219 // fixed-point, or by using a simple solver for the set of simultaneous220 // equations.221 // Another case when we might miss an instruction embedding is when222 // the operand instruction is in a different basic block that has not223 // been processed yet. This can be fixed by processing the basic blocks224 // in a topological order.225 if (DefIt != InstVecMap.end())226 ArgEmb += DefIt->second;227 else228 ArgEmb += Vocab[*Op];229 }230 // If the operand is not defined by an instruction, we use the231 // vocabulary232 else {233 LLVM_DEBUG(errs() << "Using embedding from vocabulary for operand: "234 << *Op << "=" << Vocab[*Op][0] << "\n");235 ArgEmb += Vocab[*Op];236 }237 }238 // Create the instruction vector by combining opcode, type, and arguments239 // embeddings240 auto InstVector =241 Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;242 if (const auto *IC = dyn_cast<CmpInst>(&I))243 InstVector += Vocab[IC->getPredicate()];244 InstVecMap[&I] = InstVector;245 return InstVector;246}247 248// ==----------------------------------------------------------------------===//249// VocabStorage250//===----------------------------------------------------------------------===//251 252VocabStorage::VocabStorage(std::vector<std::vector<Embedding>> &&SectionData)253 : Sections(std::move(SectionData)), TotalSize([&] {254 assert(!Sections.empty() && "Vocabulary has no sections");255 // Compute total size across all sections256 size_t Size = 0;257 for (const auto &Section : Sections) {258 assert(!Section.empty() && "Vocabulary section is empty");259 Size += Section.size();260 }261 return Size;262 }()),263 Dimension([&] {264 // Get dimension from the first embedding in the first section - all265 // embeddings must have the same dimension266 assert(!Sections.empty() && "Vocabulary has no sections");267 assert(!Sections[0].empty() && "First section of vocabulary is empty");268 unsigned ExpectedDim = static_cast<unsigned>(Sections[0][0].size());269 270 // Verify that all embeddings across all sections have the same271 // dimension272 [[maybe_unused]] auto allSameDim =273 [ExpectedDim](const std::vector<Embedding> &Section) {274 return std::all_of(Section.begin(), Section.end(),275 [ExpectedDim](const Embedding &Emb) {276 return Emb.size() == ExpectedDim;277 });278 };279 assert(std::all_of(Sections.begin(), Sections.end(), allSameDim) &&280 "All embeddings must have the same dimension");281 282 return ExpectedDim;283 }()) {}284 285const Embedding &VocabStorage::const_iterator::operator*() const {286 assert(SectionId < Storage->Sections.size() && "Invalid section ID");287 assert(LocalIndex < Storage->Sections[SectionId].size() &&288 "Local index out of range");289 return Storage->Sections[SectionId][LocalIndex];290}291 292VocabStorage::const_iterator &VocabStorage::const_iterator::operator++() {293 ++LocalIndex;294 // Check if we need to move to the next section295 if (SectionId < Storage->getNumSections() &&296 LocalIndex >= Storage->Sections[SectionId].size()) {297 assert(LocalIndex == Storage->Sections[SectionId].size() &&298 "Local index should be at the end of the current section");299 LocalIndex = 0;300 ++SectionId;301 }302 return *this;303}304 305bool VocabStorage::const_iterator::operator==(306 const const_iterator &Other) const {307 return Storage == Other.Storage && SectionId == Other.SectionId &&308 LocalIndex == Other.LocalIndex;309}310 311bool VocabStorage::const_iterator::operator!=(312 const const_iterator &Other) const {313 return !(*this == Other);314}315 316Error VocabStorage::parseVocabSection(StringRef Key,317 const json::Value &ParsedVocabValue,318 VocabMap &TargetVocab, unsigned &Dim) {319 json::Path::Root Path("");320 const json::Object *RootObj = ParsedVocabValue.getAsObject();321 if (!RootObj)322 return createStringError(errc::invalid_argument,323 "JSON root is not an object");324 325 const json::Value *SectionValue = RootObj->get(Key);326 if (!SectionValue)327 return createStringError(errc::invalid_argument,328 "Missing '" + std::string(Key) +329 "' section in vocabulary file");330 if (!json::fromJSON(*SectionValue, TargetVocab, Path))331 return createStringError(errc::illegal_byte_sequence,332 "Unable to parse '" + std::string(Key) +333 "' section from vocabulary");334 335 Dim = TargetVocab.begin()->second.size();336 if (Dim == 0)337 return createStringError(errc::illegal_byte_sequence,338 "Dimension of '" + std::string(Key) +339 "' section of the vocabulary is zero");340 341 if (!std::all_of(TargetVocab.begin(), TargetVocab.end(),342 [Dim](const std::pair<StringRef, Embedding> &Entry) {343 return Entry.second.size() == Dim;344 }))345 return createStringError(346 errc::illegal_byte_sequence,347 "All vectors in the '" + std::string(Key) +348 "' section of the vocabulary are not of the same dimension");349 350 return Error::success();351}352 353// ==----------------------------------------------------------------------===//354// Vocabulary355//===----------------------------------------------------------------------===//356 357StringRef Vocabulary::getVocabKeyForOpcode(unsigned Opcode) {358 assert(Opcode >= 1 && Opcode <= MaxOpcodes && "Invalid opcode");359#define HANDLE_INST(NUM, OPCODE, CLASS) \360 if (Opcode == NUM) { \361 return #OPCODE; \362 }363#include "llvm/IR/Instruction.def"364#undef HANDLE_INST365 return "UnknownOpcode";366}367 368// Helper function to classify an operand into OperandKind369Vocabulary::OperandKind Vocabulary::getOperandKind(const Value *Op) {370 if (isa<Function>(Op))371 return OperandKind::FunctionID;372 if (isa<PointerType>(Op->getType()))373 return OperandKind::PointerID;374 if (isa<Constant>(Op))375 return OperandKind::ConstantID;376 return OperandKind::VariableID;377}378 379unsigned Vocabulary::getPredicateLocalIndex(CmpInst::Predicate P) {380 if (P >= CmpInst::FIRST_FCMP_PREDICATE && P <= CmpInst::LAST_FCMP_PREDICATE)381 return P - CmpInst::FIRST_FCMP_PREDICATE;382 else383 return P - CmpInst::FIRST_ICMP_PREDICATE +384 (CmpInst::LAST_FCMP_PREDICATE - CmpInst::FIRST_FCMP_PREDICATE + 1);385}386 387CmpInst::Predicate Vocabulary::getPredicateFromLocalIndex(unsigned LocalIndex) {388 unsigned fcmpRange =389 CmpInst::LAST_FCMP_PREDICATE - CmpInst::FIRST_FCMP_PREDICATE + 1;390 if (LocalIndex < fcmpRange)391 return static_cast<CmpInst::Predicate>(CmpInst::FIRST_FCMP_PREDICATE +392 LocalIndex);393 else394 return static_cast<CmpInst::Predicate>(CmpInst::FIRST_ICMP_PREDICATE +395 LocalIndex - fcmpRange);396}397 398StringRef Vocabulary::getVocabKeyForPredicate(CmpInst::Predicate Pred) {399 static SmallString<16> PredNameBuffer;400 if (Pred < CmpInst::FIRST_ICMP_PREDICATE)401 PredNameBuffer = "FCMP_";402 else403 PredNameBuffer = "ICMP_";404 PredNameBuffer += CmpInst::getPredicateName(Pred);405 return PredNameBuffer;406}407 408StringRef Vocabulary::getStringKey(unsigned Pos) {409 assert(Pos < NumCanonicalEntries && "Position out of bounds in vocabulary");410 // Opcode411 if (Pos < MaxOpcodes)412 return getVocabKeyForOpcode(Pos + 1);413 // Type414 if (Pos < OperandBaseOffset)415 return getVocabKeyForCanonicalTypeID(416 static_cast<CanonicalTypeID>(Pos - MaxOpcodes));417 // Operand418 if (Pos < PredicateBaseOffset)419 return getVocabKeyForOperandKind(420 static_cast<OperandKind>(Pos - OperandBaseOffset));421 // Predicates422 return getVocabKeyForPredicate(getPredicate(Pos - PredicateBaseOffset));423}424 425// For now, assume vocabulary is stable unless explicitly invalidated.426bool Vocabulary::invalidate(Module &M, const PreservedAnalyses &PA,427 ModuleAnalysisManager::Invalidator &Inv) const {428 auto PAC = PA.getChecker<IR2VecVocabAnalysis>();429 return !(PAC.preservedWhenStateless());430}431 432VocabStorage Vocabulary::createDummyVocabForTest(unsigned Dim) {433 float DummyVal = 0.1f;434 435 // Create sections for opcodes, types, operands, and predicates436 // Order must match Vocabulary::Section enum437 std::vector<std::vector<Embedding>> Sections;438 Sections.reserve(4);439 440 // Opcodes section441 std::vector<Embedding> OpcodeSec;442 OpcodeSec.reserve(MaxOpcodes);443 for (unsigned I = 0; I < MaxOpcodes; ++I) {444 OpcodeSec.emplace_back(Dim, DummyVal);445 DummyVal += 0.1f;446 }447 Sections.push_back(std::move(OpcodeSec));448 449 // Types section450 std::vector<Embedding> TypeSec;451 TypeSec.reserve(MaxCanonicalTypeIDs);452 for (unsigned I = 0; I < MaxCanonicalTypeIDs; ++I) {453 TypeSec.emplace_back(Dim, DummyVal);454 DummyVal += 0.1f;455 }456 Sections.push_back(std::move(TypeSec));457 458 // Operands section459 std::vector<Embedding> OperandSec;460 OperandSec.reserve(MaxOperandKinds);461 for (unsigned I = 0; I < MaxOperandKinds; ++I) {462 OperandSec.emplace_back(Dim, DummyVal);463 DummyVal += 0.1f;464 }465 Sections.push_back(std::move(OperandSec));466 467 // Predicates section468 std::vector<Embedding> PredicateSec;469 PredicateSec.reserve(MaxPredicateKinds);470 for (unsigned I = 0; I < MaxPredicateKinds; ++I) {471 PredicateSec.emplace_back(Dim, DummyVal);472 DummyVal += 0.1f;473 }474 Sections.push_back(std::move(PredicateSec));475 476 return VocabStorage(std::move(Sections));477}478 479// ==----------------------------------------------------------------------===//480// IR2VecVocabAnalysis481//===----------------------------------------------------------------------===//482 483// FIXME: Make this optional. We can avoid file reads484// by auto-generating a default vocabulary during the build time.485Error IR2VecVocabAnalysis::readVocabulary(VocabMap &OpcVocab,486 VocabMap &TypeVocab,487 VocabMap &ArgVocab) {488 auto BufOrError = MemoryBuffer::getFileOrSTDIN(VocabFile, /*IsText=*/true);489 if (!BufOrError)490 return createFileError(VocabFile, BufOrError.getError());491 492 auto Content = BufOrError.get()->getBuffer();493 494 Expected<json::Value> ParsedVocabValue = json::parse(Content);495 if (!ParsedVocabValue)496 return ParsedVocabValue.takeError();497 498 unsigned OpcodeDim = 0, TypeDim = 0, ArgDim = 0;499 if (auto Err = VocabStorage::parseVocabSection("Opcodes", *ParsedVocabValue,500 OpcVocab, OpcodeDim))501 return Err;502 503 if (auto Err = VocabStorage::parseVocabSection("Types", *ParsedVocabValue,504 TypeVocab, TypeDim))505 return Err;506 507 if (auto Err = VocabStorage::parseVocabSection("Arguments", *ParsedVocabValue,508 ArgVocab, ArgDim))509 return Err;510 511 if (!(OpcodeDim == TypeDim && TypeDim == ArgDim))512 return createStringError(errc::illegal_byte_sequence,513 "Vocabulary sections have different dimensions");514 515 return Error::success();516}517 518void IR2VecVocabAnalysis::generateVocabStorage(VocabMap &OpcVocab,519 VocabMap &TypeVocab,520 VocabMap &ArgVocab) {521 522 // Helper for handling missing entities in the vocabulary.523 // Currently, we use a zero vector. In the future, we will throw an error to524 // ensure that *all* known entities are present in the vocabulary.525 auto handleMissingEntity = [](const std::string &Val) {526 LLVM_DEBUG(errs() << Val527 << " is not in vocabulary, using zero vector; This "528 "would result in an error in future.\n");529 ++VocabMissCounter;530 };531 532 unsigned Dim = OpcVocab.begin()->second.size();533 assert(Dim > 0 && "Vocabulary dimension must be greater than zero");534 535 // Handle Opcodes536 std::vector<Embedding> NumericOpcodeEmbeddings(Vocabulary::MaxOpcodes,537 Embedding(Dim));538 for (unsigned Opcode : seq(0u, Vocabulary::MaxOpcodes)) {539 StringRef VocabKey = Vocabulary::getVocabKeyForOpcode(Opcode + 1);540 auto It = OpcVocab.find(VocabKey.str());541 if (It != OpcVocab.end())542 NumericOpcodeEmbeddings[Opcode] = It->second;543 else544 handleMissingEntity(VocabKey.str());545 }546 547 // Handle Types - only canonical types are present in vocabulary548 std::vector<Embedding> NumericTypeEmbeddings(Vocabulary::MaxCanonicalTypeIDs,549 Embedding(Dim));550 for (unsigned CTypeID : seq(0u, Vocabulary::MaxCanonicalTypeIDs)) {551 StringRef VocabKey = Vocabulary::getVocabKeyForCanonicalTypeID(552 static_cast<Vocabulary::CanonicalTypeID>(CTypeID));553 if (auto It = TypeVocab.find(VocabKey.str()); It != TypeVocab.end()) {554 NumericTypeEmbeddings[CTypeID] = It->second;555 continue;556 }557 handleMissingEntity(VocabKey.str());558 }559 560 // Handle Arguments/Operands561 std::vector<Embedding> NumericArgEmbeddings(Vocabulary::MaxOperandKinds,562 Embedding(Dim));563 for (unsigned OpKind : seq(0u, Vocabulary::MaxOperandKinds)) {564 Vocabulary::OperandKind Kind = static_cast<Vocabulary::OperandKind>(OpKind);565 StringRef VocabKey = Vocabulary::getVocabKeyForOperandKind(Kind);566 auto It = ArgVocab.find(VocabKey.str());567 if (It != ArgVocab.end()) {568 NumericArgEmbeddings[OpKind] = It->second;569 continue;570 }571 handleMissingEntity(VocabKey.str());572 }573 574 // Handle Predicates: part of Operands section. We look up predicate keys575 // in ArgVocab.576 std::vector<Embedding> NumericPredEmbeddings(Vocabulary::MaxPredicateKinds,577 Embedding(Dim, 0));578 for (unsigned PK : seq(0u, Vocabulary::MaxPredicateKinds)) {579 StringRef VocabKey =580 Vocabulary::getVocabKeyForPredicate(Vocabulary::getPredicate(PK));581 auto It = ArgVocab.find(VocabKey.str());582 if (It != ArgVocab.end()) {583 NumericPredEmbeddings[PK] = It->second;584 continue;585 }586 handleMissingEntity(VocabKey.str());587 }588 589 // Create section-based storage instead of flat vocabulary590 // Order must match Vocabulary::Section enum591 std::vector<std::vector<Embedding>> Sections(4);592 Sections[static_cast<unsigned>(Vocabulary::Section::Opcodes)] =593 std::move(NumericOpcodeEmbeddings); // Section::Opcodes594 Sections[static_cast<unsigned>(Vocabulary::Section::CanonicalTypes)] =595 std::move(NumericTypeEmbeddings); // Section::CanonicalTypes596 Sections[static_cast<unsigned>(Vocabulary::Section::Operands)] =597 std::move(NumericArgEmbeddings); // Section::Operands598 Sections[static_cast<unsigned>(Vocabulary::Section::Predicates)] =599 std::move(NumericPredEmbeddings); // Section::Predicates600 601 // Create VocabStorage from organized sections602 Vocab.emplace(std::move(Sections));603}604 605void IR2VecVocabAnalysis::emitError(Error Err, LLVMContext &Ctx) {606 handleAllErrors(std::move(Err), [&](const ErrorInfoBase &EI) {607 Ctx.emitError("Error reading vocabulary: " + EI.message());608 });609}610 611IR2VecVocabAnalysis::Result612IR2VecVocabAnalysis::run(Module &M, ModuleAnalysisManager &AM) {613 auto Ctx = &M.getContext();614 // If vocabulary is already populated by the constructor, use it.615 if (Vocab.has_value())616 return Vocabulary(std::move(Vocab.value()));617 618 // Otherwise, try to read from the vocabulary file.619 if (VocabFile.empty()) {620 // FIXME: Use default vocabulary621 Ctx->emitError("IR2Vec vocabulary file path not specified; You may need to "622 "set it using --ir2vec-vocab-path");623 return Vocabulary(); // Return invalid result624 }625 626 VocabMap OpcVocab, TypeVocab, ArgVocab;627 if (auto Err = readVocabulary(OpcVocab, TypeVocab, ArgVocab)) {628 emitError(std::move(Err), *Ctx);629 return Vocabulary();630 }631 632 // Scale the vocabulary sections based on the provided weights633 auto scaleVocabSection = [](VocabMap &Vocab, double Weight) {634 for (auto &Entry : Vocab)635 Entry.second *= Weight;636 };637 scaleVocabSection(OpcVocab, OpcWeight);638 scaleVocabSection(TypeVocab, TypeWeight);639 scaleVocabSection(ArgVocab, ArgWeight);640 641 // Generate the numeric lookup vocabulary642 generateVocabStorage(OpcVocab, TypeVocab, ArgVocab);643 644 return Vocabulary(std::move(Vocab.value()));645}646 647// ==----------------------------------------------------------------------===//648// Printer Passes649//===----------------------------------------------------------------------===//650 651PreservedAnalyses IR2VecPrinterPass::run(Module &M,652 ModuleAnalysisManager &MAM) {653 auto &Vocabulary = MAM.getResult<IR2VecVocabAnalysis>(M);654 assert(Vocabulary.isValid() && "IR2Vec Vocabulary is invalid");655 656 for (Function &F : M) {657 auto Emb = Embedder::create(IR2VecEmbeddingKind, F, Vocabulary);658 if (!Emb) {659 OS << "Error creating IR2Vec embeddings \n";660 continue;661 }662 663 OS << "IR2Vec embeddings for function " << F.getName() << ":\n";664 OS << "Function vector: ";665 Emb->getFunctionVector().print(OS);666 667 OS << "Basic block vectors:\n";668 for (const BasicBlock &BB : F) {669 OS << "Basic block: " << BB.getName() << ":\n";670 Emb->getBBVector(BB).print(OS);671 }672 673 OS << "Instruction vectors:\n";674 for (const BasicBlock &BB : F) {675 for (const Instruction &I : BB) {676 OS << "Instruction: ";677 I.print(OS);678 Emb->getInstVector(I).print(OS);679 }680 }681 }682 return PreservedAnalyses::all();683}684 685PreservedAnalyses IR2VecVocabPrinterPass::run(Module &M,686 ModuleAnalysisManager &MAM) {687 auto &IR2VecVocabulary = MAM.getResult<IR2VecVocabAnalysis>(M);688 assert(IR2VecVocabulary.isValid() && "IR2Vec Vocabulary is invalid");689 690 // Print each entry691 unsigned Pos = 0;692 for (const auto &Entry : IR2VecVocabulary) {693 OS << "Key: " << IR2VecVocabulary.getStringKey(Pos++) << ": ";694 Entry.print(OS);695 }696 return PreservedAnalyses::all();697}698