Adding IR2Vec as an analysis pass

llvm/docs/MLGO.rst

@@ -174,3 +174,151 @@ clang.
     TODO(mtrofin): 
         - logging, and the use in interactive mode.
         - discuss an example (like the inliner)
+
+IR2Vec Embeddings
+=================
+
+IR2Vec is a program embedding approach designed specifically for LLVM IR. It
+is implemented as a function analysis pass in LLVM. The IR2Vec embeddings
+capture syntactic, semantic, and structural properties of the IR through 
+learned representations. These representations are obtained as a JSON 
+vocabulary that maps the entities of the IR (opcodes, types, operands) to 
+n-dimensional floating point vectors (embeddings). 
+
+With IR2Vec, representation at different granularities of IR, such as
+instructions, functions, and basic blocks, can be obtained. Representations 
+of loops and regions can be derived from these representations, which can be
+useful in different scenarios. The representations can be useful for various
+downstream tasks, including ML-guided compiler optimizations.
+
+Currently, to use IR2Vec embeddings, the JSON vocabulary first needs to be read
+and used to obtain the vocabulary mapping. Then, use this mapping to
+derive the representations. In LLVM, this process is implemented using two
+independent passes: ``IR2VecVocabAnalysis`` and ``IR2VecAnalysis``. The former
+reads the JSON vocabulary and populates ``IR2VecVocabResult``, which is then used
+by ``IR2VecAnalysis``.
+
+It is recommended to run ``IR2VecVocabAnalysis`` once, as the
+vocabulary typically does not change. In the future, we plan
+to improve this process by automatically generating the vocabulary mappings
+during build time, eliminating the need for a separate file read.
+
+IR2VecAnalysis Usage
+--------------------
+
+To use IR2Vec in an LLVM-based tool or pass, interaction with the analysis 
+results can be done through the following APIs:
+
+1. **Including the Header:**
+
+   First, include the necessary header file in the source code:
+
+   .. code-block:: c++
+
+      #include "llvm/Analysis/IR2VecAnalysis.h"
+
+2. **Accessing the Analysis Results:**
+
+   To access the IR2Vec embeddings, obtain the ``IR2VecAnalysis``
+   result from the Function Analysis Manager (FAM).
+
+   .. code-block:: c++
+
+      llvm::FunctionAnalysisManager &FAM = ...; // The FAM instance
+      llvm::Function &F = ...; // The function to analyze
+      auto &IR2VecResult = FAM.getResult<llvm::IR2VecAnalysis>(F);
+
+3. **Checking for Valid Results:**
+
+   Ensure that the analysis result is valid before accessing the embeddings:
+
+   .. code-block:: c++
+
+      if (IR2VecResult.isValid()) {
+        // Proceed to access embeddings
+      }
+
+4. **Retrieving Embeddings:**
+
+   The ``IR2VecResult`` provides access to embeddings (currently) at three levels:
+
+   - **Instruction Embeddings:**
+
+     .. code-block:: c++
+
+        const auto &instVecMap = IR2VecResult.getInstVecMap();
+        // instVecMap is a SmallMapVector<const Instruction*, ir2vec::Embedding, 128>
+        for (const auto &it : instVecMap) {
+          const Instruction *I = it.first;
+          const ir2vec::Embedding &embedding = it.second;
+          // Use the instruction embedding
+        }
+   - **Basic Block Embeddings:**
+
+     .. code-block:: c++
+
+        const auto &bbVecMap = IR2VecResult.getBBVecMap();
+        // bbVecMap is a SmallMapVector<const BasicBlock*, ir2vec::Embedding, 16>
+        for (const auto &it : bbVecMap) {
+          const BasicBlock *BB = it.first;
+          const ir2vec::Embedding &embedding = it.second;
+          // Use the basic block embedding
+        }
+   - **Function Embedding:**
+
+     .. code-block:: c++
+
+        const ir2vec::Embedding &funcEmbedding = IR2VecResult.getFunctionVector();
+        // Use the function embedding
+
+5. **Working with Embeddings:**
+
+   Embeddings are represented as ``std::vector<double>``. These
+   vectors as features for machine learning models, compute similarity scores
+   between different code snippets, or perform other analyses as needed.
+
+Example Usage
+^^^^^^^^^^^^^
+
+.. code-block:: c++
+
+   #include "llvm/Analysis/IR2VecAnalysis.h"
+   #include "llvm/IR/Function.h"
+   #include "llvm/IR/Instructions.h"
+   #include "llvm/Passes/PassBuilder.h"
+
+   // ... other includes and code ...
+
+   void processFunction(llvm::Function &F, llvm::FunctionAnalysisManager &FAM) {
+     auto &IR2VecResult = FAM.getResult<llvm::IR2VecAnalysis>(F);
+
+     if (IR2VecResult.isValid()) {
+       const auto &instVecMap = IR2VecResult.getInstVecMap();
+       for (const auto &it : instVecMap) {
+         const Instruction *I = it.first;
+         const auto &embedding = it.second;
+         llvm::errs() << "Instruction: " << *I << "\n";
+         llvm::errs() << "Embedding: ";
+         for (double val : embedding) {
+           llvm::errs() << val << " ";
+         }
+         llvm::errs() << "\n";
+       }
+     } else {
+       llvm::errs() << "IR2Vec analysis failed for function " << F.getName() << "\n";
+     }
+   }
+
+   // ... rest of the pass ...
+
+   // In the pass's run method:
+   // processFunction(F, FAM);
+
+Further Details
+---------------
+
+For more detailed information about the IR2Vec algorithm, its parameters, and
+advanced usage, please refer to the original paper:
+`IR2Vec: LLVM IR Based Scalable Program Embeddings <https://doi.org/10.1145/3418463>`_.
+The LLVM source code for ``IR2VecAnalysis`` can also be explored to understand the 
+implementation details.

llvm/include/llvm/Analysis/IR2VecAnalysis.h

@@ -0,0 +1,128 @@
+//===- IR2VecAnalysis.h - IR2Vec Analysis Implementation -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions. See the LICENSE file for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains the declaration of IR2VecAnalysis that computes
+/// IR2Vec Embeddings of the program.
+///
+/// Program Embeddings are typically or derived-from a learned
+/// representation of the program. Such embeddings are used to represent the
+/// programs as input to machine learning algorithms. IR2Vec represents the
+/// LLVM IR as embeddings.
+///
+/// The IR2Vec algorithm is described in the following paper:
+///
+///   IR2Vec: LLVM IR Based Scalable Program Embeddings, S. VenkataKeerthy,
+///   Rohit Aggarwal, Shalini Jain, Maunendra Sankar Desarkar, Ramakrishna
+///   Upadrasta, and Y. N. Srikant, ACM Transactions on Architecture and
+///   Code Optimization (TACO), 2020. https://doi.org/10.1145/3418463.
+///   https://arxiv.org/abs/1909.06228
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_IR2VECANALYSIS_H
+#define LLVM_ANALYSIS_IR2VECANALYSIS_H
+
+#include "llvm/ADT/MapVector.h"
+#include "llvm/IR/PassManager.h"
+#include <map>
+
+namespace llvm {
+
+class Module;
+class BasicBlock;
+class Instruction;
+class Function;
+
+namespace ir2vec {
+using Embedding = std::vector<double>;
+// FIXME: Current the keys are strings. This can be changed to
+// use integers for cheaper lookups.
+using Vocab = std::map<std::string, Embedding>;
+} // namespace ir2vec
+
+class IR2VecVocabResult;
+class IR2VecResult;
+
+/// This analysis provides the vocabulary for IR2Vec. The vocabulary provides a
+/// mapping between an entity of the IR (like opcode, type, argument, etc.) and
+/// its corresponding embedding.
+class IR2VecVocabAnalysis : public AnalysisInfoMixin<IR2VecVocabAnalysis> {
+  ir2vec::Vocab Vocabulary;
+  Error readVocabulary();
+
+public:
+  static AnalysisKey Key;
+  IR2VecVocabAnalysis() = default;
+  using Result = IR2VecVocabResult;
+  Result run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+class IR2VecVocabResult {
+  ir2vec::Vocab Vocabulary;
+  bool Valid = false;
+
+public:
+  IR2VecVocabResult() = default;
+  IR2VecVocabResult(ir2vec::Vocab &&Vocabulary);
+
+  // Helper functions
+  bool isValid() const { return Valid; }
+  const ir2vec::Vocab &getVocabulary() const;
+  unsigned getDimension() const;
+  bool invalidate(Module &M, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &Inv);
+};
+
+class IR2VecResult {
+  SmallMapVector<const Instruction *, ir2vec::Embedding, 128> InstVecMap;
+  SmallMapVector<const BasicBlock *, ir2vec::Embedding, 16> BBVecMap;
+  ir2vec::Embedding FuncVector;
+  bool Valid = false;
+
+public:
+  IR2VecResult() = default;
+  IR2VecResult(
+      const SmallMapVector<const Instruction *, ir2vec::Embedding, 128>
+          &&InstMap,
+      const SmallMapVector<const BasicBlock *, ir2vec::Embedding, 16> &&BBMap,
+      const ir2vec::Embedding &&FuncVector);
+  bool isValid() const { return Valid; }
+
+  const SmallMapVector<const Instruction *, ir2vec::Embedding, 128> &
+  getInstVecMap() const;
+  const SmallMapVector<const BasicBlock *, ir2vec::Embedding, 16> &
+  getBBVecMap() const;
+  const ir2vec::Embedding &getFunctionVector() const;
+};
+
+/// This analysis provides the IR2Vec embeddings for instructions, basic blocks,
+/// and functions.
+class IR2VecAnalysis : public AnalysisInfoMixin<IR2VecAnalysis> {
+public:
+  IR2VecAnalysis() = default;
+  static AnalysisKey Key;
+  using Result = IR2VecResult;
+  Result run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+/// This pass prints the IR2Vec embeddings for instructions, basic blocks, and
+/// functions.
+class IR2VecPrinterPass : public PassInfoMixin<IR2VecPrinterPass> {
+  raw_ostream &OS;
+  void printVector(const ir2vec::Embedding &Vec) const;
+
+public:
+  explicit IR2VecPrinterPass(raw_ostream &OS) : OS(OS) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+  static bool isRequired() { return true; }
+};
+
+} // namespace llvm
+
+#endif // LLVM_ANALYSIS_IR2VECANALYSIS_H

llvm/lib/Analysis/CMakeLists.txt

@@ -78,6 +78,7 @@ add_llvm_component_library(LLVMAnalysis
   GlobalsModRef.cpp
   GuardUtils.cpp
   HeatUtils.cpp
+  IR2VecAnalysis.cpp
   IRSimilarityIdentifier.cpp
   IVDescriptors.cpp
   IVUsers.cpp