[IR2Vec][llvm-ir2vec] Revamp triplet generation and add entity mapping mode #149214
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This was referenced Jul 16, 2025
This was referenced Jul 16, 2025
svkeerthy
changed the title
svkeerthy
requested review from
boomanaiden154,
kazutakahirata,
mtrofin and
snehasish
|
@llvm/pr-subscribers-mlgo @llvm/pr-subscribers-llvm-binary-utilities Author: S. VenkataKeerthy (svkeerthy) ChangesAdd entity mapping mode to llvm-ir2vec and improve triplet generation format for knowledge graph embedding training. This change streamlines the workflow for training the vocabulary embeddings with IR2Vec by:
These improvements make IR2Vec more compatible with standard knowledge graph embedding training pipelines and reduce the preprocessing steps needed before training. See #149215 for more details on how it is used. (Tracking issue - #141817) Patch is 20.46 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/149214.diff 4 Files Affected:
diff --git a/llvm/docs/CommandGuide/llvm-ir2vec.rst b/llvm/docs/CommandGuide/llvm-ir2vec.rst
index 13fe4996b968f..56ece4f509f6e 100644
--- a/llvm/docs/CommandGuide/llvm-ir2vec.rst
+++ b/llvm/docs/CommandGuide/llvm-ir2vec.rst
@@ -13,17 +13,21 @@ DESCRIPTION
:program:`llvm-ir2vec` is a standalone command-line tool for IR2Vec. It
generates IR2Vec embeddings for LLVM IR and supports triplet generation
-for vocabulary training. It provides two main operation modes:
+for vocabulary training. It provides three main operation modes:
-1. **Triplet Mode**: Generates triplets (opcode, type, operands) for vocabulary
+1. **Triplet Mode**: Generates numeric triplets in train2id format for vocabulary
training from LLVM IR.
-2. **Embedding Mode**: Generates IR2Vec embeddings using a trained vocabulary
+2. **Entity Mode**: Generates entity mapping files (entity2id.txt) for vocabulary
+ training.
+
+3. **Embedding Mode**: Generates IR2Vec embeddings using a trained vocabulary
at different granularity levels (instruction, basic block, or function).
The tool is designed to facilitate machine learning applications that work with
LLVM IR by converting the IR into numerical representations that can be used by
-ML models.
+ML models. The triplet mode generates numeric IDs directly instead of string
+triplets, streamlining the training data preparation workflow.
.. note::
@@ -34,18 +38,46 @@ ML models.
OPERATION MODES
---------------
+Triplet Generation and Entity Mapping Modes are used for preparing
+vocabulary and training data for knowledge graph embeddings. The Embedding Mode
+is used for generating embeddings from LLVM IR using a pre-trained vocabulary.
+
+The Seed Embedding Vocabulary of IR2Vec is trained on a large corpus of LLVM IR
+by modeling the relationships between opcodes, types, and operands as a knowledge
+graph. For this purpose, Triplet Generation and Entity Mapping Modes generate
+triplets and entity mappings in the standard format used for knowledge graph
+embedding training (see
+<https://github.com/thunlp/OpenKE/tree/OpenKE-PyTorch?tab=readme-ov-file#data-format>
+for details).
+
Triplet Generation Mode
~~~~~~~~~~~~~~~~~~~~~~~
-In triplet mode, :program:`llvm-ir2vec` analyzes LLVM IR and extracts triplets
-consisting of opcodes, types, and operands. These triplets can be used to train
-vocabularies for embedding generation.
+In triplet mode, :program:`llvm-ir2vec` analyzes LLVM IR and extracts numeric
+triplets consisting of opcode IDs, type IDs, and operand IDs. These triplets
+are generated in train2id format. The tool outputs numeric IDs directly using
+the ir2vec::Vocabulary mapping infrastructure, eliminating the need for
+string-to-ID preprocessing.
+
+Usage:
+
+.. code-block:: bash
+
+ llvm-ir2vec --mode=triplets input.bc -o triplets_train2id.txt
+
+Entity Mapping Generation Mode
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+In entity mode, :program:`llvm-ir2vec` generates the entity mappings supported by
+IR2Vec in entity2id format. This mode outputs all supported entities (opcodes,
+types, and operands) with their corresponding numeric IDs, and is not specific for
+an LLVM IR file.
Usage:
.. code-block:: bash
- llvm-ir2vec --mode=triplets input.bc -o triplets.txt
+ llvm-ir2vec --mode=entities -o entity2id.txt
Embedding Generation Mode
~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -67,6 +99,7 @@ OPTIONS
Specify the operation mode. Valid values are:
* ``triplets`` - Generate triplets for vocabulary training
+ * ``entities`` - Generate entity mappings for vocabulary training
* ``embeddings`` - Generate embeddings using trained vocabulary (default)
.. option:: --level=<level>
@@ -115,7 +148,7 @@ OPTIONS
``--level``, ``--function``, ``--ir2vec-vocab-path``, ``--ir2vec-opc-weight``,
``--ir2vec-type-weight``, and ``--ir2vec-arg-weight`` are only used in embedding
- mode. These options are ignored in triplet mode.
+ mode. These options are ignored in triplet and entity modes.
INPUT FILE FORMAT
-----------------
@@ -129,14 +162,34 @@ OUTPUT FORMAT
Triplet Mode Output
~~~~~~~~~~~~~~~~~~~
-In triplet mode, the output consists of lines containing space-separated triplets:
+In triplet mode, the output consists of numeric triplets in train2id format with
+metadata headers. The format includes:
+
+.. code-block:: text
+
+ MAX_RELATIONS=<max_relations_count>
+ <head_entity_id> <tail_entity_id> <relation_id>
+ <head_entity_id> <tail_entity_id> <relation_id>
+ ...
+
+Each line after the metadata header represents one instruction relationship,
+with numeric IDs for head entity, relation, and tail entity. The metadata
+header (MAX_RELATIONS) provides counts for post-processing and training setup.
+
+Entity Mode Output
+~~~~~~~~~~~~~~~~~~
+
+In entity mode, the output consists of entity mapping in the format:
.. code-block:: text
- <opcode> <type> <operand1> <operand2> ...
+ <total_entities>
+ <entity_string> <numeric_id>
+ <entity_string> <numeric_id>
+ ...
-Each line represents the information of one instruction, with the opcode, type,
-and operands.
+The first line contains the total number of entities, followed by one entity
+mapping per line with tab-separated entity string and numeric ID.
Embedding Mode Output
~~~~~~~~~~~~~~~~~~~~~
diff --git a/llvm/test/tools/llvm-ir2vec/entities.ll b/llvm/test/tools/llvm-ir2vec/entities.ll
new file mode 100644
index 0000000000000..57c3d6fa6d6c4
--- /dev/null
+++ b/llvm/test/tools/llvm-ir2vec/entities.ll
@@ -0,0 +1,95 @@
+; RUN: llvm-ir2vec --mode=entities | FileCheck %s
+
+CHECK: 92
+CHECK-NEXT: Ret 0
+CHECK-NEXT: Br 1
+CHECK-NEXT: Switch 2
+CHECK-NEXT: IndirectBr 3
+CHECK-NEXT: Invoke 4
+CHECK-NEXT: Resume 5
+CHECK-NEXT: Unreachable 6
+CHECK-NEXT: CleanupRet 7
+CHECK-NEXT: CatchRet 8
+CHECK-NEXT: CatchSwitch 9
+CHECK-NEXT: CallBr 10
+CHECK-NEXT: FNeg 11
+CHECK-NEXT: Add 12
+CHECK-NEXT: FAdd 13
+CHECK-NEXT: Sub 14
+CHECK-NEXT: FSub 15
+CHECK-NEXT: Mul 16
+CHECK-NEXT: FMul 17
+CHECK-NEXT: UDiv 18
+CHECK-NEXT: SDiv 19
+CHECK-NEXT: FDiv 20
+CHECK-NEXT: URem 21
+CHECK-NEXT: SRem 22
+CHECK-NEXT: FRem 23
+CHECK-NEXT: Shl 24
+CHECK-NEXT: LShr 25
+CHECK-NEXT: AShr 26
+CHECK-NEXT: And 27
+CHECK-NEXT: Or 28
+CHECK-NEXT: Xor 29
+CHECK-NEXT: Alloca 30
+CHECK-NEXT: Load 31
+CHECK-NEXT: Store 32
+CHECK-NEXT: GetElementPtr 33
+CHECK-NEXT: Fence 34
+CHECK-NEXT: AtomicCmpXchg 35
+CHECK-NEXT: AtomicRMW 36
+CHECK-NEXT: Trunc 37
+CHECK-NEXT: ZExt 38
+CHECK-NEXT: SExt 39
+CHECK-NEXT: FPToUI 40
+CHECK-NEXT: FPToSI 41
+CHECK-NEXT: UIToFP 42
+CHECK-NEXT: SIToFP 43
+CHECK-NEXT: FPTrunc 44
+CHECK-NEXT: FPExt 45
+CHECK-NEXT: PtrToInt 46
+CHECK-NEXT: IntToPtr 47
+CHECK-NEXT: BitCast 48
+CHECK-NEXT: AddrSpaceCast 49
+CHECK-NEXT: CleanupPad 50
+CHECK-NEXT: CatchPad 51
+CHECK-NEXT: ICmp 52
+CHECK-NEXT: FCmp 53
+CHECK-NEXT: PHI 54
+CHECK-NEXT: Call 55
+CHECK-NEXT: Select 56
+CHECK-NEXT: UserOp1 57
+CHECK-NEXT: UserOp2 58
+CHECK-NEXT: VAArg 59
+CHECK-NEXT: ExtractElement 60
+CHECK-NEXT: InsertElement 61
+CHECK-NEXT: ShuffleVector 62
+CHECK-NEXT: ExtractValue 63
+CHECK-NEXT: InsertValue 64
+CHECK-NEXT: LandingPad 65
+CHECK-NEXT: Freeze 66
+CHECK-NEXT: FloatTy 67
+CHECK-NEXT: FloatTy 68
+CHECK-NEXT: FloatTy 69
+CHECK-NEXT: FloatTy 70
+CHECK-NEXT: FloatTy 71
+CHECK-NEXT: FloatTy 72
+CHECK-NEXT: FloatTy 73
+CHECK-NEXT: VoidTy 74
+CHECK-NEXT: LabelTy 75
+CHECK-NEXT: MetadataTy 76
+CHECK-NEXT: UnknownTy 77
+CHECK-NEXT: TokenTy 78
+CHECK-NEXT: IntegerTy 79
+CHECK-NEXT: FunctionTy 80
+CHECK-NEXT: PointerTy 81
+CHECK-NEXT: StructTy 82
+CHECK-NEXT: ArrayTy 83
+CHECK-NEXT: VectorTy 84
+CHECK-NEXT: VectorTy 85
+CHECK-NEXT: PointerTy 86
+CHECK-NEXT: UnknownTy 87
+CHECK-NEXT: Function 88
+CHECK-NEXT: Pointer 89
+CHECK-NEXT: Constant 90
+CHECK-NEXT: Variable 91
diff --git a/llvm/test/tools/llvm-ir2vec/triplets.ll b/llvm/test/tools/llvm-ir2vec/triplets.ll
index d1ef5b388e258..dcd1dc9afb478 100644
--- a/llvm/test/tools/llvm-ir2vec/triplets.ll
+++ b/llvm/test/tools/llvm-ir2vec/triplets.ll
@@ -24,15 +24,42 @@ entry:
ret i32 %result
}
-; TRIPLETS: Add IntegerTy Variable Variable
-; TRIPLETS-NEXT: Ret VoidTy Variable
-; TRIPLETS-NEXT: Mul IntegerTy Variable Variable
-; TRIPLETS-NEXT: Ret VoidTy Variable
-; TRIPLETS-NEXT: Alloca PointerTy Constant
-; TRIPLETS-NEXT: Alloca PointerTy Constant
-; TRIPLETS-NEXT: Store VoidTy Variable Pointer
-; TRIPLETS-NEXT: Store VoidTy Variable Pointer
-; TRIPLETS-NEXT: Load IntegerTy Pointer
-; TRIPLETS-NEXT: Load IntegerTy Pointer
-; TRIPLETS-NEXT: Add IntegerTy Variable Variable
-; TRIPLETS-NEXT: Ret VoidTy Variable
+; TRIPLETS: MAX_RELATION=3
+; TRIPLETS-NEXT: 12 79 0
+; TRIPLETS-NEXT: 12 91 2
+; TRIPLETS-NEXT: 12 91 3
+; TRIPLETS-NEXT: 12 0 1
+; TRIPLETS-NEXT: 0 74 0
+; TRIPLETS-NEXT: 0 91 2
+; TRIPLETS-NEXT: 16 79 0
+; TRIPLETS-NEXT: 16 91 2
+; TRIPLETS-NEXT: 16 91 3
+; TRIPLETS-NEXT: 16 0 1
+; TRIPLETS-NEXT: 0 74 0
+; TRIPLETS-NEXT: 0 91 2
+; TRIPLETS-NEXT: 30 81 0
+; TRIPLETS-NEXT: 30 90 2
+; TRIPLETS-NEXT: 30 30 1
+; TRIPLETS-NEXT: 30 81 0
+; TRIPLETS-NEXT: 30 90 2
+; TRIPLETS-NEXT: 30 32 1
+; TRIPLETS-NEXT: 32 74 0
+; TRIPLETS-NEXT: 32 91 2
+; TRIPLETS-NEXT: 32 89 3
+; TRIPLETS-NEXT: 32 32 1
+; TRIPLETS-NEXT: 32 74 0
+; TRIPLETS-NEXT: 32 91 2
+; TRIPLETS-NEXT: 32 89 3
+; TRIPLETS-NEXT: 32 31 1
+; TRIPLETS-NEXT: 31 79 0
+; TRIPLETS-NEXT: 31 89 2
+; TRIPLETS-NEXT: 31 31 1
+; TRIPLETS-NEXT: 31 79 0
+; TRIPLETS-NEXT: 31 89 2
+; TRIPLETS-NEXT: 31 12 1
+; TRIPLETS-NEXT: 12 79 0
+; TRIPLETS-NEXT: 12 91 2
+; TRIPLETS-NEXT: 12 91 3
+; TRIPLETS-NEXT: 12 0 1
+; TRIPLETS-NEXT: 0 74 0
+; TRIPLETS-NEXT: 0 91 2
diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
index 3e6cb4b64fde5..40257c0d6aba4 100644
--- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
+++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
@@ -9,13 +9,20 @@
/// \file
/// This file implements the IR2Vec embedding generation tool.
///
-/// This tool provides two main functionalities:
+/// This tool provides three main modes:
///
/// 1. Triplet Generation Mode (--mode=triplets):
-/// Generates triplets (opcode, type, operands) for vocabulary training.
-/// Usage: llvm-ir2vec --mode=triplets input.bc -o triplets.txt
+/// Generates numeric triplets (head, tail, relation) for vocabulary
+/// training. Output format: MAX_RELATION=N header followed by
+/// head\ttail\trelation lines. Relations: 0=Type, 1=Next, 2+=Arg0,Arg1,...
+/// Usage: llvm-ir2vec --mode=triplets input.bc -o train2id.txt
///
-/// 2. Embedding Generation Mode (--mode=embeddings):
+/// 2. Entities Generation Mode (--mode=entities):
+/// Generates entity mappings for vocabulary training.
+/// Output format: <total_entities> header followed by entity\tid lines.
+/// Usage: llvm-ir2vec --mode=entities input.bc -o entity2id.txt
+///
+/// 3. Embedding Generation Mode (--mode=embeddings):
/// Generates IR2Vec embeddings using a trained vocabulary.
/// Usage: llvm-ir2vec --mode=embeddings --ir2vec-vocab-path=vocab.json
/// --level=func input.bc -o embeddings.txt Levels: --level=inst
@@ -61,16 +68,19 @@ static cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),
enum ToolMode {
TripletMode, // Generate triplets for vocabulary training
+ EntityMode, // Generate entity mappings for vocabulary training
EmbeddingMode // Generate embeddings using trained vocabulary
};
-static cl::opt<ToolMode>
- Mode("mode", cl::desc("Tool operation mode:"),
- cl::values(clEnumValN(TripletMode, "triplets",
- "Generate triplets for vocabulary training"),
- clEnumValN(EmbeddingMode, "embeddings",
- "Generate embeddings using trained vocabulary")),
- cl::init(EmbeddingMode), cl::cat(IR2VecToolCategory));
+static cl::opt<ToolMode> Mode(
+ "mode", cl::desc("Tool operation mode:"),
+ cl::values(clEnumValN(TripletMode, "triplets",
+ "Generate triplets for vocabulary training"),
+ clEnumValN(EntityMode, "entities",
+ "Generate entity mappings for vocabulary training"),
+ clEnumValN(EmbeddingMode, "embeddings",
+ "Generate embeddings using trained vocabulary")),
+ cl::init(EmbeddingMode), cl::cat(IR2VecToolCategory));
static cl::opt<std::string>
FunctionName("function", cl::desc("Process specific function only"),
@@ -95,6 +105,13 @@ static cl::opt<EmbeddingLevel>
namespace {
+/// Relation types for triplet generation
+enum RelationType {
+ TypeRelation = 0, ///< Instruction to type relationship
+ NextRelation = 1, ///< Sequential instruction relationship
+ ArgRelation = 2 ///< Instruction to operand relationship (ArgRelation + N)
+};
+
/// Helper class for collecting IR triplets and generating embeddings
class IR2VecTool {
private:
@@ -116,25 +133,96 @@ class IR2VecTool {
return Vocab->isValid();
}
- /// Generate triplets for the entire module
+ /// Generate triplets for the module
+ /// Output format: MAX_RELATION=N header followed by relationships
void generateTriplets(raw_ostream &OS) const {
- for (const Function &F : M)
- generateTriplets(F, OS);
+ unsigned MaxRelation = NextRelation; // Track maximum relation ID
+ std::string Relationships;
+ raw_string_ostream RelOS(Relationships);
+
+ for (const Function &F : M) {
+ unsigned FuncMaxRelation = generateTriplets(F, RelOS);
+ MaxRelation = std::max(MaxRelation, FuncMaxRelation);
+ }
+
+ RelOS.flush();
+
+ // Write metadata header followed by relationships
+ OS << "MAX_RELATION=" << MaxRelation << '\n';
+ OS << Relationships;
}
/// Generate triplets for a single function
- void generateTriplets(const Function &F, raw_ostream &OS) const {
+ /// Returns the maximum relation ID used in this function
+ unsigned generateTriplets(const Function &F, raw_ostream &OS) const {
if (F.isDeclaration())
- return;
+ return 0;
+
+ unsigned MaxRelation = 1;
+ unsigned PrevOpcode = 0;
+ bool HasPrevOpcode = false;
+
+ for (const BasicBlock &BB : F) {
+ for (const auto &I : BB.instructionsWithoutDebug()) {
+ unsigned Opcode = Vocabulary::getNumericID(I.getOpcode());
+ unsigned TypeID = Vocabulary::getNumericID(I.getType()->getTypeID());
+
+ // Add "Next" relationship with previous instruction
+ if (HasPrevOpcode) {
+ OS << PrevOpcode << '\t' << Opcode << '\t' << NextRelation << '\n';
+ LLVM_DEBUG(dbgs()
+ << Vocabulary::getVocabKeyForOpcode(PrevOpcode + 1) << '\t'
+ << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'
+ << "Next\n");
+ }
- std::string LocalOutput;
- raw_string_ostream LocalOS(LocalOutput);
+ // Add "Type" relationship
+ OS << Opcode << '\t' << TypeID << '\t' << TypeRelation << '\n';
+ LLVM_DEBUG(
+ dbgs() << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'
+ << Vocabulary::getVocabKeyForTypeID(I.getType()->getTypeID())
+ << '\t' << "Type\n");
+
+ // Add "Arg" relationships
+ unsigned ArgIndex = 0;
+ for (const Use &U : I.operands()) {
+ unsigned OperandID = Vocabulary::getNumericID(U.get());
+ unsigned RelationID = ArgRelation + ArgIndex;
+ OS << Opcode << '\t' << OperandID << '\t' << RelationID << '\n';
+
+ LLVM_DEBUG({
+ StringRef OperandStr = Vocabulary::getVocabKeyForOperandKind(
+ Vocabulary::getOperandKind(U.get()));
+ dbgs() << Vocabulary::getVocabKeyForOpcode(Opcode + 1) << '\t'
+ << OperandStr << '\t' << "Arg" << ArgIndex << '\n';
+ });
+
+ ArgIndex++;
+ }
+ // Only update MaxRelation if there were operands
+ if (ArgIndex > 0) {
+ MaxRelation = std::max(MaxRelation, ArgRelation + ArgIndex - 1);
+ }
+ PrevOpcode = Opcode;
+ HasPrevOpcode = true;
+ }
+ }
- for (const BasicBlock &BB : F)
- traverseBasicBlock(BB, LocalOS);
+ return MaxRelation;
+ }
- LocalOS.flush();
- OS << LocalOutput;
+ /// Dump entity ID to string mappings
+ static void generateEntityMappings(raw_ostream &OS) {
+ // FIXME: Currently, the generated entity mappings are not one-to-one;
+ // Multiple TypeIDs map to same string key (Like Half, BFloat, etc. map to
+ // FloatTy). This would hinder learning good seed embeddings.
+ // We should fix this in the future by ensuring unique string keys either by
+ // post-processing here without changing the mapping in ir2vec::Vocabulary,
+ // or by changing the Vocabulary generation logic to ensure unique keys.
+ auto EntityLen = Vocabulary::expectedSize();
+ OS << EntityLen << "\n";
+ for (unsigned EntityID = 0; EntityID < EntityLen; ++EntityID)
+ OS << Vocabulary::getStringKey(EntityID) << '\t' << EntityID << '\n';
}
/// Generate embeddings for the entire module
@@ -198,27 +286,6 @@ class IR2VecTool {
}
}
}
-
-private:
- /// Process a single basic block for triplet generation
- void traverseBasicBlock(const BasicBlock &BB, raw_string_ostream &OS) const {
- // Consider only non-debug and non-pseudo instructions
- for (const auto &I : BB.instructionsWithoutDebug()) {
- StringRef OpcStr = Vocabulary::getVocabKeyForOpcode(I.getOpcode());
- StringRef TypeStr =
- Vocabulary::getVocabKeyForTypeID(I.getType()->getTypeID());
-
- OS << '\n' << OpcStr << ' ' << TypeStr << ' ';
-
- LLVM_DEBUG(I.print(dbgs()); dbgs() << "\n");
- LLVM_DEBUG(I.getType()->print(dbgs()); dbgs() << " Type\n");
-
- for (const Use &U : I.operands())
- OS << Vocabulary::getVocabKeyForOperandKind(
- Vocabulary::getOperandKind(U.get()))
- << ' ';
- }
- }
};
Error processModule(Module &M, raw_ostream &OS) {
@@ -246,18 +313,7 @@ Error processModule(Module &M, raw_ostream &OS) {
Tool.generateEmbeddings(OS);
}
} else {
- // Triplet generation mode - no vocabulary needed
- if (!FunctionName.empty())
- // Process single function
- if (const Function *F = M.getFunction(FunctionName))
- Tool.generateTriplets(*F, OS);
- else
- return createStringError(errc::invalid_argument,
- "Function '%s' not found",
- FunctionName.c_str());
- else
- // Process all functions
- Tool.generateTriplets(OS);
+ Tool.generateTriplets(OS);
}
return Error::success();
}
@@ -280,9 +336,32 @@ int main(int argc, char **argv) {
"See https://llvm.org/docs/CommandGuide/llvm-ir2vec.html for more "
"information.\n");
+ // Validate input file requirement
+ if (InputFilename.empty() && Mode != EntityMode) {
+ errs() << "Error: Input file (.bc/.ll) or stdin (-) is required\n";
+ return 1;
+ }
+
// Validate command line options
- if (Mode == TripletMode && Level.getNumOccurrences() > 0)
- errs() << "Warning: --level option is ignored in triplet mode\n";
+ if (Mode != EmbeddingMode) {
+ if (Level.getNumOccurrences() > 0)
+ errs() << "Warning: --level option is ignored\n";
+ if (FunctionName.getNumOccurrences() > 0)
+ errs() << "Warning: --function option is ignored\n";
+ }
+
+ std::error_code EC;
+ raw_fd_ostream OS(OutputFilename, EC);
+ if (EC) {
+ errs() << "Error opening output file: " << EC.message() << "\n";
+ return 1;
+ }
+
+ if (Mode == EntityMode) {
+ // Just dump entity mappings without processing any IR
+ IR2VecTool::generateEntityMappings(OS);
+ return 0;
+ }
...
[truncated]
|
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
from
42f9479 to
528ac7b
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-support-stdin-input-llvm-ir2vec
branch
from
6efc8a8 to
9e17794
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-support-stdin-input-llvm-ir2vec
branch
from
9e17794 to
1e22261
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
2 times, most recently
from
09d483a to
3f8c21f
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-support-stdin-input-llvm-ir2vec
branch
2 times, most recently
from
0903552 to
7fee589
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-support-stdin-input-llvm-ir2vec
branch
from
7fee589 to
36ecab5
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
2 times, most recently
from
ea01937 to
b209998
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-support-stdin-input-llvm-ir2vec
branch
from
32275ce to
a47b7f7
Compare
Base automatically changed from
users/svkeerthy/07-16-support-stdin-input-llvm-ir2vec
to
main
July 17, 2025 20:43
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
from
b209998 to
2f6e0b4
Compare
boomanaiden154
approved these changes
Jul 22, 2025
|
|
||
| // Validate input file requirement | ||
| if (InputFilename.empty() && Mode != EntityMode) { | ||
| errs() << "Error: Input file (.bc/.ll) or stdin (-) is required\n"; |
There was a problem hiding this comment.
You should probably have a test for this case.
There was a problem hiding this comment.
Seems that this check is not needed. When input files are not provided it "just waits" for input (like other llvm tools). So removed the check.
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
from
2f6e0b4 to
a1e45d7
Compare
mtrofin
reviewed
Jul 29, 2025
| @@ -0,0 +1,95 @@ | |||
| ; RUN: llvm-ir2vec --mode=entities | FileCheck %s | |||
|
|
|||
| CHECK: 92 | |||
There was a problem hiding this comment.
you could also have dumped the expected output to a reference file under Inputs and diff.
But I wouldn't bother doing that now, it's just more work, and my point was that there's a "less work" alternative - for next time.
mtrofin
reviewed
Jul 29, 2025
| clEnumValN(EmbeddingMode, "embeddings", | ||
| "Generate embeddings using trained vocabulary")), | ||
| cl::init(EmbeddingMode), cl::cat(IR2VecToolCategory)); | ||
| static cl::opt<ToolMode> Mode( |
There was a problem hiding this comment.
would cl::SubCommand work here? not for this change, but something to consider.
There was a problem hiding this comment.
Just read about it; using SubCommand here would help. Thanks the pointer! Will refactor the code in the next patch.
mtrofin
reviewed
Jul 29, 2025
mtrofin
approved these changes
Jul 29, 2025
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
2 times, most recently
from
3bfc9f5 to
e3f6820
Compare
Merge activity
|
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
2 times, most recently
from
0194360 to
b980082
Compare
|
@svkeerthy It looks like Graphite is stuck in a loop. You should probably cancel it and retry only merging this PR. See #151135 for more details. |
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
3 times, most recently
from
41f655c to
36543d6
Compare
svkeerthy
force-pushed
the
users/svkeerthy/07-16-revamp-triplet-gen
branch
from
36543d6 to
e609434
Compare
|
Umm.. Somehow it seem to have worked automatically! Is it just taking more time? |
Probably a gap in commit traffic assuming the hypothesis in #151135 is correct. Still best to not just let it automatically go through because you will burn a ton of CI time. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add entity mapping mode to llvm-ir2vec and improve triplet generation format for knowledge graph embedding training.
This change streamlines the workflow for training the vocabulary embeddings with IR2Vec by:
These improvements make IR2Vec more compatible with standard knowledge graph embedding training pipelines and reduce the preprocessing steps needed before training.
See #149215 for more details on how it is used.
(Tracking issues - #141817, #141834)