[llvm-branch-commits] [llvm] [IR2Vec] Add embeddings mode to llvm-ir2vec tool (PR #147844)

[Aiden Grossman via llvm-branch-commits]

================
@@ -50,16 +58,63 @@ static cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),
                                            cl::init("-"),
                                            cl::cat(IR2VecToolCategory));
 
+enum ToolMode {
+  TripletMode,  // Generate triplets 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<std::string>
+    FunctionName("function", cl::desc("Process specific function only"),
+                 cl::value_desc("name"), cl::Optional, cl::init(""),
+                 cl::cat(IR2VecToolCategory));
+
+enum EmbeddingLevel {
+  InstructionLevel, // Generate instruction-level embeddings
+  BasicBlockLevel,  // Generate basic block-level embeddings
+  FunctionLevel     // Generate function-level embeddings
+};
+
+static cl::opt<EmbeddingLevel>
+    Level("level", cl::desc("Embedding generation level (for embedding mode):"),
+          cl::values(clEnumValN(InstructionLevel, "inst",
+                                "Generate instruction-level embeddings"),
+                     clEnumValN(BasicBlockLevel, "bb",
+                                "Generate basic block-level embeddings"),
+                     clEnumValN(FunctionLevel, "func",
+                                "Generate function-level embeddings")),
+          cl::init(FunctionLevel), cl::cat(IR2VecToolCategory));
+
 namespace {
 
-/// Helper class for collecting IR information and generating triplets
+/// Helper class for collecting IR information and generating embeddings
----------------
boomanaiden154 wrote:

This comment seems incorrect given this can do embeddings and triplets?

https://github.com/llvm/llvm-project/pull/147844