1//===- IR2Vec.cpp - Implementation of IR2Vec -----------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM
4// Exceptions. See the LICENSE file for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// \file
10/// 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 specified
43cl::opt<std::string>
44 VocabFile("ir2vec-vocab-path", cl::Optional,
45 cl::desc("Path to the vocabulary file for IR2Vec"), cl::init(Val: ""),
46 cl::cat(IR2VecCategory));
47cl::opt<float> OpcWeight("ir2vec-opc-weight", cl::Optional, cl::init(Val: 1.0),
48 cl::desc("Weight for opcode embeddings"),
49 cl::cat(IR2VecCategory));
50cl::opt<float> TypeWeight("ir2vec-type-weight", cl::Optional, cl::init(Val: 0.5),
51 cl::desc("Weight for type embeddings"),
52 cl::cat(IR2VecCategory));
53cl::opt<float> ArgWeight("ir2vec-arg-weight", cl::Optional, cl::init(Val: 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(Val: IR2VecKind::Symbolic), cl::desc("IR2Vec embedding kind"),
63 cl::cat(IR2VecCategory));
64
65} // namespace ir2vec
66} // namespace llvm
67
68AnalysisKey IR2VecVocabAnalysis::Key;
69
70// ==----------------------------------------------------------------------===//
71// Local helper functions
72//===----------------------------------------------------------------------===//
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, Out&: TempOut, P))
78 return false;
79 Out = Embedding(std::move(TempOut));
80 return true;
81}
82} // namespace llvm::json
83
84// ==----------------------------------------------------------------------===//
85// Embedding
86//===----------------------------------------------------------------------===//
87Embedding &Embedding::operator+=(const Embedding &RHS) {
88 assert(this->size() == RHS.size() && "Vectors must have the same dimension");
89 std::transform(first1: this->begin(), last1: this->end(), first2: RHS.begin(), result: this->begin(),
90 binary_op: 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(first1: this->begin(), last1: this->end(), first2: RHS.begin(), result: this->begin(),
103 binary_op: 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(first: this->begin(), last: this->end(), result: this->begin(),
115 unary_op: [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(x: (*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(Fmt: "%.2f", Vals: Elem) << " ";
149 OS << "]\n";
150}
151
152// ==----------------------------------------------------------------------===//
153// Embedder and its subclasses
154//===----------------------------------------------------------------------===//
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>(args: F, args: Vocab);
161 case IR2VecKind::FlowAware:
162 return std::make_unique<FlowAwareEmbedder>(args: F, args: 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 entry
174 for (const BasicBlock *BB : depth_first(G: &F))
175 FuncVector += computeEmbeddings(BB: *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 instructions
183 for (const auto &I : BB)
184 if (!I.isDebugOrPseudoInst())
185 BBVector += computeEmbeddings(I);
186 return BBVector;
187}
188
189Embedding SymbolicEmbedder::computeEmbeddings(const Instruction &I) const {
190 // Currently, we always (re)compute the embeddings for symbolic embedder.
191 // This is cheaper than caching the vectors.
192 Embedding ArgEmb(Dimension, 0);
193 for (const auto &Op : I.operands())
194 ArgEmb += Vocab[*Op];
195 auto InstVector =
196 Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;
197 if (const auto *IC = dyn_cast<CmpInst>(Val: &I))
198 InstVector += Vocab[IC->getPredicate()];
199 return InstVector;
200}
201
202Embedding FlowAwareEmbedder::computeEmbeddings(const Instruction &I) const {
203 // If we have already computed the embedding for this instruction, return it
204 auto It = InstVecMap.find(Val: &I);
205 if (It != InstVecMap.end())
206 return It->second;
207
208 // TODO: Handle call instructions differently.
209 // For now, we treat them like other instructions
210 Embedding ArgEmb(Dimension, 0);
211 for (const auto &Op : I.operands()) {
212 // If the operand is defined elsewhere, we use its embedding
213 if (const auto *DefInst = dyn_cast<Instruction>(Val: Op)) {
214 auto DefIt = InstVecMap.find(Val: DefInst);
215 // Fixme (#159171): Ideally we should never miss an instruction
216 // embedding here.
217 // But when we have cyclic dependencies (e.g., phi
218 // nodes), we might miss the embedding. In such cases, we fall back to
219 // using the vocabulary embedding. This can be fixed by iterating to a
220 // fixed-point, or by using a simple solver for the set of simultaneous
221 // equations.
222 // Another case when we might miss an instruction embedding is when
223 // the operand instruction is in a different basic block that has not
224 // been processed yet. This can be fixed by processing the basic blocks
225 // in a topological order.
226 if (DefIt != InstVecMap.end())
227 ArgEmb += DefIt->second;
228 else
229 ArgEmb += Vocab[*Op];
230 }
231 // If the operand is not defined by an instruction, we use the
232 // vocabulary
233 else {
234 LLVM_DEBUG(errs() << "Using embedding from vocabulary for operand: "
235 << *Op << "=" << Vocab[*Op][0] << "\n");
236 ArgEmb += Vocab[*Op];
237 }
238 }
239 // Create the instruction vector by combining opcode, type, and arguments
240 // embeddings
241 auto InstVector =
242 Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;
243 if (const auto *IC = dyn_cast<CmpInst>(Val: &I))
244 InstVector += Vocab[IC->getPredicate()];
245 InstVecMap[&I] = InstVector;
246 return InstVector;
247}
248
249// ==----------------------------------------------------------------------===//
250// VocabStorage
251//===----------------------------------------------------------------------===//
252
253VocabStorage::VocabStorage(std::vector<std::vector<Embedding>> &&SectionData)
254 : Sections(std::move(SectionData)), TotalSize([&] {
255 assert(!Sections.empty() && "Vocabulary has no sections");
256 // Compute total size across all sections
257 size_t Size = 0;
258 for (const auto &Section : Sections) {
259 assert(!Section.empty() && "Vocabulary section is empty");
260 Size += Section.size();
261 }
262 return Size;
263 }()),
264 Dimension([&] {
265 // Get dimension from the first embedding in the first section - all
266 // embeddings must have the same dimension
267 assert(!Sections.empty() && "Vocabulary has no sections");
268 assert(!Sections[0].empty() && "First section of vocabulary is empty");
269 unsigned ExpectedDim = static_cast<unsigned>(Sections[0][0].size());
270
271 // Verify that all embeddings across all sections have the same
272 // dimension
273 [[maybe_unused]] auto allSameDim =
274 [ExpectedDim](const std::vector<Embedding> &Section) {
275 return std::all_of(first: Section.begin(), last: Section.end(),
276 pred: [ExpectedDim](const Embedding &Emb) {
277 return Emb.size() == ExpectedDim;
278 });
279 };
280 assert(std::all_of(Sections.begin(), Sections.end(), allSameDim) &&
281 "All embeddings must have the same dimension");
282
283 return ExpectedDim;
284 }()) {}
285
286const Embedding &VocabStorage::const_iterator::operator*() const {
287 assert(SectionId < Storage->Sections.size() && "Invalid section ID");
288 assert(LocalIndex < Storage->Sections[SectionId].size() &&
289 "Local index out of range");
290 return Storage->Sections[SectionId][LocalIndex];
291}
292
293VocabStorage::const_iterator &VocabStorage::const_iterator::operator++() {
294 ++LocalIndex;
295 // Check if we need to move to the next section
296 if (SectionId < Storage->getNumSections() &&
297 LocalIndex >= Storage->Sections[SectionId].size()) {
298 assert(LocalIndex == Storage->Sections[SectionId].size() &&
299 "Local index should be at the end of the current section");
300 LocalIndex = 0;
301 ++SectionId;
302 }
303 return *this;
304}
305
306bool VocabStorage::const_iterator::operator==(
307 const const_iterator &Other) const {
308 return Storage == Other.Storage && SectionId == Other.SectionId &&
309 LocalIndex == Other.LocalIndex;
310}
311
312bool VocabStorage::const_iterator::operator!=(
313 const const_iterator &Other) const {
314 return !(*this == Other);
315}
316
317Error VocabStorage::parseVocabSection(StringRef Key,
318 const json::Value &ParsedVocabValue,
319 VocabMap &TargetVocab, unsigned &Dim) {
320 json::Path::Root Path("");
321 const json::Object *RootObj = ParsedVocabValue.getAsObject();
322 if (!RootObj)
323 return createStringError(EC: errc::invalid_argument,
324 S: "JSON root is not an object");
325
326 const json::Value *SectionValue = RootObj->get(K: Key);
327 if (!SectionValue)
328 return createStringError(EC: errc::invalid_argument,
329 S: "Missing '" + std::string(Key) +
330 "' section in vocabulary file");
331 if (!json::fromJSON(E: *SectionValue, Out&: TargetVocab, P: Path))
332 return createStringError(EC: errc::illegal_byte_sequence,
333 S: "Unable to parse '" + std::string(Key) +
334 "' section from vocabulary");
335
336 Dim = TargetVocab.begin()->second.size();
337 if (Dim == 0)
338 return createStringError(EC: errc::illegal_byte_sequence,
339 S: "Dimension of '" + std::string(Key) +
340 "' section of the vocabulary is zero");
341
342 if (!std::all_of(first: TargetVocab.begin(), last: TargetVocab.end(),
343 pred: [Dim](const std::pair<StringRef, Embedding> &Entry) {
344 return Entry.second.size() == Dim;
345 }))
346 return createStringError(
347 EC: errc::illegal_byte_sequence,
348 S: "All vectors in the '" + std::string(Key) +
349 "' section of the vocabulary are not of the same dimension");
350
351 return Error::success();
352}
353
354// ==----------------------------------------------------------------------===//
355// Vocabulary
356//===----------------------------------------------------------------------===//
357
358StringRef Vocabulary::getVocabKeyForOpcode(unsigned Opcode) {
359 assert(Opcode >= 1 && Opcode <= MaxOpcodes && "Invalid opcode");
360#define HANDLE_INST(NUM, OPCODE, CLASS) \
361 if (Opcode == NUM) { \
362 return #OPCODE; \
363 }
364#include "llvm/IR/Instruction.def"
365#undef HANDLE_INST
366 return "UnknownOpcode";
367}
368
369// Helper function to classify an operand into OperandKind
370Vocabulary::OperandKind Vocabulary::getOperandKind(const Value *Op) {
371 if (isa<Function>(Val: Op))
372 return OperandKind::FunctionID;
373 if (isa<PointerType>(Val: Op->getType()))
374 return OperandKind::PointerID;
375 if (isa<Constant>(Val: Op))
376 return OperandKind::ConstantID;
377 return OperandKind::VariableID;
378}
379
380unsigned Vocabulary::getPredicateLocalIndex(CmpInst::Predicate P) {
381 if (P >= CmpInst::FIRST_FCMP_PREDICATE && P <= CmpInst::LAST_FCMP_PREDICATE)
382 return P - CmpInst::FIRST_FCMP_PREDICATE;
383 else
384 return P - CmpInst::FIRST_ICMP_PREDICATE +
385 (CmpInst::LAST_FCMP_PREDICATE - CmpInst::FIRST_FCMP_PREDICATE + 1);
386}
387
388CmpInst::Predicate Vocabulary::getPredicateFromLocalIndex(unsigned LocalIndex) {
389 unsigned fcmpRange =
390 CmpInst::LAST_FCMP_PREDICATE - CmpInst::FIRST_FCMP_PREDICATE + 1;
391 if (LocalIndex < fcmpRange)
392 return static_cast<CmpInst::Predicate>(CmpInst::FIRST_FCMP_PREDICATE +
393 LocalIndex);
394 else
395 return static_cast<CmpInst::Predicate>(CmpInst::FIRST_ICMP_PREDICATE +
396 LocalIndex - fcmpRange);
397}
398
399StringRef Vocabulary::getVocabKeyForPredicate(CmpInst::Predicate Pred) {
400 static SmallString<16> PredNameBuffer;
401 if (Pred < CmpInst::FIRST_ICMP_PREDICATE)
402 PredNameBuffer = "FCMP_";
403 else
404 PredNameBuffer = "ICMP_";
405 PredNameBuffer += CmpInst::getPredicateName(P: Pred);
406 return PredNameBuffer;
407}
408
409StringRef Vocabulary::getStringKey(unsigned Pos) {
410 assert(Pos < NumCanonicalEntries && "Position out of bounds in vocabulary");
411 // Opcode
412 if (Pos < MaxOpcodes)
413 return getVocabKeyForOpcode(Opcode: Pos + 1);
414 // Type
415 if (Pos < OperandBaseOffset)
416 return getVocabKeyForCanonicalTypeID(
417 CType: static_cast<CanonicalTypeID>(Pos - MaxOpcodes));
418 // Operand
419 if (Pos < PredicateBaseOffset)
420 return getVocabKeyForOperandKind(
421 Kind: static_cast<OperandKind>(Pos - OperandBaseOffset));
422 // Predicates
423 return getVocabKeyForPredicate(Pred: getPredicate(Index: Pos - PredicateBaseOffset));
424}
425
426// For now, assume vocabulary is stable unless explicitly invalidated.
427bool Vocabulary::invalidate(Module &M, const PreservedAnalyses &PA,
428 ModuleAnalysisManager::Invalidator &Inv) const {
429 auto PAC = PA.getChecker<IR2VecVocabAnalysis>();
430 return !(PAC.preservedWhenStateless());
431}
432
433VocabStorage Vocabulary::createDummyVocabForTest(unsigned Dim) {
434 float DummyVal = 0.1f;
435
436 // Create sections for opcodes, types, operands, and predicates
437 // Order must match Vocabulary::Section enum
438 std::vector<std::vector<Embedding>> Sections;
439 Sections.reserve(n: 4);
440
441 // Opcodes section
442 std::vector<Embedding> OpcodeSec;
443 OpcodeSec.reserve(n: MaxOpcodes);
444 for (unsigned I = 0; I < MaxOpcodes; ++I) {
445 OpcodeSec.emplace_back(args&: Dim, args&: DummyVal);
446 DummyVal += 0.1f;
447 }
448 Sections.push_back(x: std::move(OpcodeSec));
449
450 // Types section
451 std::vector<Embedding> TypeSec;
452 TypeSec.reserve(n: MaxCanonicalTypeIDs);
453 for (unsigned I = 0; I < MaxCanonicalTypeIDs; ++I) {
454 TypeSec.emplace_back(args&: Dim, args&: DummyVal);
455 DummyVal += 0.1f;
456 }
457 Sections.push_back(x: std::move(TypeSec));
458
459 // Operands section
460 std::vector<Embedding> OperandSec;
461 OperandSec.reserve(n: MaxOperandKinds);
462 for (unsigned I = 0; I < MaxOperandKinds; ++I) {
463 OperandSec.emplace_back(args&: Dim, args&: DummyVal);
464 DummyVal += 0.1f;
465 }
466 Sections.push_back(x: std::move(OperandSec));
467
468 // Predicates section
469 std::vector<Embedding> PredicateSec;
470 PredicateSec.reserve(n: MaxPredicateKinds);
471 for (unsigned I = 0; I < MaxPredicateKinds; ++I) {
472 PredicateSec.emplace_back(args&: Dim, args&: DummyVal);
473 DummyVal += 0.1f;
474 }
475 Sections.push_back(x: std::move(PredicateSec));
476
477 return VocabStorage(std::move(Sections));
478}
479
480namespace {
481using VocabMap = std::map<std::string, Embedding>;
482
483/// Read vocabulary JSON file and populate the section maps.
484Error readVocabularyFromFile(StringRef VocabFilePath, VocabMap &OpcVocab,
485 VocabMap &TypeVocab, VocabMap &ArgVocab) {
486 auto BufOrError =
487 MemoryBuffer::getFileOrSTDIN(Filename: VocabFilePath, /*IsText=*/true);
488 if (!BufOrError)
489 return createFileError(F: VocabFilePath, EC: BufOrError.getError());
490
491 auto Content = BufOrError.get()->getBuffer();
492
493 Expected<json::Value> ParsedVocabValue = json::parse(JSON: Content);
494 if (!ParsedVocabValue)
495 return ParsedVocabValue.takeError();
496
497 unsigned OpcodeDim = 0, TypeDim = 0, ArgDim = 0;
498 if (auto Err = VocabStorage::parseVocabSection(Key: "Opcodes", ParsedVocabValue: *ParsedVocabValue,
499 TargetVocab&: OpcVocab, Dim&: OpcodeDim))
500 return Err;
501
502 if (auto Err = VocabStorage::parseVocabSection(Key: "Types", ParsedVocabValue: *ParsedVocabValue,
503 TargetVocab&: TypeVocab, Dim&: TypeDim))
504 return Err;
505
506 if (auto Err = VocabStorage::parseVocabSection(Key: "Arguments", ParsedVocabValue: *ParsedVocabValue,
507 TargetVocab&: ArgVocab, Dim&: ArgDim))
508 return Err;
509
510 if (!(OpcodeDim == TypeDim && TypeDim == ArgDim))
511 return createStringError(EC: errc::illegal_byte_sequence,
512 S: "Vocabulary sections have different dimensions");
513
514 return Error::success();
515}
516} // anonymous namespace
517
518/// Generate VocabStorage from vocabulary maps.
519VocabStorage Vocabulary::buildVocabStorage(const VocabMap &OpcVocab,
520 const VocabMap &TypeVocab,
521 const VocabMap &ArgVocab) {
522
523 // Helper for handling missing entities in the vocabulary.
524 // Currently, we use a zero vector. In the future, we will throw an error to
525 // ensure that *all* known entities are present in the vocabulary.
526 auto handleMissingEntity = [](const std::string &Val) {
527 LLVM_DEBUG(errs() << Val
528 << " is not in vocabulary, using zero vector; This "
529 "would result in an error in future.\n");
530 ++VocabMissCounter;
531 };
532
533 unsigned Dim = OpcVocab.begin()->second.size();
534 assert(Dim > 0 && "Vocabulary dimension must be greater than zero");
535
536 // Handle Opcodes
537 std::vector<Embedding> NumericOpcodeEmbeddings(Vocabulary::MaxOpcodes,
538 Embedding(Dim));
539 for (unsigned Opcode : seq(Begin: 0u, End: Vocabulary::MaxOpcodes)) {
540 StringRef VocabKey = Vocabulary::getVocabKeyForOpcode(Opcode: Opcode + 1);
541 auto It = OpcVocab.find(x: VocabKey.str());
542 if (It != OpcVocab.end())
543 NumericOpcodeEmbeddings[Opcode] = It->second;
544 else
545 handleMissingEntity(VocabKey.str());
546 }
547
548 // Handle Types - only canonical types are present in vocabulary
549 std::vector<Embedding> NumericTypeEmbeddings(Vocabulary::MaxCanonicalTypeIDs,
550 Embedding(Dim));
551 for (unsigned CTypeID : seq(Begin: 0u, End: Vocabulary::MaxCanonicalTypeIDs)) {
552 StringRef VocabKey = Vocabulary::getVocabKeyForCanonicalTypeID(
553 CType: static_cast<Vocabulary::CanonicalTypeID>(CTypeID));
554 if (auto It = TypeVocab.find(x: VocabKey.str()); It != TypeVocab.end()) {
555 NumericTypeEmbeddings[CTypeID] = It->second;
556 continue;
557 }
558 handleMissingEntity(VocabKey.str());
559 }
560
561 // Handle Arguments/Operands
562 std::vector<Embedding> NumericArgEmbeddings(Vocabulary::MaxOperandKinds,
563 Embedding(Dim));
564 for (unsigned OpKind : seq(Begin: 0u, End: Vocabulary::MaxOperandKinds)) {
565 Vocabulary::OperandKind Kind = static_cast<Vocabulary::OperandKind>(OpKind);
566 StringRef VocabKey = Vocabulary::getVocabKeyForOperandKind(Kind);
567 auto It = ArgVocab.find(x: VocabKey.str());
568 if (It != ArgVocab.end()) {
569 NumericArgEmbeddings[OpKind] = It->second;
570 continue;
571 }
572 handleMissingEntity(VocabKey.str());
573 }
574
575 // Handle Predicates: part of Operands section. We look up predicate keys
576 // in ArgVocab.
577 std::vector<Embedding> NumericPredEmbeddings(Vocabulary::MaxPredicateKinds,
578 Embedding(Dim, 0));
579 for (unsigned PK : seq(Begin: 0u, End: Vocabulary::MaxPredicateKinds)) {
580 StringRef VocabKey =
581 Vocabulary::getVocabKeyForPredicate(Pred: Vocabulary::getPredicate(Index: PK));
582 auto It = ArgVocab.find(x: VocabKey.str());
583 if (It != ArgVocab.end()) {
584 NumericPredEmbeddings[PK] = It->second;
585 continue;
586 }
587 handleMissingEntity(VocabKey.str());
588 }
589
590 // Create section-based storage instead of flat vocabulary
591 // Order must match Vocabulary::Section enum
592 std::vector<std::vector<Embedding>> Sections(4);
593 Sections[static_cast<unsigned>(Section::Opcodes)] =
594 std::move(NumericOpcodeEmbeddings); // Section::Opcodes
595 Sections[static_cast<unsigned>(Section::CanonicalTypes)] =
596 std::move(NumericTypeEmbeddings); // Section::CanonicalTypes
597 Sections[static_cast<unsigned>(Section::Operands)] =
598 std::move(NumericArgEmbeddings); // Section::Operands
599 Sections[static_cast<unsigned>(Section::Predicates)] =
600 std::move(NumericPredEmbeddings); // Section::Predicates
601
602 // Create VocabStorage from organized sections
603 return VocabStorage(std::move(Sections));
604}
605
606// ==----------------------------------------------------------------------===//
607// Vocabulary
608//===----------------------------------------------------------------------===//
609
610Expected<Vocabulary> Vocabulary::fromFile(StringRef VocabFilePath,
611 float OpcWeight, float TypeWeight,
612 float ArgWeight) {
613 VocabMap OpcVocab, TypeVocab, ArgVocab;
614 if (auto Err =
615 readVocabularyFromFile(VocabFilePath, OpcVocab, TypeVocab, ArgVocab))
616 return std::move(Err);
617
618 // Scale the vocabulary sections based on the provided weights
619 auto scaleVocabSection = [](VocabMap &Vocab, float Weight) {
620 for (auto &Entry : Vocab)
621 Entry.second *= Weight;
622 };
623 scaleVocabSection(OpcVocab, OpcWeight);
624 scaleVocabSection(TypeVocab, TypeWeight);
625 scaleVocabSection(ArgVocab, ArgWeight);
626
627 // Generate the numeric lookup vocabulary
628 return Vocabulary(buildVocabStorage(OpcVocab, TypeVocab, ArgVocab));
629}
630
631// ==----------------------------------------------------------------------===//
632// IR2VecVocabAnalysis
633//===----------------------------------------------------------------------===//
634
635void IR2VecVocabAnalysis::emitError(Error Err, LLVMContext &Ctx) {
636 handleAllErrors(E: std::move(Err), Handlers: [&](const ErrorInfoBase &EI) {
637 Ctx.emitError(ErrorStr: "Error reading vocabulary: " + EI.message());
638 });
639}
640
641IR2VecVocabAnalysis::Result
642IR2VecVocabAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
643 auto Ctx = &M.getContext();
644 // If vocabulary is already populated by the constructor, use it.
645 if (Vocab.has_value())
646 return Vocabulary(std::move(Vocab.value()));
647
648 // Otherwise, try to read from the vocabulary file specified via CLI.
649 if (VocabFile.empty()) {
650 // FIXME: Use default vocabulary
651 Ctx->emitError(ErrorStr: "IR2Vec vocabulary file path not specified; You may need to "
652 "set it using --ir2vec-vocab-path");
653 return Vocabulary(); // Return invalid result
654 }
655
656 // Use the static factory method to load the vocabulary.
657 auto VocabOrErr =
658 Vocabulary::fromFile(VocabFilePath: VocabFile, OpcWeight, TypeWeight, ArgWeight);
659 if (!VocabOrErr) {
660 emitError(Err: VocabOrErr.takeError(), Ctx&: *Ctx);
661 return Vocabulary();
662 }
663
664 return std::move(*VocabOrErr);
665}
666
667// ==----------------------------------------------------------------------===//
668// Printer Passes
669//===----------------------------------------------------------------------===//
670
671PreservedAnalyses IR2VecPrinterPass::run(Module &M,
672 ModuleAnalysisManager &MAM) {
673 auto &Vocabulary = MAM.getResult<IR2VecVocabAnalysis>(IR&: M);
674 assert(Vocabulary.isValid() && "IR2Vec Vocabulary is invalid");
675
676 for (Function &F : M) {
677 auto Emb = Embedder::create(Mode: IR2VecEmbeddingKind, F, Vocab: Vocabulary);
678 if (!Emb) {
679 OS << "Error creating IR2Vec embeddings \n";
680 continue;
681 }
682
683 OS << "IR2Vec embeddings for function " << F.getName() << ":\n";
684 OS << "Function vector: ";
685 Emb->getFunctionVector().print(OS);
686
687 OS << "Basic block vectors:\n";
688 for (const BasicBlock &BB : F) {
689 OS << "Basic block: " << BB.getName() << ":\n";
690 Emb->getBBVector(BB).print(OS);
691 }
692
693 OS << "Instruction vectors:\n";
694 for (const BasicBlock &BB : F) {
695 for (const Instruction &I : BB) {
696 OS << "Instruction: ";
697 I.print(O&: OS);
698 Emb->getInstVector(I).print(OS);
699 }
700 }
701 }
702 return PreservedAnalyses::all();
703}
704
705PreservedAnalyses IR2VecVocabPrinterPass::run(Module &M,
706 ModuleAnalysisManager &MAM) {
707 auto &IR2VecVocabulary = MAM.getResult<IR2VecVocabAnalysis>(IR&: M);
708 assert(IR2VecVocabulary.isValid() && "IR2Vec Vocabulary is invalid");
709
710 // Print each entry
711 unsigned Pos = 0;
712 for (const auto &Entry : IR2VecVocabulary) {
713 OS << "Key: " << IR2VecVocabulary.getStringKey(Pos: Pos++) << ": ";
714 Entry.print(OS);
715 }
716 return PreservedAnalyses::all();
717}
718