1//===- IRUtils.cpp - IR2Vec Embedding Generation ----------------------===//
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-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// \file
10/// This file implements the IR2VecTool class for IR2Vec embedding generation
11/// from LLVM IR. It has no dependency on Machine IR.
12///
13//===----------------------------------------------------------------------===//
14
15#include "IRUtils.h"
16#include "llvm/Analysis/IR2Vec.h"
17#include "llvm/IR/BasicBlock.h"
18#include "llvm/IR/Function.h"
19#include "llvm/IR/InstIterator.h"
20#include "llvm/IR/Instructions.h"
21#include "llvm/IR/IntrinsicInst.h"
22#include "llvm/IR/LLVMContext.h"
23#include "llvm/IR/Module.h"
24#include "llvm/IR/PassInstrumentation.h"
25#include "llvm/IR/PassManager.h"
26#include "llvm/IR/Type.h"
27#include "llvm/Support/Debug.h"
28#include "llvm/Support/Errc.h"
29#include "llvm/Support/Error.h"
30#include "llvm/Support/WithColor.h"
31#include "llvm/Support/raw_ostream.h"
32
33#define DEBUG_TYPE "ir2vec"
34
35namespace llvm {
36namespace ir2vec {
37
38Expected<std::shared_ptr<Vocabulary>> loadVocabulary(StringRef VocabPath) {
39 auto VocabOrErr = Vocabulary::fromFile(VocabFilePath: VocabPath);
40 if (!VocabOrErr)
41 return VocabOrErr.takeError();
42
43 auto V = std::make_shared<Vocabulary>(args: std::move(*VocabOrErr));
44
45 if (!V->isValid())
46 return createStringError(EC: errc::invalid_argument,
47 S: "Failed to initialize IR2Vec vocabulary");
48 return V;
49}
50
51Error IR2VecTool::setVocabulary(std::shared_ptr<Vocabulary> V) {
52 if (!V)
53 return createStringError(EC: errc::invalid_argument,
54 S: "Null pointer provided for vocabulary. Will not "
55 "set IR2VecTool vocabulary.");
56 if (!V->isValid())
57 return createStringError(
58 EC: errc::invalid_argument,
59 S: "Vocabulary is not valid. Will not set IR2VecTool vocabulary.");
60 Vocab = std::move(V);
61 return Error::success();
62}
63
64TripletResult IR2VecTool::generateTriplets(const Function &F) const {
65 if (F.isDeclaration())
66 return {};
67
68 TripletResult Result;
69 Result.MaxRelation = 0;
70
71 unsigned MaxRelation = NextRelation;
72 unsigned PrevOpcode = 0;
73 bool HasPrevOpcode = false;
74
75 for (const BasicBlock &BB : F) {
76 for (const auto &I : BB) {
77 if (isa<PseudoProbeInst>(Val: I))
78 continue;
79 unsigned Opcode = Vocabulary::getIndex(Opcode: I.getOpcode());
80 unsigned TypeID = Vocabulary::getIndex(TypeID: I.getType()->getTypeID());
81
82 // Add "Next" relationship with previous instruction
83 if (HasPrevOpcode) {
84 Result.Triplets.push_back(x: {.Head: PrevOpcode, .Tail: Opcode, .Relation: NextRelation});
85 LLVM_DEBUG(dbgs() << Vocabulary::getVocabKeyForOpcode(PrevOpcode + 1)
86 << '\t'
87 << Vocabulary::getVocabKeyForOpcode(Opcode + 1)
88 << '\t' << "Next\n");
89 }
90
91 // Add "Type" relationship
92 Result.Triplets.push_back(x: {.Head: Opcode, .Tail: TypeID, .Relation: TypeRelation});
93 LLVM_DEBUG(
94 dbgs() << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'
95 << Vocabulary::getVocabKeyForTypeID(I.getType()->getTypeID())
96 << '\t' << "Type\n");
97
98 // Add "Arg" relationships
99 unsigned ArgIndex = 0;
100 for (const Use &U : I.operands()) {
101 unsigned OperandID = Vocabulary::getIndex(Op: *U.get());
102 unsigned RelationID = ArgRelation + ArgIndex;
103 Result.Triplets.push_back(x: {.Head: Opcode, .Tail: OperandID, .Relation: RelationID});
104
105 LLVM_DEBUG({
106 StringRef OperandStr = Vocabulary::getVocabKeyForOperandKind(
107 Vocabulary::getOperandKind(U.get()));
108 dbgs() << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'
109 << OperandStr << '\t' << "Arg" << ArgIndex << '\n';
110 });
111
112 ++ArgIndex;
113 }
114 // Only update MaxRelation if there were operands
115 if (ArgIndex > 0)
116 MaxRelation = std::max(a: MaxRelation, b: ArgRelation + ArgIndex - 1);
117 PrevOpcode = Opcode;
118 HasPrevOpcode = true;
119 }
120 }
121
122 Result.MaxRelation = MaxRelation;
123 return Result;
124}
125
126TripletResult IR2VecTool::generateTriplets() const {
127 TripletResult Result;
128 Result.MaxRelation = NextRelation;
129
130 for (const Function &F : M.getFunctionDefs()) {
131 TripletResult FuncResult = generateTriplets(F);
132 Result.MaxRelation = std::max(a: Result.MaxRelation, b: FuncResult.MaxRelation);
133 Result.Triplets.insert(position: Result.Triplets.end(), first: FuncResult.Triplets.begin(),
134 last: FuncResult.Triplets.end());
135 }
136
137 return Result;
138}
139
140void IR2VecTool::writeTripletsToStream(raw_ostream &OS) const {
141 auto Result = generateTriplets();
142 OS << "MAX_RELATION=" << Result.MaxRelation << '\n';
143 for (const auto &T : Result.Triplets)
144 OS << T.Head << '\t' << T.Tail << '\t' << T.Relation << '\n';
145}
146
147EntityList IR2VecTool::collectEntityMappings() {
148 auto EntityLen = Vocabulary::getCanonicalSize();
149 EntityList Result;
150 for (unsigned EntityID = 0; EntityID < EntityLen; ++EntityID)
151 Result.push_back(x: Vocabulary::getStringKey(Pos: EntityID).str());
152 return Result;
153}
154
155void IR2VecTool::writeEntitiesToStream(raw_ostream &OS) {
156 auto Entities = collectEntityMappings();
157 OS << Entities.size() << "\n";
158 for (unsigned EntityID = 0; EntityID < Entities.size(); ++EntityID)
159 OS << Entities[EntityID] << '\t' << EntityID << '\n';
160}
161
162Expected<std::unique_ptr<Embedder>>
163IR2VecTool::createIR2VecEmbedder(const Function &F, IR2VecKind Kind) const {
164 if (!Vocab || !Vocab->isValid())
165 return createStringError(
166 EC: errc::invalid_argument,
167 S: "Vocabulary is not valid. IR2VecTool not initialized.");
168
169 if (F.isDeclaration())
170 return createStringError(EC: errc::invalid_argument,
171 S: "Function is a declaration.");
172
173 auto Emb = Embedder::create(Mode: Kind, F, Vocab: *Vocab);
174 if (!Emb)
175 return createStringError(EC: errc::invalid_argument,
176 Fmt: "Failed to create embedder for function '%s'.",
177 Vals: F.getName().str().c_str());
178
179 return std::move(Emb);
180}
181
182Expected<Embedding> IR2VecTool::getFunctionEmbedding(const Function &F,
183 IR2VecKind Kind) const {
184 auto Emb = createIR2VecEmbedder(F, Kind);
185 if (!Emb)
186 return Emb.takeError();
187
188 return (*Emb)->getFunctionVector();
189}
190
191Expected<FuncEmbMap>
192IR2VecTool::getFunctionEmbeddingsMap(IR2VecKind Kind) const {
193 FuncEmbMap Result;
194
195 for (const Function &F : M.getFunctionDefs()) {
196 auto Emb = getFunctionEmbedding(F, Kind);
197 if (!Emb)
198 return Emb.takeError();
199 Result.try_emplace(Key: &F, Args: std::move(*Emb));
200 }
201
202 return Result;
203}
204
205Expected<BBEmbeddingsMap>
206IR2VecTool::getBBEmbeddingsMap(const Function &F, IR2VecKind Kind) const {
207 auto Emb = createIR2VecEmbedder(F, Kind);
208 if (!Emb)
209 return Emb.takeError();
210
211 BBEmbeddingsMap Result;
212
213 for (const BasicBlock &BB : F)
214 Result.try_emplace(Key: &BB, Args: (*Emb)->getBBVector(BB));
215
216 return Result;
217}
218
219Expected<InstEmbeddingsMap>
220IR2VecTool::getInstEmbeddingsMap(const Function &F, IR2VecKind Kind) const {
221 auto Emb = createIR2VecEmbedder(F, Kind);
222 if (!Emb)
223 return Emb.takeError();
224
225 InstEmbeddingsMap Result;
226
227 for (const Instruction &I : instructions(F))
228 Result.try_emplace(Key: &I, Args: (*Emb)->getInstVector(I));
229
230 return Result;
231}
232
233void IR2VecTool::writeEmbeddingsToStream(raw_ostream &OS,
234 EmbeddingLevel Level) const {
235 for (const Function &F : M.getFunctionDefs())
236 writeEmbeddingsToStream(F, OS, Level);
237}
238
239void IR2VecTool::writeEmbeddingsToStream(const Function &F, raw_ostream &OS,
240 EmbeddingLevel Level) const {
241 auto IR2VecEmbedderObj = createIR2VecEmbedder(F, Kind: IR2VecEmbeddingKind);
242 if (!IR2VecEmbedderObj) {
243 WithColor::error(OS&: errs(), Prefix: ToolName)
244 << toString(E: IR2VecEmbedderObj.takeError()) << "\n";
245 return;
246 }
247 auto Emb = std::move(*IR2VecEmbedderObj);
248
249 OS << "Function: " << F.getName() << "\n";
250
251 // Generate embeddings based on the specified level
252 switch (Level) {
253 case FunctionLevel:
254 Emb->getFunctionVector().print(OS);
255 break;
256 case BasicBlockLevel:
257 for (const BasicBlock &BB : F) {
258 OS << BB.getName() << ":";
259 Emb->getBBVector(BB).print(OS);
260 }
261 break;
262 case InstructionLevel:
263 for (const Instruction &I : instructions(F)) {
264 OS << I;
265 Emb->getInstVector(I).print(OS);
266 }
267 break;
268 }
269}
270
271} // namespace ir2vec
272} // namespace llvm
273