766 lines · cpp
1//===- llvm-ir2vec.cpp - IR2Vec/MIR2Vec Embedding Generation Tool --------===//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-exception6//7//===----------------------------------------------------------------------===//8///9/// \file10/// This file implements the IR2Vec and MIR2Vec embedding generation tool.11///12/// This tool supports two modes:13/// - LLVM IR mode (-mode=llvm): Process LLVM IR14/// - Machine IR mode (-mode=mir): Process Machine IR15///16/// Available subcommands:17///18/// 1. Triplet Generation (triplets):19/// Generates numeric triplets (head, tail, relation) for vocabulary20/// training. Output format: MAX_RELATION=N header followed by21/// head\ttail\trelation lines. Relations: 0=Type, 1=Next, 2+=Arg0,Arg1,...22///23/// For LLVM IR:24/// llvm-ir2vec triplets input.bc -o train2id.txt25///26/// For Machine IR:27/// llvm-ir2vec triplets -mode=mir input.mir -o train2id.txt28///29/// 2. Entity Mappings (entities):30/// Generates entity mappings for vocabulary training.31/// Output format: <total_entities> header followed by entity\tid lines.32///33/// For LLVM IR:34/// llvm-ir2vec entities input.bc -o entity2id.txt35///36/// For Machine IR:37/// llvm-ir2vec entities -mode=mir input.mir -o entity2id.txt38///39/// 3. Embedding Generation (embeddings):40/// Generates IR2Vec/MIR2Vec embeddings using a trained vocabulary.41///42/// For LLVM IR:43/// llvm-ir2vec embeddings --ir2vec-vocab-path=vocab.json44/// --ir2vec-kind=<kind> --level=<level> input.bc -o embeddings.txt45/// Kind: --ir2vec-kind=symbolic (default), --ir2vec-kind=flow-aware46///47/// For Machine IR:48/// llvm-ir2vec embeddings -mode=mir --mir2vec-vocab-path=vocab.json49/// --level=<level> input.mir -o embeddings.txt50///51/// Levels: --level=inst (instructions), --level=bb (basic blocks),52/// --level=func (functions) (See IR2Vec.cpp/MIR2Vec.cpp for more embedding53/// generation options)54///55//===----------------------------------------------------------------------===//56 57#include "llvm/ADT/ArrayRef.h"58#include "llvm/Analysis/IR2Vec.h"59#include "llvm/IR/BasicBlock.h"60#include "llvm/IR/Function.h"61#include "llvm/IR/Instructions.h"62#include "llvm/IR/LLVMContext.h"63#include "llvm/IR/Module.h"64#include "llvm/IR/PassInstrumentation.h"65#include "llvm/IR/PassManager.h"66#include "llvm/IR/Type.h"67#include "llvm/IRReader/IRReader.h"68#include "llvm/Support/CommandLine.h"69#include "llvm/Support/Debug.h"70#include "llvm/Support/Errc.h"71#include "llvm/Support/InitLLVM.h"72#include "llvm/Support/SourceMgr.h"73#include "llvm/Support/raw_ostream.h"74 75#include "llvm/CodeGen/CommandFlags.h"76#include "llvm/CodeGen/MIR2Vec.h"77#include "llvm/CodeGen/MIRParser/MIRParser.h"78#include "llvm/CodeGen/MachineFunction.h"79#include "llvm/CodeGen/MachineModuleInfo.h"80#include "llvm/CodeGen/TargetInstrInfo.h"81#include "llvm/CodeGen/TargetRegisterInfo.h"82#include "llvm/MC/TargetRegistry.h"83#include "llvm/Support/TargetSelect.h"84#include "llvm/Support/WithColor.h"85#include "llvm/Target/TargetMachine.h"86#include "llvm/TargetParser/Host.h"87 88#define DEBUG_TYPE "ir2vec"89 90namespace llvm {91 92static const char *ToolName = "llvm-ir2vec";93 94// Common option category for options shared between IR2Vec and MIR2Vec95static cl::OptionCategory CommonCategory("Common Options",96 "Options applicable to both IR2Vec "97 "and MIR2Vec modes");98 99enum IRKind {100 LLVMIR = 0, ///< LLVM IR101 MIR ///< Machine IR102};103 104static cl::opt<IRKind>105 IRMode("mode", cl::desc("Tool operation mode:"),106 cl::values(clEnumValN(LLVMIR, "llvm", "Process LLVM IR"),107 clEnumValN(MIR, "mir", "Process Machine IR")),108 cl::init(LLVMIR), cl::cat(CommonCategory));109 110// Subcommands111static cl::SubCommand112 TripletsSubCmd("triplets", "Generate triplets for vocabulary training");113static cl::SubCommand114 EntitiesSubCmd("entities",115 "Generate entity mappings for vocabulary training");116static cl::SubCommand117 EmbeddingsSubCmd("embeddings",118 "Generate embeddings using trained vocabulary");119 120// Common options121static cl::opt<std::string> InputFilename(122 cl::Positional, cl::desc("<input bitcode/MIR file or '-' for stdin>"),123 cl::init("-"), cl::sub(TripletsSubCmd), cl::sub(EntitiesSubCmd),124 cl::sub(EmbeddingsSubCmd), cl::cat(CommonCategory));125 126static cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),127 cl::value_desc("filename"),128 cl::init("-"),129 cl::cat(CommonCategory));130 131// Embedding-specific options132static cl::opt<std::string>133 FunctionName("function", cl::desc("Process specific function only"),134 cl::value_desc("name"), cl::Optional, cl::init(""),135 cl::sub(EmbeddingsSubCmd), cl::cat(CommonCategory));136 137enum EmbeddingLevel {138 InstructionLevel, // Generate instruction-level embeddings139 BasicBlockLevel, // Generate basic block-level embeddings140 FunctionLevel // Generate function-level embeddings141};142 143static cl::opt<EmbeddingLevel>144 Level("level", cl::desc("Embedding generation level:"),145 cl::values(clEnumValN(InstructionLevel, "inst",146 "Generate instruction-level embeddings"),147 clEnumValN(BasicBlockLevel, "bb",148 "Generate basic block-level embeddings"),149 clEnumValN(FunctionLevel, "func",150 "Generate function-level embeddings")),151 cl::init(FunctionLevel), cl::sub(EmbeddingsSubCmd),152 cl::cat(CommonCategory));153 154namespace ir2vec {155 156/// Relation types for triplet generation157enum RelationType {158 TypeRelation = 0, ///< Instruction to type relationship159 NextRelation = 1, ///< Sequential instruction relationship160 ArgRelation = 2 ///< Instruction to operand relationship (ArgRelation + N)161};162 163/// Helper class for collecting IR triplets and generating embeddings164class IR2VecTool {165private:166 Module &M;167 ModuleAnalysisManager MAM;168 const Vocabulary *Vocab = nullptr;169 170public:171 explicit IR2VecTool(Module &M) : M(M) {}172 173 /// Initialize the IR2Vec vocabulary analysis174 bool initializeVocabulary() {175 // Register and run the IR2Vec vocabulary analysis176 // The vocabulary file path is specified via --ir2vec-vocab-path global177 // option178 MAM.registerPass([&] { return PassInstrumentationAnalysis(); });179 MAM.registerPass([&] { return IR2VecVocabAnalysis(); });180 // This will throw an error if vocab is not found or invalid181 Vocab = &MAM.getResult<IR2VecVocabAnalysis>(M);182 return Vocab->isValid();183 }184 185 /// Generate triplets for the module186 /// Output format: MAX_RELATION=N header followed by relationships187 void generateTriplets(raw_ostream &OS) const {188 unsigned MaxRelation = NextRelation; // Track maximum relation ID189 std::string Relationships;190 raw_string_ostream RelOS(Relationships);191 192 for (const Function &F : M) {193 unsigned FuncMaxRelation = generateTriplets(F, RelOS);194 MaxRelation = std::max(MaxRelation, FuncMaxRelation);195 }196 197 RelOS.flush();198 199 // Write metadata header followed by relationships200 OS << "MAX_RELATION=" << MaxRelation << '\n';201 OS << Relationships;202 }203 204 /// Generate triplets for a single function205 /// Returns the maximum relation ID used in this function206 unsigned generateTriplets(const Function &F, raw_ostream &OS) const {207 if (F.isDeclaration())208 return 0;209 210 unsigned MaxRelation = 1;211 unsigned PrevOpcode = 0;212 bool HasPrevOpcode = false;213 214 for (const BasicBlock &BB : F) {215 for (const auto &I : BB.instructionsWithoutDebug()) {216 unsigned Opcode = Vocabulary::getIndex(I.getOpcode());217 unsigned TypeID = Vocabulary::getIndex(I.getType()->getTypeID());218 219 // Add "Next" relationship with previous instruction220 if (HasPrevOpcode) {221 OS << PrevOpcode << '\t' << Opcode << '\t' << NextRelation << '\n';222 LLVM_DEBUG(dbgs()223 << Vocabulary::getVocabKeyForOpcode(PrevOpcode + 1) << '\t'224 << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'225 << "Next\n");226 }227 228 // Add "Type" relationship229 OS << Opcode << '\t' << TypeID << '\t' << TypeRelation << '\n';230 LLVM_DEBUG(231 dbgs() << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'232 << Vocabulary::getVocabKeyForTypeID(I.getType()->getTypeID())233 << '\t' << "Type\n");234 235 // Add "Arg" relationships236 unsigned ArgIndex = 0;237 for (const Use &U : I.operands()) {238 unsigned OperandID = Vocabulary::getIndex(*U.get());239 unsigned RelationID = ArgRelation + ArgIndex;240 OS << Opcode << '\t' << OperandID << '\t' << RelationID << '\n';241 242 LLVM_DEBUG({243 StringRef OperandStr = Vocabulary::getVocabKeyForOperandKind(244 Vocabulary::getOperandKind(U.get()));245 dbgs() << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'246 << OperandStr << '\t' << "Arg" << ArgIndex << '\n';247 });248 249 ++ArgIndex;250 }251 // Only update MaxRelation if there were operands252 if (ArgIndex > 0) {253 MaxRelation = std::max(MaxRelation, ArgRelation + ArgIndex - 1);254 }255 PrevOpcode = Opcode;256 HasPrevOpcode = true;257 }258 }259 260 return MaxRelation;261 }262 263 /// Dump entity ID to string mappings264 static void generateEntityMappings(raw_ostream &OS) {265 auto EntityLen = Vocabulary::getCanonicalSize();266 OS << EntityLen << "\n";267 for (unsigned EntityID = 0; EntityID < EntityLen; ++EntityID)268 OS << Vocabulary::getStringKey(EntityID) << '\t' << EntityID << '\n';269 }270 271 /// Generate embeddings for the entire module272 void generateEmbeddings(raw_ostream &OS) const {273 if (!Vocab->isValid()) {274 WithColor::error(errs(), ToolName)275 << "Vocabulary is not valid. IR2VecTool not initialized.\n";276 return;277 }278 279 for (const Function &F : M)280 generateEmbeddings(F, OS);281 }282 283 /// Generate embeddings for a single function284 void generateEmbeddings(const Function &F, raw_ostream &OS) const {285 if (F.isDeclaration()) {286 OS << "Function " << F.getName() << " is a declaration, skipping.\n";287 return;288 }289 290 // Create embedder for this function291 assert(Vocab->isValid() && "Vocabulary is not valid");292 auto Emb = Embedder::create(IR2VecEmbeddingKind, F, *Vocab);293 if (!Emb) {294 WithColor::error(errs(), ToolName)295 << "Failed to create embedder for function " << F.getName() << "\n";296 return;297 }298 299 OS << "Function: " << F.getName() << "\n";300 301 // Generate embeddings based on the specified level302 switch (Level) {303 case FunctionLevel: {304 Emb->getFunctionVector().print(OS);305 break;306 }307 case BasicBlockLevel: {308 for (const BasicBlock &BB : F) {309 OS << BB.getName() << ":";310 Emb->getBBVector(BB).print(OS);311 }312 break;313 }314 case InstructionLevel: {315 for (const BasicBlock &BB : F) {316 for (const Instruction &I : BB) {317 I.print(OS);318 Emb->getInstVector(I).print(OS);319 }320 }321 break;322 }323 }324 }325};326 327Error processModule(Module &M, raw_ostream &OS) {328 IR2VecTool Tool(M);329 330 if (EmbeddingsSubCmd) {331 // Initialize vocabulary for embedding generation332 // Note: Requires --ir2vec-vocab-path option to be set333 auto VocabStatus = Tool.initializeVocabulary();334 assert(VocabStatus && "Failed to initialize IR2Vec vocabulary");335 (void)VocabStatus;336 337 if (!FunctionName.empty()) {338 // Process single function339 if (const Function *F = M.getFunction(FunctionName))340 Tool.generateEmbeddings(*F, OS);341 else342 return createStringError(errc::invalid_argument,343 "Function '%s' not found",344 FunctionName.c_str());345 } else {346 // Process all functions347 Tool.generateEmbeddings(OS);348 }349 } else {350 // Both triplets and entities use triplet generation351 Tool.generateTriplets(OS);352 }353 return Error::success();354}355} // namespace ir2vec356 357namespace mir2vec {358 359/// Relation types for MIR2Vec triplet generation360enum MIRRelationType {361 MIRNextRelation = 0, ///< Sequential instruction relationship362 MIRArgRelation = 1 ///< Instruction to operand relationship (ArgRelation + N)363};364 365/// Helper class for MIR2Vec embedding generation366class MIR2VecTool {367private:368 MachineModuleInfo &MMI;369 std::unique_ptr<MIRVocabulary> Vocab;370 371public:372 explicit MIR2VecTool(MachineModuleInfo &MMI) : MMI(MMI) {}373 374 /// Initialize MIR2Vec vocabulary from file (for embeddings generation)375 bool initializeVocabulary(const Module &M) {376 MIR2VecVocabProvider Provider(MMI);377 auto VocabOrErr = Provider.getVocabulary(M);378 if (!VocabOrErr) {379 WithColor::error(errs(), ToolName)380 << "Failed to load MIR2Vec vocabulary - "381 << toString(VocabOrErr.takeError()) << "\n";382 return false;383 }384 Vocab = std::make_unique<MIRVocabulary>(std::move(*VocabOrErr));385 return true;386 }387 388 /// Initialize vocabulary with layout information only.389 /// This creates a minimal vocabulary with correct layout but no actual390 /// embeddings. Sufficient for generating training data and entity mappings.391 ///392 /// Note: Requires target-specific information from the first machine function393 /// to determine the vocabulary layout (number of opcodes, register classes).394 ///395 /// FIXME: Use --target option to get target info directly, avoiding the need396 /// to parse machine functions for pre-training operations.397 bool initializeVocabularyForLayout(const Module &M) {398 for (const Function &F : M) {399 if (F.isDeclaration())400 continue;401 402 MachineFunction *MF = MMI.getMachineFunction(F);403 if (!MF)404 continue;405 406 const TargetInstrInfo &TII = *MF->getSubtarget().getInstrInfo();407 const TargetRegisterInfo &TRI = *MF->getSubtarget().getRegisterInfo();408 const MachineRegisterInfo &MRI = MF->getRegInfo();409 410 auto VocabOrErr =411 MIRVocabulary::createDummyVocabForTest(TII, TRI, MRI, 1);412 if (!VocabOrErr) {413 WithColor::error(errs(), ToolName)414 << "Failed to create dummy vocabulary - "415 << toString(VocabOrErr.takeError()) << "\n";416 return false;417 }418 Vocab = std::make_unique<MIRVocabulary>(std::move(*VocabOrErr));419 return true;420 }421 422 WithColor::error(errs(), ToolName)423 << "No machine functions found to initialize vocabulary\n";424 return false;425 }426 427 /// Generate triplets for the module428 /// Output format: MAX_RELATION=N header followed by relationships429 void generateTriplets(const Module &M, raw_ostream &OS) const {430 unsigned MaxRelation = MIRNextRelation; // Track maximum relation ID431 std::string Relationships;432 raw_string_ostream RelOS(Relationships);433 434 for (const Function &F : M) {435 if (F.isDeclaration())436 continue;437 438 MachineFunction *MF = MMI.getMachineFunction(F);439 if (!MF) {440 WithColor::warning(errs(), ToolName)441 << "No MachineFunction for " << F.getName() << "\n";442 continue;443 }444 445 unsigned FuncMaxRelation = generateTriplets(*MF, RelOS);446 MaxRelation = std::max(MaxRelation, FuncMaxRelation);447 }448 449 RelOS.flush();450 451 // Write metadata header followed by relationships452 OS << "MAX_RELATION=" << MaxRelation << '\n';453 OS << Relationships;454 }455 456 /// Generate triplets for a single machine function457 /// Returns the maximum relation ID used in this function458 unsigned generateTriplets(const MachineFunction &MF, raw_ostream &OS) const {459 unsigned MaxRelation = MIRNextRelation;460 unsigned PrevOpcode = 0;461 bool HasPrevOpcode = false;462 463 if (!Vocab) {464 WithColor::error(errs(), ToolName)465 << "MIR Vocabulary must be initialized for triplet generation.\n";466 return MaxRelation;467 }468 469 for (const MachineBasicBlock &MBB : MF) {470 for (const MachineInstr &MI : MBB) {471 // Skip debug instructions472 if (MI.isDebugInstr())473 continue;474 475 // Get opcode entity ID476 unsigned OpcodeID = Vocab->getEntityIDForOpcode(MI.getOpcode());477 478 // Add "Next" relationship with previous instruction479 if (HasPrevOpcode) {480 OS << PrevOpcode << '\t' << OpcodeID << '\t' << MIRNextRelation481 << '\n';482 LLVM_DEBUG(dbgs()483 << Vocab->getStringKey(PrevOpcode) << '\t'484 << Vocab->getStringKey(OpcodeID) << '\t' << "Next\n");485 }486 487 // Add "Arg" relationships for operands488 unsigned ArgIndex = 0;489 for (const MachineOperand &MO : MI.operands()) {490 auto OperandID = Vocab->getEntityIDForMachineOperand(MO);491 unsigned RelationID = MIRArgRelation + ArgIndex;492 OS << OpcodeID << '\t' << OperandID << '\t' << RelationID << '\n';493 LLVM_DEBUG({494 std::string OperandStr = Vocab->getStringKey(OperandID);495 dbgs() << Vocab->getStringKey(OpcodeID) << '\t' << OperandStr496 << '\t' << "Arg" << ArgIndex << '\n';497 });498 499 ++ArgIndex;500 }501 502 // Update MaxRelation if there were operands503 if (ArgIndex > 0)504 MaxRelation = std::max(MaxRelation, MIRArgRelation + ArgIndex - 1);505 506 PrevOpcode = OpcodeID;507 HasPrevOpcode = true;508 }509 }510 511 return MaxRelation;512 }513 514 /// Generate entity mappings with vocabulary515 void generateEntityMappings(raw_ostream &OS) const {516 if (!Vocab) {517 WithColor::error(errs(), ToolName)518 << "Vocabulary must be initialized for entity mappings.\n";519 return;520 }521 522 const unsigned EntityCount = Vocab->getCanonicalSize();523 OS << EntityCount << "\n";524 for (unsigned EntityID = 0; EntityID < EntityCount; ++EntityID)525 OS << Vocab->getStringKey(EntityID) << '\t' << EntityID << '\n';526 }527 528 /// Generate embeddings for all machine functions in the module529 void generateEmbeddings(const Module &M, raw_ostream &OS) const {530 if (!Vocab) {531 WithColor::error(errs(), ToolName) << "Vocabulary not initialized.\n";532 return;533 }534 535 for (const Function &F : M) {536 if (F.isDeclaration())537 continue;538 539 MachineFunction *MF = MMI.getMachineFunction(F);540 if (!MF) {541 WithColor::warning(errs(), ToolName)542 << "No MachineFunction for " << F.getName() << "\n";543 continue;544 }545 546 generateEmbeddings(*MF, OS);547 }548 }549 550 /// Generate embeddings for a specific machine function551 void generateEmbeddings(MachineFunction &MF, raw_ostream &OS) const {552 if (!Vocab) {553 WithColor::error(errs(), ToolName) << "Vocabulary not initialized.\n";554 return;555 }556 557 auto Emb = MIREmbedder::create(MIR2VecKind::Symbolic, MF, *Vocab);558 if (!Emb) {559 WithColor::error(errs(), ToolName)560 << "Failed to create embedder for " << MF.getName() << "\n";561 return;562 }563 564 OS << "MIR2Vec embeddings for machine function " << MF.getName() << ":\n";565 566 // Generate embeddings based on the specified level567 switch (Level) {568 case FunctionLevel: {569 OS << "Function vector: ";570 Emb->getMFunctionVector().print(OS);571 break;572 }573 case BasicBlockLevel: {574 OS << "Basic block vectors:\n";575 for (const MachineBasicBlock &MBB : MF) {576 OS << "MBB " << MBB.getName() << ": ";577 Emb->getMBBVector(MBB).print(OS);578 }579 break;580 }581 case InstructionLevel: {582 OS << "Instruction vectors:\n";583 for (const MachineBasicBlock &MBB : MF) {584 for (const MachineInstr &MI : MBB) {585 OS << MI << " -> ";586 Emb->getMInstVector(MI).print(OS);587 }588 }589 break;590 }591 }592 }593 594 const MIRVocabulary *getVocabulary() const { return Vocab.get(); }595};596 597} // namespace mir2vec598 599} // namespace llvm600 601int main(int argc, char **argv) {602 using namespace llvm;603 using namespace llvm::ir2vec;604 using namespace llvm::mir2vec;605 606 InitLLVM X(argc, argv);607 // Show Common, IR2Vec and MIR2Vec option categories608 cl::HideUnrelatedOptions(ArrayRef<const cl::OptionCategory *>{609 &CommonCategory, &ir2vec::IR2VecCategory, &mir2vec::MIR2VecCategory});610 cl::ParseCommandLineOptions(611 argc, argv,612 "IR2Vec/MIR2Vec - Embedding Generation Tool\n"613 "Generates embeddings for a given LLVM IR or MIR and "614 "supports triplet generation for vocabulary "615 "training and embedding generation.\n\n"616 "See https://llvm.org/docs/CommandGuide/llvm-ir2vec.html for more "617 "information.\n");618 619 std::error_code EC;620 raw_fd_ostream OS(OutputFilename, EC);621 if (EC) {622 WithColor::error(errs(), ToolName)623 << "opening output file: " << EC.message() << "\n";624 return 1;625 }626 627 if (IRMode == IRKind::LLVMIR) {628 if (EntitiesSubCmd) {629 // Just dump entity mappings without processing any IR630 IR2VecTool::generateEntityMappings(OS);631 return 0;632 }633 634 // Parse the input LLVM IR file or stdin635 SMDiagnostic Err;636 LLVMContext Context;637 std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);638 if (!M) {639 Err.print(ToolName, errs());640 return 1;641 }642 643 if (Error Err = processModule(*M, OS)) {644 handleAllErrors(std::move(Err), [&](const ErrorInfoBase &EIB) {645 WithColor::error(errs(), ToolName) << EIB.message() << "\n";646 });647 return 1;648 }649 return 0;650 }651 if (IRMode == IRKind::MIR) {652 // Initialize targets for Machine IR processing653 InitializeAllTargets();654 InitializeAllTargetMCs();655 InitializeAllAsmParsers();656 InitializeAllAsmPrinters();657 static codegen::RegisterCodeGenFlags CGF;658 659 // Parse MIR input file660 SMDiagnostic Err;661 LLVMContext Context;662 std::unique_ptr<TargetMachine> TM;663 664 auto MIR = createMIRParserFromFile(InputFilename, Err, Context);665 if (!MIR) {666 Err.print(ToolName, errs());667 return 1;668 }669 670 auto SetDataLayout = [&](StringRef DataLayoutTargetTriple,671 StringRef OldDLStr) -> std::optional<std::string> {672 std::string IRTargetTriple = DataLayoutTargetTriple.str();673 Triple TheTriple = Triple(IRTargetTriple);674 if (TheTriple.getTriple().empty())675 TheTriple.setTriple(sys::getDefaultTargetTriple());676 auto TMOrErr = codegen::createTargetMachineForTriple(TheTriple.str());677 if (!TMOrErr) {678 Err.print(ToolName, errs());679 exit(1);680 }681 TM = std::move(*TMOrErr);682 return TM->createDataLayout().getStringRepresentation();683 };684 685 std::unique_ptr<Module> M = MIR->parseIRModule(SetDataLayout);686 if (!M) {687 Err.print(ToolName, errs());688 return 1;689 }690 691 // Parse machine functions692 auto MMI = std::make_unique<MachineModuleInfo>(TM.get());693 if (!MMI || MIR->parseMachineFunctions(*M, *MMI)) {694 Err.print(ToolName, errs());695 return 1;696 }697 698 // Create MIR2Vec tool699 MIR2VecTool Tool(*MMI);700 701 // Initialize vocabulary. For triplet/entity generation, only layout is702 // needed For embedding generation, the full vocabulary is needed.703 //704 // Note: Unlike IR2Vec, MIR2Vec vocabulary initialization requires705 // target-specific information for generating the vocabulary layout. So, we706 // always initialize the vocabulary in this case.707 if (TripletsSubCmd || EntitiesSubCmd) {708 if (!Tool.initializeVocabularyForLayout(*M)) {709 WithColor::error(errs(), ToolName)710 << "Failed to initialize MIR2Vec vocabulary for layout.\n";711 return 1;712 }713 } else {714 if (!Tool.initializeVocabulary(*M)) {715 WithColor::error(errs(), ToolName)716 << "Failed to initialize MIR2Vec vocabulary.\n";717 return 1;718 }719 }720 assert(Tool.getVocabulary() &&721 "MIR2Vec vocabulary should be initialized at this point");722 LLVM_DEBUG(dbgs() << "MIR2Vec vocabulary loaded successfully.\n"723 << "Vocabulary dimension: "724 << Tool.getVocabulary()->getDimension() << "\n"725 << "Vocabulary size: "726 << Tool.getVocabulary()->getCanonicalSize() << "\n");727 728 // Handle subcommands729 if (TripletsSubCmd) {730 Tool.generateTriplets(*M, OS);731 } else if (EntitiesSubCmd) {732 Tool.generateEntityMappings(OS);733 } else if (EmbeddingsSubCmd) {734 if (!FunctionName.empty()) {735 // Process single function736 Function *F = M->getFunction(FunctionName);737 if (!F) {738 WithColor::error(errs(), ToolName)739 << "Function '" << FunctionName << "' not found\n";740 return 1;741 }742 743 MachineFunction *MF = MMI->getMachineFunction(*F);744 if (!MF) {745 WithColor::error(errs(), ToolName)746 << "No MachineFunction for " << FunctionName << "\n";747 return 1;748 }749 750 Tool.generateEmbeddings(*MF, OS);751 } else {752 // Process all functions753 Tool.generateEmbeddings(*M, OS);754 }755 } else {756 WithColor::error(errs(), ToolName)757 << "Please specify a subcommand: triplets, entities, or embeddings\n";758 return 1;759 }760 761 return 0;762 }763 764 return 0;765}766