MIR2Vec embedding

Author: svkeerthy

llvm/include/llvm/CodeGen/MIR2Vec.h

@@ -51,11 +51,21 @@ class LLVMContext;
 class MIR2VecVocabLegacyAnalysis;
 class TargetInstrInfo;
 
+enum class MIR2VecKind { Symbolic };
+
 namespace mir2vec {
+
+// Forward declarations
+class MIREmbedder;
+class SymbolicMIREmbedder;
+
 extern llvm::cl::OptionCategory MIR2VecCategory;
 extern cl::opt<float> OpcWeight;
 
 using Embedding = ir2vec::Embedding;
+using MachineInstEmbeddingsMap = DenseMap<const MachineInstr *, Embedding>;
+using MachineBlockEmbeddingsMap =
+    DenseMap<const MachineBasicBlock *, Embedding>;
 
 /// Class for storing and accessing the MIR2Vec vocabulary.
 /// The MIRVocabulary class manages seed embeddings for LLVM Machine IR
@@ -132,6 +142,79 @@ class MIRVocabulary {
     assert(isValid() && "Invalid vocabulary");
     return Storage.size();
   }
+
+  /// Create a dummy vocabulary for testing purposes.
+  static MIRVocabulary createDummyVocabForTest(const TargetInstrInfo &TII,
+                                               unsigned Dim = 1);
+};
+
+/// Base class for MIR embedders
+class MIREmbedder {
+protected:
+  const MachineFunction &MF;
+  const MIRVocabulary &Vocab;
+
+  /// Dimension of the embeddings; Captured from the vocabulary
+  const unsigned Dimension;
+
+  /// Weight for opcode embeddings
+  const float OpcWeight;
+
+  // Utility maps - these are used to store the vector representations of
+  // instructions, basic blocks and functions.
+  mutable Embedding MFuncVector;
+  mutable MachineBlockEmbeddingsMap MBBVecMap;
+  mutable MachineInstEmbeddingsMap MInstVecMap;
+
+  MIREmbedder(const MachineFunction &MF, const MIRVocabulary &Vocab);
+
+  /// Function to compute embeddings. It generates embeddings for all
+  /// the instructions and basic blocks in the function F.
+  void computeEmbeddings() const;
+
+  /// Function to compute the embedding for a given basic block.
+  /// Specific to the kind of embeddings being computed.
+  virtual void computeEmbeddings(const MachineBasicBlock &MBB) const = 0;
+
+public:
+  virtual ~MIREmbedder() = default;
+
+  /// Factory method to create an Embedder object of the specified kind
+  /// Returns nullptr if the requested kind is not supported.
+  static std::unique_ptr<MIREmbedder> create(MIR2VecKind Mode,
+                                             const MachineFunction &MF,
+                                             const MIRVocabulary &Vocab);
+
+  /// Returns a map containing machine instructions and the corresponding
+  /// embeddings for the machine function MF if it has been computed. If not, it
+  /// computes the embeddings for MF and returns the map.
+  const MachineInstEmbeddingsMap &getMInstVecMap() const;
+
+  /// Returns a map containing machine basic block and the corresponding
+  /// embeddings for the machine function MF if it has been computed. If not, it
+  /// computes the embeddings for MF and returns the map.
+  const MachineBlockEmbeddingsMap &getMBBVecMap() const;
+
+  /// Returns the embedding for a given machine basic block in the machine
+  /// function MF if it has been computed. If not, it computes the embedding for
+  /// MBB and returns it.
+  const Embedding &getMBBVector(const MachineBasicBlock &MBB) const;
+
+  /// Computes and returns the embedding for the current machine function.
+  const Embedding &getMFunctionVector() const;
+};
+
+/// Class for computing Symbolic embeddings
+/// Symbolic embeddings are constructed based on the entity-level
+/// representations obtained from the MIR Vocabulary.
+class SymbolicMIREmbedder : public MIREmbedder {
+private:
+  void computeEmbeddings(const MachineBasicBlock &MBB) const override;
+
+public:
+  SymbolicMIREmbedder(const MachineFunction &F, const MIRVocabulary &Vocab);
+  static std::unique_ptr<SymbolicMIREmbedder>
+  create(const MachineFunction &MF, const MIRVocabulary &Vocab);
 };
 
 } // namespace mir2vec
@@ -181,6 +264,31 @@ class MIR2VecVocabPrinterLegacyPass : public MachineFunctionPass {
   }
 };
 
+/// This pass prints the MIR2Vec embeddings for machine functions, basic blocks,
+/// and instructions
+class MIR2VecPrinterLegacyPass : public MachineFunctionPass {
+  raw_ostream &OS;
+
+public:
+  static char ID;
+  explicit MIR2VecPrinterLegacyPass(raw_ostream &OS)
+      : MachineFunctionPass(ID), OS(OS) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MIR2VecVocabLegacyAnalysis>();
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  StringRef getPassName() const override {
+    return "MIR2Vec Embedder Printer Pass";
+  }
+};
+
+/// Create a machine pass that prints MIR2Vec embeddings
+MachineFunctionPass *createMIR2VecPrinterLegacyPass(raw_ostream &OS);
+
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_MIR2VEC_H

llvm/include/llvm/CodeGen/Passes.h

@@ -93,6 +93,10 @@ createMachineFunctionPrinterPass(raw_ostream &OS,
 LLVM_ABI MachineFunctionPass *
 createMIR2VecVocabPrinterLegacyPass(raw_ostream &OS);
 
+/// MIR2VecPrinter pass - This pass prints out the MIR2Vec embeddings for
+/// machine functions, basic blocks and instructions.
+LLVM_ABI MachineFunctionPass *createMIR2VecPrinterLegacyPass(raw_ostream &OS);
+
 /// StackFramePrinter pass - This pass prints out the machine function's
 /// stack frame to the given stream as a debugging tool.
 LLVM_ABI MachineFunctionPass *createStackFrameLayoutAnalysisPass();

llvm/include/llvm/InitializePasses.h

@@ -222,6 +222,7 @@ LLVM_ABI void
 initializeMachineSanitizerBinaryMetadataLegacyPass(PassRegistry &);
 LLVM_ABI void initializeMIR2VecVocabLegacyAnalysisPass(PassRegistry &);
 LLVM_ABI void initializeMIR2VecVocabPrinterLegacyPassPass(PassRegistry &);
+LLVM_ABI void initializeMIR2VecPrinterLegacyPassPass(PassRegistry &);
 LLVM_ABI void initializeMachineSchedulerLegacyPass(PassRegistry &);
 LLVM_ABI void initializeMachineSinkingLegacyPass(PassRegistry &);
 LLVM_ABI void initializeMachineTraceMetricsWrapperPassPass(PassRegistry &);

llvm/lib/CodeGen/CodeGen.cpp

@@ -98,6 +98,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeMachineUniformityAnalysisPassPass(Registry);
   initializeMIR2VecVocabLegacyAnalysisPass(Registry);
   initializeMIR2VecVocabPrinterLegacyPassPass(Registry);
+  initializeMIR2VecPrinterLegacyPassPass(Registry);
   initializeMachineUniformityInfoPrinterPassPass(Registry);
   initializeMachineVerifierLegacyPassPass(Registry);
   initializeObjCARCContractLegacyPassPass(Registry);

llvm/lib/CodeGen/MIR2Vec.cpp

@@ -41,11 +41,18 @@ static cl::opt<std::string>
 cl::opt<float> OpcWeight("mir2vec-opc-weight", cl::Optional, cl::init(1.0),
                          cl::desc("Weight for machine opcode embeddings"),
                          cl::cat(MIR2VecCategory));
+cl::opt<MIR2VecKind> MIR2VecEmbeddingKind(
+    "mir2vec-kind", cl::Optional,
+    cl::values(clEnumValN(MIR2VecKind::Symbolic, "symbolic",
+                          "Generate symbolic embeddings for MIR")),
+    cl::init(MIR2VecKind::Symbolic), cl::desc("MIR2Vec embedding kind"),
+    cl::cat(MIR2VecCategory));
+
 } // namespace mir2vec
 } // namespace llvm
 
 //===----------------------------------------------------------------------===//
-// Vocabulary Implementation
+// Vocabulary
 //===----------------------------------------------------------------------===//
 
 MIRVocabulary::MIRVocabulary(VocabMap &&OpcodeEntries,
@@ -190,6 +197,29 @@ void MIRVocabulary::buildCanonicalOpcodeMapping() {
                     << " unique base opcodes\n");
 }
 
+MIRVocabulary MIRVocabulary::createDummyVocabForTest(const TargetInstrInfo &TII,
+                                                     unsigned Dim) {
+  assert(Dim > 0 && "Dimension must be greater than zero");
+
+  float DummyVal = 0.1f;
+
+  // Create a temporary vocabulary instance to build canonical mapping
+  MIRVocabulary TempVocab({}, &TII);
+  TempVocab.buildCanonicalOpcodeMapping();
+
+  // Create dummy embeddings for all canonical opcode names
+  VocabMap DummyVocabMap;
+  for (const auto &COpcodeName : TempVocab.UniqueBaseOpcodeNames) {
+    // Create dummy embedding filled with DummyVal
+    Embedding DummyEmbedding(Dim, DummyVal);
+    DummyVocabMap[COpcodeName] = DummyEmbedding;
+    DummyVal += 0.1f;
+  }
+
+  // Create and return vocabulary with dummy embeddings
+  return MIRVocabulary(std::move(DummyVocabMap), &TII);
+}
+
 //===----------------------------------------------------------------------===//
 // MIR2VecVocabLegacyAnalysis Implementation
 //===----------------------------------------------------------------------===//
@@ -267,7 +297,104 @@ MIR2VecVocabLegacyAnalysis::getMIR2VecVocabulary(const Module &M) {
 }
 
 //===----------------------------------------------------------------------===//
-// Printer Passes Implementation
+// MIREmbedder and its subclasses
+//===----------------------------------------------------------------------===//
+
+MIREmbedder::MIREmbedder(const MachineFunction &MF, const MIRVocabulary &Vocab)
+    : MF(MF), Vocab(Vocab), Dimension(Vocab.getDimension()),
+      OpcWeight(::OpcWeight), MFuncVector(Embedding(Dimension)) {}
+
+std::unique_ptr<MIREmbedder> MIREmbedder::create(MIR2VecKind Mode,
+                                                 const MachineFunction &MF,
+                                                 const MIRVocabulary &Vocab) {
+  switch (Mode) {
+  case MIR2VecKind::Symbolic:
+    return std::make_unique<SymbolicMIREmbedder>(MF, Vocab);
+  }
+  return nullptr;
+}
+
+const MachineInstEmbeddingsMap &MIREmbedder::getMInstVecMap() const {
+  if (MInstVecMap.empty())
+    computeEmbeddings();
+  return MInstVecMap;
+}
+
+const MachineBlockEmbeddingsMap &MIREmbedder::getMBBVecMap() const {
+  if (MBBVecMap.empty())
+    computeEmbeddings();
+  return MBBVecMap;
+}
+
+const Embedding &MIREmbedder::getMBBVector(const MachineBasicBlock &BB) const {
+  auto It = MBBVecMap.find(&BB);
+  if (It != MBBVecMap.end())
+    return It->second;
+  computeEmbeddings(BB);
+  return MBBVecMap[&BB];
+}
+
+const Embedding &MIREmbedder::getMFunctionVector() const {
+  // Currently, we always (re)compute the embeddings for the function.
+  // This is cheaper than caching the vector.
+  computeEmbeddings();
+  return MFuncVector;
+}
+
+void MIREmbedder::computeEmbeddings() const {
+  // Reset function vector to zero before recomputing
+  MFuncVector = Embedding(Dimension, 0.0);
+
+  // Consider all machine basic blocks in the function
+  for (const auto &MBB : MF) {
+    computeEmbeddings(MBB);
+    MFuncVector += MBBVecMap[&MBB];
+  }
+}
+
+SymbolicMIREmbedder::SymbolicMIREmbedder(const MachineFunction &MF,
+                                         const MIRVocabulary &Vocab)
+    : MIREmbedder(MF, Vocab) {}
+
+std::unique_ptr<SymbolicMIREmbedder>
+SymbolicMIREmbedder::create(const MachineFunction &MF,
+                            const MIRVocabulary &Vocab) {
+  return std::make_unique<SymbolicMIREmbedder>(MF, Vocab);
+}
+
+void SymbolicMIREmbedder::computeEmbeddings(
+    const MachineBasicBlock &MBB) const {
+  Embedding MBBVector(Dimension, 0);
+
+  // Get instruction info for opcode name resolution
+  const auto &Subtarget = MF.getSubtarget();
+  const auto *TII = Subtarget.getInstrInfo();
+  if (!TII) {
+    MF.getFunction().getContext().emitError(
+        "MIR2Vec: No TargetInstrInfo available; cannot compute embeddings");
+    return;
+  }
+
+  // Process each machine instruction in the basic block
+  for (const auto &MI : MBB) {
+    // Skip debug instructions and other metadata
+    if (MI.isDebugInstr())
+      continue;
+
+    // Todo: Add operand/argument contributions
+
+    // Store the instruction embedding
+    auto InstVector = Vocab[MI.getOpcode()];
+    MInstVecMap[&MI] = InstVector;
+    MBBVector += InstVector;
+  }
+
+  // Store the basic block embedding
+  MBBVecMap[&MBB] = MBBVector;
+}
+
+//===----------------------------------------------------------------------===//
+// Printer Passes
 //===----------------------------------------------------------------------===//
 
 char MIR2VecVocabPrinterLegacyPass::ID = 0;
@@ -304,3 +431,67 @@ MachineFunctionPass *
 llvm::createMIR2VecVocabPrinterLegacyPass(raw_ostream &OS) {
   return new MIR2VecVocabPrinterLegacyPass(OS);
 }
+
+char MIR2VecPrinterLegacyPass::ID = 0;
+INITIALIZE_PASS_BEGIN(MIR2VecPrinterLegacyPass, "print-mir2vec",
+                      "MIR2Vec Embedder Printer Pass", false, true)
+INITIALIZE_PASS_DEPENDENCY(MIR2VecVocabLegacyAnalysis)
+INITIALIZE_PASS_DEPENDENCY(MachineModuleInfoWrapperPass)
+INITIALIZE_PASS_END(MIR2VecPrinterLegacyPass, "print-mir2vec",
+                    "MIR2Vec Embedder Printer Pass", false, true)
+
+bool MIR2VecPrinterLegacyPass::runOnMachineFunction(MachineFunction &MF) {
+  auto &Analysis = getAnalysis<MIR2VecVocabLegacyAnalysis>();
+  auto MIRVocab = Analysis.getMIR2VecVocabulary(*MF.getFunction().getParent());
+
+  if (!MIRVocab.isValid()) {
+    OS << "MIR2Vec Embedder Printer: Invalid vocabulary for function "
+       << MF.getName() << "\n";
+    return false;
+  }
+
+  auto Emb = mir2vec::MIREmbedder::create(MIR2VecEmbeddingKind, MF, MIRVocab);
+  if (!Emb) {
+    OS << "Error creating MIR2Vec embeddings for function " << MF.getName()
+       << "\n";
+    return false;
+  }
+
+  OS << "MIR2Vec embeddings for machine function " << MF.getName() << ":\n";
+  OS << "Machine Function vector: ";
+  Emb->getMFunctionVector().print(OS);
+
+  OS << "Machine basic block vectors:\n";
+  const auto &MBBMap = Emb->getMBBVecMap();
+  for (const MachineBasicBlock &MBB : MF) {
+    auto It = MBBMap.find(&MBB);
+    if (It != MBBMap.end()) {
+      OS << "Machine basic block: " << MBB.getFullName() << ":\n";
+      It->second.print(OS);
+    }
+  }
+
+  OS << "Machine instruction vectors:\n";
+  const auto &MInstMap = Emb->getMInstVecMap();
+  for (const MachineBasicBlock &MBB : MF) {
+    for (const MachineInstr &MI : MBB) {
+      // Skip debug instructions as they are not
+      // embedded
+      if (MI.isDebugInstr())
+        continue;
+
+      auto It = MInstMap.find(&MI);
+      if (It != MInstMap.end()) {
+        OS << "Machine instruction: ";
+        MI.print(OS);
+        It->second.print(OS);
+      }
+    }
+  }
+
+  return false;
+}
+
+MachineFunctionPass *llvm::createMIR2VecPrinterLegacyPass(raw_ostream &OS) {
+  return new MIR2VecPrinterLegacyPass(OS);
+}