1//===- MIR2Vec.cpp - Implementation of MIR2Vec ---------------------------===//
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 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");
32STATISTIC(MIRClasslessRegCounter,
33 "Number of register operands with no register class");
34
35namespace llvm {
36namespace mir2vec {
37cl::OptionCategory MIR2VecCategory("MIR2Vec Options");
38
39// FIXME: Use a default vocab when not specified
40static cl::opt<std::string>
41 VocabFile("mir2vec-vocab-path", cl::Optional,
42 cl::desc("Path to the vocabulary file for MIR2Vec"), cl::init(Val: ""),
43 cl::cat(MIR2VecCategory));
44cl::opt<float> OpcWeight("mir2vec-opc-weight", cl::Optional, cl::init(Val: 1.0),
45 cl::desc("Weight for machine opcode embeddings"),
46 cl::cat(MIR2VecCategory));
47cl::opt<float> CommonOperandWeight(
48 "mir2vec-common-operand-weight", cl::Optional, cl::init(Val: 1.0),
49 cl::desc("Weight for common operand embeddings"), cl::cat(MIR2VecCategory));
50cl::opt<float>
51 RegOperandWeight("mir2vec-reg-operand-weight", cl::Optional, cl::init(Val: 1.0),
52 cl::desc("Weight for register operand embeddings"),
53 cl::cat(MIR2VecCategory));
54cl::opt<MIR2VecKind> MIR2VecEmbeddingKind(
55 "mir2vec-kind", cl::Optional,
56 cl::values(clEnumValN(MIR2VecKind::Symbolic, "symbolic",
57 "Generate symbolic embeddings for MIR")),
58 cl::init(Val: MIR2VecKind::Symbolic), cl::desc("MIR2Vec embedding kind"),
59 cl::cat(MIR2VecCategory));
60
61static cl::opt<bool> PrintAllVocabEntries(
62 "mir2vec-print-all-vocab-entries", cl::Optional, cl::init(Val: false),
63 cl::desc("Print all vocabulary entries including zero embeddings"),
64 cl::cat(MIR2VecCategory));
65
66} // namespace mir2vec
67} // namespace llvm
68
69//===----------------------------------------------------------------------===//
70// Vocabulary
71//===----------------------------------------------------------------------===//
72
73MIRVocabulary::MIRVocabulary(VocabMap &&OpcodeMap, VocabMap &&CommonOperandMap,
74 VocabMap &&PhysicalRegisterMap,
75 VocabMap &&VirtualRegisterMap,
76 const TargetInstrInfo &TII,
77 const TargetRegisterInfo &TRI,
78 const MachineRegisterInfo &MRI)
79 : TII(TII), TRI(TRI), MRI(MRI) {
80 buildCanonicalOpcodeMapping();
81 unsigned CanonicalOpcodeCount = UniqueBaseOpcodeNames.size();
82 assert(CanonicalOpcodeCount > 0 &&
83 "No canonical opcodes found for target - invalid vocabulary");
84
85 buildRegisterOperandMapping();
86
87 // Define layout of vocabulary sections
88 Layout.OpcodeBase = 0;
89 Layout.CommonOperandBase = CanonicalOpcodeCount;
90 // We expect same classes for physical and virtual registers
91 Layout.PhyRegBase = Layout.CommonOperandBase + std::size(CommonOperandNames);
92 Layout.VirtRegBase = Layout.PhyRegBase + RegisterOperandNames.size();
93
94 generateStorage(OpcodeMap, CommonOperandMap, PhyRegMap: PhysicalRegisterMap,
95 VirtRegMap: VirtualRegisterMap);
96 Layout.TotalEntries = Storage.size();
97}
98
99Expected<MIRVocabulary>
100MIRVocabulary::create(VocabMap &&OpcodeMap, VocabMap &&CommonOperandMap,
101 VocabMap &&PhyRegMap, VocabMap &&VirtRegMap,
102 const TargetInstrInfo &TII, const TargetRegisterInfo &TRI,
103 const MachineRegisterInfo &MRI) {
104 if (OpcodeMap.empty() || CommonOperandMap.empty() || PhyRegMap.empty() ||
105 VirtRegMap.empty())
106 return createStringError(EC: errc::invalid_argument,
107 S: "Empty vocabulary entries provided");
108
109 MIRVocabulary Vocab(std::move(OpcodeMap), std::move(CommonOperandMap),
110 std::move(PhyRegMap), std::move(VirtRegMap), TII, TRI,
111 MRI);
112
113 // Validate Storage after construction
114 if (!Vocab.Storage.isValid())
115 return createStringError(EC: errc::invalid_argument,
116 S: "Failed to create valid vocabulary storage");
117 Vocab.ZeroEmbedding = Embedding(Vocab.Storage.getDimension(), 0.0);
118 return std::move(Vocab);
119}
120
121std::string MIRVocabulary::extractBaseOpcodeName(StringRef InstrName) {
122 // Extract base instruction name using regex to capture letters and
123 // underscores Examples: "ADD32rr" -> "ADD", "ARITH_FENCE" -> "ARITH_FENCE"
124 //
125 // TODO: Consider more sophisticated extraction:
126 // - Handle complex prefixes like "AVX1_SETALLONES" correctly (Currently, it
127 // would naively map to "AVX")
128 // - Extract width suffixes (8,16,32,64) as separate features
129 // - Capture addressing mode suffixes (r,i,m,ri,etc.) for better analysis
130 // (Currently, instances like "MOV32mi" map to "MOV", but "ADDPDrr" would map
131 // to "ADDPDrr")
132
133 assert(!InstrName.empty() && "Instruction name should not be empty");
134
135 // Use regex to extract initial sequence of letters and underscores
136 static const Regex BaseOpcodeRegex("([a-zA-Z_]+)");
137 SmallVector<StringRef, 2> Matches;
138
139 if (BaseOpcodeRegex.match(String: InstrName, Matches: &Matches) && Matches.size() > 1) {
140 StringRef Match = Matches[1];
141 // Trim trailing underscores
142 while (!Match.empty() && Match.back() == '_')
143 Match = Match.drop_back();
144 return Match.str();
145 }
146
147 // Fallback to original name if no pattern matches
148 return InstrName.str();
149}
150
151unsigned MIRVocabulary::getCanonicalIndexForBaseName(StringRef BaseName) const {
152 assert(!UniqueBaseOpcodeNames.empty() && "Canonical mapping not built");
153 auto It = std::find(first: UniqueBaseOpcodeNames.begin(),
154 last: UniqueBaseOpcodeNames.end(), val: BaseName.str());
155 assert(It != UniqueBaseOpcodeNames.end() &&
156 "Base name not found in unique opcodes");
157 return std::distance(first: UniqueBaseOpcodeNames.begin(), last: It);
158}
159
160unsigned MIRVocabulary::getCanonicalOpcodeIndex(unsigned Opcode) const {
161 auto BaseOpcode = extractBaseOpcodeName(InstrName: TII.getName(Opcode));
162 return getCanonicalIndexForBaseName(BaseName: BaseOpcode);
163}
164
165unsigned
166MIRVocabulary::getCanonicalIndexForOperandName(StringRef OperandName) const {
167 auto It = std::find(first: std::begin(arr: CommonOperandNames),
168 last: std::end(arr: CommonOperandNames), val: OperandName);
169 assert(It != std::end(CommonOperandNames) &&
170 "Operand name not found in common operands");
171 return Layout.CommonOperandBase +
172 std::distance(first: std::begin(arr: CommonOperandNames), last: It);
173}
174
175unsigned
176MIRVocabulary::getCanonicalIndexForRegisterClass(StringRef RegName,
177 bool IsPhysical) const {
178 auto It = std::find(first: RegisterOperandNames.begin(), last: RegisterOperandNames.end(),
179 val: RegName);
180 assert(It != RegisterOperandNames.end() &&
181 "Register name not found in register operands");
182 unsigned LocalIndex = std::distance(first: RegisterOperandNames.begin(), last: It);
183 return (IsPhysical ? Layout.PhyRegBase : Layout.VirtRegBase) + LocalIndex;
184}
185
186std::string MIRVocabulary::getStringKey(unsigned Pos) const {
187 assert(Pos < Layout.TotalEntries && "Position out of bounds in vocabulary");
188
189 // Handle opcodes section
190 if (Pos < Layout.CommonOperandBase) {
191 // Convert canonical index back to base opcode name
192 auto It = UniqueBaseOpcodeNames.begin();
193 std::advance(i&: It, n: Pos);
194 assert(It != UniqueBaseOpcodeNames.end() &&
195 "Canonical index out of bounds in opcode section");
196 return *It;
197 }
198
199 auto getLocalIndex = [](unsigned Pos, size_t BaseOffset, size_t Bound,
200 const char *Msg) {
201 unsigned LocalIndex = Pos - BaseOffset;
202 assert(LocalIndex < Bound && Msg);
203 return LocalIndex;
204 };
205
206 // Handle common operands section
207 if (Pos < Layout.PhyRegBase) {
208 unsigned LocalIndex = getLocalIndex(
209 Pos, Layout.CommonOperandBase, std::size(CommonOperandNames),
210 "Local index out of bounds in common operands");
211 return CommonOperandNames[LocalIndex].str();
212 }
213
214 // Handle physical registers section
215 if (Pos < Layout.VirtRegBase) {
216 unsigned LocalIndex =
217 getLocalIndex(Pos, Layout.PhyRegBase, RegisterOperandNames.size(),
218 "Local index out of bounds in physical registers");
219 return "PhyReg_" + RegisterOperandNames[LocalIndex];
220 }
221
222 // Handle virtual registers section
223 unsigned LocalIndex =
224 getLocalIndex(Pos, Layout.VirtRegBase, RegisterOperandNames.size(),
225 "Local index out of bounds in virtual registers");
226 return "VirtReg_" + RegisterOperandNames[LocalIndex];
227}
228
229void MIRVocabulary::generateStorage(const VocabMap &OpcodeMap,
230 const VocabMap &CommonOperandsMap,
231 const VocabMap &PhyRegMap,
232 const VocabMap &VirtRegMap) {
233
234 // Helper for handling missing entities in the vocabulary.
235 // Currently, we use a zero vector. In the future, we will throw an error to
236 // ensure that *all* known entities are present in the vocabulary.
237 auto handleMissingEntity = [](StringRef Key) {
238 LLVM_DEBUG(errs() << "MIR2Vec: Missing vocabulary entry for " << Key
239 << "; using zero vector. This will result in an error "
240 "in the future.\n");
241 ++MIRVocabMissCounter;
242 };
243
244 // Initialize opcode embeddings section
245 unsigned EmbeddingDim = OpcodeMap.begin()->second.size();
246 std::vector<Embedding> OpcodeEmbeddings(Layout.CommonOperandBase,
247 Embedding(EmbeddingDim));
248
249 // Populate opcode embeddings using canonical mapping
250 for (auto COpcodeName : UniqueBaseOpcodeNames) {
251 if (auto It = OpcodeMap.find(x: COpcodeName); It != OpcodeMap.end()) {
252 auto COpcodeIndex = getCanonicalIndexForBaseName(BaseName: COpcodeName);
253 assert(COpcodeIndex < Layout.CommonOperandBase &&
254 "Canonical index out of bounds");
255 OpcodeEmbeddings[COpcodeIndex] = It->second;
256 } else {
257 handleMissingEntity(COpcodeName);
258 }
259 }
260
261 // Initialize common operand embeddings section
262 std::vector<Embedding> CommonOperandEmbeddings(std::size(CommonOperandNames),
263 Embedding(EmbeddingDim));
264 unsigned OperandIndex = 0;
265 for (const auto &CommonOperandName : CommonOperandNames) {
266 if (auto It = CommonOperandsMap.find(x: CommonOperandName.str());
267 It != CommonOperandsMap.end()) {
268 CommonOperandEmbeddings[OperandIndex] = It->second;
269 } else {
270 handleMissingEntity(CommonOperandName);
271 }
272 ++OperandIndex;
273 }
274
275 // Helper lambda for creating register operand embeddings
276 auto createRegisterEmbeddings = [&](const VocabMap &RegMap) {
277 std::vector<Embedding> RegEmbeddings(TRI.getNumRegClasses(),
278 Embedding(EmbeddingDim));
279 unsigned RegOperandIndex = 0;
280 for (const auto &RegOperandName : RegisterOperandNames) {
281 if (auto It = RegMap.find(x: RegOperandName); It != RegMap.end())
282 RegEmbeddings[RegOperandIndex] = It->second;
283 else
284 handleMissingEntity(RegOperandName);
285 ++RegOperandIndex;
286 }
287 return RegEmbeddings;
288 };
289
290 // Initialize register operand embeddings sections
291 std::vector<Embedding> PhyRegEmbeddings = createRegisterEmbeddings(PhyRegMap);
292 std::vector<Embedding> VirtRegEmbeddings =
293 createRegisterEmbeddings(VirtRegMap);
294
295 // Scale the vocabulary sections based on the provided weights
296 auto scaleVocabSection = [](std::vector<Embedding> &Embeddings,
297 double Weight) {
298 for (auto &Embedding : Embeddings)
299 Embedding *= Weight;
300 };
301 scaleVocabSection(OpcodeEmbeddings, OpcWeight);
302 scaleVocabSection(CommonOperandEmbeddings, CommonOperandWeight);
303 scaleVocabSection(PhyRegEmbeddings, RegOperandWeight);
304 scaleVocabSection(VirtRegEmbeddings, RegOperandWeight);
305
306 std::vector<std::vector<Embedding>> Sections(
307 static_cast<unsigned>(Section::MaxSections));
308 Sections[static_cast<unsigned>(Section::Opcodes)] =
309 std::move(OpcodeEmbeddings);
310 Sections[static_cast<unsigned>(Section::CommonOperands)] =
311 std::move(CommonOperandEmbeddings);
312 Sections[static_cast<unsigned>(Section::PhyRegisters)] =
313 std::move(PhyRegEmbeddings);
314 Sections[static_cast<unsigned>(Section::VirtRegisters)] =
315 std::move(VirtRegEmbeddings);
316
317 Storage = ir2vec::VocabStorage(std::move(Sections));
318}
319
320void MIRVocabulary::buildCanonicalOpcodeMapping() {
321 // Check if already built
322 if (!UniqueBaseOpcodeNames.empty())
323 return;
324
325 // Build mapping from opcodes to canonical base opcode indices
326 for (unsigned Opcode = 0; Opcode < TII.getNumOpcodes(); ++Opcode) {
327 std::string BaseOpcode = extractBaseOpcodeName(InstrName: TII.getName(Opcode));
328 UniqueBaseOpcodeNames.insert(x: BaseOpcode);
329 }
330
331 LLVM_DEBUG(dbgs() << "MIR2Vec: Built canonical mapping for target with "
332 << UniqueBaseOpcodeNames.size()
333 << " unique base opcodes\n");
334}
335
336void MIRVocabulary::buildRegisterOperandMapping() {
337 // Check if already built
338 if (!RegisterOperandNames.empty())
339 return;
340
341 for (unsigned RC = 0; RC < TRI.getNumRegClasses(); ++RC) {
342 const TargetRegisterClass *RegClass = TRI.getRegClass(i: RC);
343 if (!RegClass)
344 continue;
345
346 // Get the register class name
347 StringRef ClassName = TRI.getRegClassName(Class: RegClass);
348 RegisterOperandNames.push_back(Elt: ClassName.str());
349 }
350}
351
352unsigned MIRVocabulary::getCommonOperandIndex(
353 MachineOperand::MachineOperandType OperandType) const {
354 assert(OperandType != MachineOperand::MO_Register &&
355 "Expected non-register operand type");
356 assert(OperandType > MachineOperand::MO_Register &&
357 OperandType < MachineOperand::MO_Last && "Operand type out of bounds");
358 return static_cast<unsigned>(OperandType) - 1;
359}
360
361std::optional<unsigned>
362MIRVocabulary::getRegisterOperandIndex(Register Reg) const {
363 assert(!RegisterOperandNames.empty() && "Register operand mapping not built");
364 assert(Reg.isValid() && "Invalid register; not expected here");
365 assert((Reg.isPhysical() || Reg.isVirtual()) &&
366 "Expected a physical or virtual register");
367
368 const TargetRegisterClass *RegClass = nullptr;
369
370 // For physical registers, use TRI to get minimal register class as a
371 // physical register can belong to multiple classes. For virtual
372 // registers, use MRI to uniquely identify the assigned register class.
373 if (Reg.isPhysical())
374 RegClass = TRI.getMinimalPhysRegClass(Reg);
375 else
376 RegClass = MRI.getRegClassOrNull(Reg);
377
378 // Not every register belongs to a register class. This can happen for
379 // physical registers, e.g. X86's $mxcsr and $fpcw or AMDGPU's $mode, for
380 // which getMinimalPhysRegClass() returns nullptr. It can also happen for
381 // generic virtual registers that have not yet been through (or completed)
382 // GlobalISel's register bank selection, and thus carry an LLT or a
383 // RegisterBank instead of a TargetRegisterClass, for which
384 // getRegClassOrNull() returns nullptr.
385 // TODO: Avoid special-casing these registers at every use site. Classless
386 // registers currently fall back to a zero embedding in operator[] and to
387 // VirtRegBase in getEntityIDForRegister(), which is the same ad-hoc handling
388 // the invalid/stack-slot cases already get. Give them a real vocabulary
389 // representation instead -- e.g. an explicit "no register class" entry, or
390 // keying generic vregs on their LLT/RegisterBank -- so that the lookup is
391 // total and the callers need no fallbacks.
392 if (!RegClass) {
393 LLVM_DEBUG(errs() << "MIR2Vec: No register class for register " << Reg.id()
394 << "; using zero vector.\n");
395 ++MIRClasslessRegCounter;
396 return std::nullopt;
397 }
398
399 return RegClass->getID();
400}
401
402Expected<MIRVocabulary> MIRVocabulary::createDummyVocabForTest(
403 const TargetInstrInfo &TII, const TargetRegisterInfo &TRI,
404 const MachineRegisterInfo &MRI, unsigned Dim) {
405 assert(Dim > 0 && "Dimension must be greater than zero");
406
407 float DummyVal = 0.1f;
408
409 VocabMap DummyOpcMap, DummyOperandMap, DummyPhyRegMap, DummyVirtRegMap;
410
411 // Process opcodes directly without creating temporary vocabulary
412 for (unsigned Opcode = 0; Opcode < TII.getNumOpcodes(); ++Opcode) {
413 std::string BaseOpcode = extractBaseOpcodeName(InstrName: TII.getName(Opcode));
414 if (DummyOpcMap.count(x: BaseOpcode) == 0) { // Only add if not already present
415 DummyOpcMap[BaseOpcode] = Embedding(Dim, DummyVal);
416 DummyVal += 0.1f;
417 }
418 }
419
420 // Add common operands
421 for (const auto &CommonOperandName : CommonOperandNames) {
422 DummyOperandMap[CommonOperandName.str()] = Embedding(Dim, DummyVal);
423 DummyVal += 0.1f;
424 }
425
426 // Process register classes directly
427 for (unsigned RC = 0; RC < TRI.getNumRegClasses(); ++RC) {
428 const TargetRegisterClass *RegClass = TRI.getRegClass(i: RC);
429 if (!RegClass)
430 continue;
431
432 std::string ClassName = TRI.getRegClassName(Class: RegClass);
433 DummyPhyRegMap[ClassName] = Embedding(Dim, DummyVal);
434 DummyVirtRegMap[ClassName] = Embedding(Dim, DummyVal);
435 DummyVal += 0.1f;
436 }
437
438 // Create vocabulary directly without temporary instance
439 return MIRVocabulary::create(
440 OpcodeMap: std::move(DummyOpcMap), CommonOperandMap: std::move(DummyOperandMap),
441 PhyRegMap: std::move(DummyPhyRegMap), VirtRegMap: std::move(DummyVirtRegMap), TII, TRI, MRI);
442}
443
444//===----------------------------------------------------------------------===//
445// MIR2VecVocabProvider and MIR2VecVocabLegacyAnalysis
446//===----------------------------------------------------------------------===//
447
448Expected<mir2vec::MIRVocabulary>
449MIR2VecVocabProvider::getVocabulary(const Module &M) {
450 VocabMap OpcVocab, CommonOperandVocab, PhyRegVocabMap, VirtRegVocabMap;
451
452 if (Error Err = readVocabulary(OpcVocab, CommonOperandVocab, PhyRegVocabMap,
453 VirtRegVocabMap))
454 return std::move(Err);
455
456 for (const auto &F : M) {
457 if (F.isDeclaration())
458 continue;
459
460 if (auto *MF = MMI.getMachineFunction(F)) {
461 auto &Subtarget = MF->getSubtarget();
462 if (const auto *TII = Subtarget.getInstrInfo())
463 if (const auto *TRI = Subtarget.getRegisterInfo())
464 return mir2vec::MIRVocabulary::create(
465 OpcodeMap: std::move(OpcVocab), CommonOperandMap: std::move(CommonOperandVocab),
466 PhyRegMap: std::move(PhyRegVocabMap), VirtRegMap: std::move(VirtRegVocabMap), TII: *TII, TRI: *TRI,
467 MRI: MF->getRegInfo());
468 }
469 }
470 return createStringError(EC: errc::invalid_argument,
471 S: "No machine functions found in module");
472}
473
474Error MIR2VecVocabProvider::readVocabulary(VocabMap &OpcodeVocab,
475 VocabMap &CommonOperandVocab,
476 VocabMap &PhyRegVocabMap,
477 VocabMap &VirtRegVocabMap) {
478 if (VocabFile.empty())
479 return createStringError(
480 EC: errc::invalid_argument,
481 S: "MIR2Vec vocabulary file path not specified; set it "
482 "using --mir2vec-vocab-path");
483
484 auto BufOrError = MemoryBuffer::getFileOrSTDIN(Filename: VocabFile, /*IsText=*/true);
485 if (!BufOrError)
486 return createFileError(F: VocabFile, EC: BufOrError.getError());
487
488 auto Content = BufOrError.get()->getBuffer();
489
490 Expected<json::Value> ParsedVocabValue = json::parse(JSON: Content);
491 if (!ParsedVocabValue)
492 return ParsedVocabValue.takeError();
493
494 unsigned OpcodeDim = 0, CommonOperandDim = 0, PhyRegOperandDim = 0,
495 VirtRegOperandDim = 0;
496 if (auto Err = ir2vec::VocabStorage::parseVocabSection(
497 Key: "Opcodes", ParsedVocabValue: *ParsedVocabValue, TargetVocab&: OpcodeVocab, Dim&: OpcodeDim))
498 return Err;
499
500 if (auto Err = ir2vec::VocabStorage::parseVocabSection(
501 Key: "CommonOperands", ParsedVocabValue: *ParsedVocabValue, TargetVocab&: CommonOperandVocab,
502 Dim&: CommonOperandDim))
503 return Err;
504
505 if (auto Err = ir2vec::VocabStorage::parseVocabSection(
506 Key: "PhysicalRegisters", ParsedVocabValue: *ParsedVocabValue, TargetVocab&: PhyRegVocabMap,
507 Dim&: PhyRegOperandDim))
508 return Err;
509
510 if (auto Err = ir2vec::VocabStorage::parseVocabSection(
511 Key: "VirtualRegisters", ParsedVocabValue: *ParsedVocabValue, TargetVocab&: VirtRegVocabMap,
512 Dim&: VirtRegOperandDim))
513 return Err;
514
515 // All sections must have the same embedding dimension
516 if (!(OpcodeDim == CommonOperandDim && CommonOperandDim == PhyRegOperandDim &&
517 PhyRegOperandDim == VirtRegOperandDim)) {
518 return createStringError(
519 EC: errc::illegal_byte_sequence,
520 S: "MIR2Vec vocabulary sections have different dimensions");
521 }
522
523 return Error::success();
524}
525
526char MIR2VecVocabLegacyAnalysis::ID = 0;
527INITIALIZE_PASS_BEGIN(MIR2VecVocabLegacyAnalysis, "mir2vec-vocab-analysis",
528 "MIR2Vec Vocabulary Analysis", false, true)
529INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)
530INITIALIZE_PASS_END(MIR2VecVocabLegacyAnalysis, "mir2vec-vocab-analysis",
531 "MIR2Vec Vocabulary Analysis", false, true)
532
533StringRef MIR2VecVocabLegacyAnalysis::getPassName() const {
534 return "MIR2Vec Vocabulary Analysis";
535}
536
537//===----------------------------------------------------------------------===//
538// MIREmbedder and its subclasses
539//===----------------------------------------------------------------------===//
540
541std::unique_ptr<MIREmbedder> MIREmbedder::create(MIR2VecKind Mode,
542 const MachineFunction &MF,
543 const MIRVocabulary &Vocab) {
544 switch (Mode) {
545 case MIR2VecKind::Symbolic:
546 return std::make_unique<SymbolicMIREmbedder>(args: MF, args: Vocab);
547 }
548 return nullptr;
549}
550
551Embedding MIREmbedder::computeEmbeddings(const MachineBasicBlock &MBB) const {
552 Embedding MBBVector(Dimension, 0);
553
554 // Get instruction info for opcode name resolution
555 const auto &Subtarget = MF.getSubtarget();
556 const auto *TII = Subtarget.getInstrInfo();
557 if (!TII) {
558 MF.getFunction().getContext().emitError(
559 ErrorStr: "MIR2Vec: No TargetInstrInfo available; cannot compute embeddings");
560 return MBBVector;
561 }
562
563 // Process each machine instruction in the basic block
564 for (const auto &MI : MBB) {
565 // Skip debug instructions and other metadata
566 if (MI.isDebugInstr())
567 continue;
568 MBBVector += computeEmbeddings(MI);
569 }
570
571 return MBBVector;
572}
573
574Embedding MIREmbedder::computeEmbeddings() const {
575 Embedding MFuncVector(Dimension, 0);
576
577 if (MF.empty())
578 return MFuncVector;
579
580 // Consider all reachable machine basic blocks in the function
581 for (const auto *MBB : depth_first(G: &MF))
582 MFuncVector += computeEmbeddings(MBB: *MBB);
583 return MFuncVector;
584}
585
586SymbolicMIREmbedder::SymbolicMIREmbedder(const MachineFunction &MF,
587 const MIRVocabulary &Vocab)
588 : MIREmbedder(MF, Vocab) {}
589
590std::unique_ptr<SymbolicMIREmbedder>
591SymbolicMIREmbedder::create(const MachineFunction &MF,
592 const MIRVocabulary &Vocab) {
593 return std::make_unique<SymbolicMIREmbedder>(args: MF, args: Vocab);
594}
595
596Embedding SymbolicMIREmbedder::computeEmbeddings(const MachineInstr &MI) const {
597 // Skip debug instructions and other metadata
598 if (MI.isDebugInstr())
599 return Embedding(Dimension, 0);
600
601 // Opcode embedding
602 Embedding InstructionEmbedding = Vocab[MI.getOpcode()];
603
604 // Add operand contributions
605 for (const MachineOperand &MO : MI.operands())
606 InstructionEmbedding += Vocab[MO];
607
608 return InstructionEmbedding;
609}
610
611//===----------------------------------------------------------------------===//
612// Printer Passes
613//===----------------------------------------------------------------------===//
614
615char MIR2VecVocabPrinterLegacyPass::ID = 0;
616INITIALIZE_PASS_BEGIN(MIR2VecVocabPrinterLegacyPass, "print-mir2vec-vocab",
617 "MIR2Vec Vocabulary Printer Pass", false, true)
618INITIALIZE_PASS_DEPENDENCY(MIR2VecVocabLegacyAnalysis)
619INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)
620INITIALIZE_PASS_END(MIR2VecVocabPrinterLegacyPass, "print-mir2vec-vocab",
621 "MIR2Vec Vocabulary Printer Pass", false, true)
622
623bool MIR2VecVocabPrinterLegacyPass::runOnMachineFunction(MachineFunction &MF) {
624 return false;
625}
626
627bool MIR2VecVocabPrinterLegacyPass::doFinalization(Module &M) {
628 auto &Analysis = getAnalysis<MIR2VecVocabLegacyAnalysis>();
629 auto MIR2VecVocabOrErr = Analysis.getMIR2VecVocabulary(M);
630
631 if (!MIR2VecVocabOrErr) {
632 OS << "MIR2Vec Vocabulary Printer: Failed to get vocabulary - "
633 << toString(E: MIR2VecVocabOrErr.takeError()) << "\n";
634 return false;
635 }
636
637 auto &MIR2VecVocab = *MIR2VecVocabOrErr;
638 unsigned Pos = 0;
639 for (const auto &Entry : MIR2VecVocab) {
640 // Skip zero embeddings to avoid printing entries not in the vocabulary.
641 // This makes the output stable across changes to the opcode list.
642 if (PrintAllVocabEntries || !Entry.isZero()) {
643 OS << "Key: " << MIR2VecVocab.getStringKey(Pos) << ": ";
644 Entry.print(OS);
645 }
646 ++Pos;
647 }
648
649 return false;
650}
651
652MachineFunctionPass *
653llvm::createMIR2VecVocabPrinterLegacyPass(raw_ostream &OS) {
654 return new MIR2VecVocabPrinterLegacyPass(OS);
655}
656
657char MIR2VecPrinterLegacyPass::ID = 0;
658INITIALIZE_PASS_BEGIN(MIR2VecPrinterLegacyPass, "print-mir2vec",
659 "MIR2Vec Embedder Printer Pass", false, true)
660INITIALIZE_PASS_DEPENDENCY(MIR2VecVocabLegacyAnalysis)
661INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)
662INITIALIZE_PASS_END(MIR2VecPrinterLegacyPass, "print-mir2vec",
663 "MIR2Vec Embedder Printer Pass", false, true)
664
665bool MIR2VecPrinterLegacyPass::runOnMachineFunction(MachineFunction &MF) {
666 auto &Analysis = getAnalysis<MIR2VecVocabLegacyAnalysis>();
667 auto VocabOrErr =
668 Analysis.getMIR2VecVocabulary(M: *MF.getFunction().getParent());
669 assert(VocabOrErr && "Failed to get MIR2Vec vocabulary");
670 auto &MIRVocab = *VocabOrErr;
671
672 auto Emb = mir2vec::MIREmbedder::create(Mode: MIR2VecEmbeddingKind, MF, Vocab: MIRVocab);
673 if (!Emb) {
674 OS << "Error creating MIR2Vec embeddings for function " << MF.getName()
675 << "\n";
676 return false;
677 }
678
679 OS << "MIR2Vec embeddings for machine function " << MF.getName() << ":\n";
680 OS << "Machine Function vector: ";
681 Emb->getMFunctionVector().print(OS);
682
683 OS << "Machine basic block vectors:\n";
684 for (const MachineBasicBlock &MBB : MF) {
685 OS << "Machine basic block: " << MBB.getFullName() << ":\n";
686 Emb->getMBBVector(MBB).print(OS);
687 }
688
689 OS << "Machine instruction vectors:\n";
690 for (const MachineBasicBlock &MBB : MF) {
691 for (const MachineInstr &MI : MBB) {
692 // Skip debug instructions as they are not
693 // embedded
694 if (MI.isDebugInstr())
695 continue;
696
697 OS << "Machine instruction: ";
698 MI.print(OS);
699 Emb->getMInstVector(MI).print(OS);
700 }
701 }
702
703 return false;
704}
705
706MachineFunctionPass *llvm::createMIR2VecPrinterLegacyPass(raw_ostream &OS) {
707 return new MIR2VecPrinterLegacyPass(OS);
708}
709