Exposing BB level representations

.github/workflows/test.yml

@@ -45,4 +45,4 @@ jobs:
     - name: make
       run: cd build && make -j8
     - name: Run-tests
-      run: cd build && ulimit -s unlimited && make check
+      run: cd build && ulimit -s unlimited && make check_ir2vec

README.md

@@ -104,7 +104,7 @@ If you're a C++ developer and require low-level control, optimization, or integr
 
 This process would generate `ir2vec` binary under `build/bin` directory, `libIR2Vec.a` and `libIR2Vec.so` under `build/lib` directory.
 
-To ensure the correctness, run `make check`
+To ensure the correctness, run `make check_ir2vec`
 
 
 

src/CMakeLists.txt

@@ -107,7 +107,7 @@ if(NOT LLVM_IR2VEC)
 
   add_subdirectory(test-suite)
 
-  add_custom_target(check
+  add_custom_target(check_ir2vec
     COMMAND python3 test-lit.py -a .
     COMMENT "Running LIT based test-suite"
     WORKING_DIRECTORY ./test-suite

src/FlowAware.cpp

@@ -31,6 +31,7 @@
 
 #include <functional>
 #include <regex>
+#include <string>
 
 using namespace llvm;
 using namespace IR2Vec;
@@ -534,7 +535,7 @@ Vector IR2Vec_FA::func2Vec(Function &F,
                        bbVector.begin(), std::plus<double>());
       }
     }
-
+    bbVecMap[b] = bbVector;
     IR2VEC_DEBUG(outs() << "-------------------------------------------\n");
     for (auto i : bbVector) {
       if ((i <= 0.0001 && i > 0) || (i < 0 && i >= -0.0001)) {

src/Symbolic.cpp

@@ -122,6 +122,7 @@ Vector IR2Vec_Symbolic::func2Vec(Function &F,
     std::transform(funcVector.begin(), funcVector.end(),
                    weightedBBVector.begin(), funcVector.begin(),
                    std::plus<double>());
+    bbVecMap[b] = weightedBBVector;
   }
 
   funcStack.pop_back();
@@ -130,6 +131,10 @@ Vector IR2Vec_Symbolic::func2Vec(Function &F,
 
 Vector IR2Vec_Symbolic::bb2Vec(BasicBlock &B,
                                SmallVector<Function *, 15> &funcStack) {
+  auto It = bbVecMap.find(&B);
+  if (It != bbVecMap.end()) {
+    return It->second;
+  }
   Vector bbVector(DIM, 0);
 
   for (auto &I : B) {

src/include/FlowAware.h

@@ -39,6 +39,7 @@ class IR2Vec_FA {
 
   llvm::SmallMapVector<const llvm::Instruction *, IR2Vec::Vector, 128>
       instVecMap;
+  llvm::SmallMapVector<const llvm::BasicBlock *, IR2Vec::Vector, 16> bbVecMap;
   llvm::SmallMapVector<const llvm::Function *, IR2Vec::Vector, 16> funcVecMap;
 
   llvm::SmallMapVector<const llvm::Function *,
@@ -183,6 +184,11 @@ class IR2Vec_FA {
     return instVecMap;
   }
 
+  llvm::SmallMapVector<const llvm::BasicBlock *, IR2Vec::Vector, 16>
+  getBBVecMap() {
+    return bbVecMap;
+  }
+
   llvm::SmallMapVector<const llvm::Function *, IR2Vec::Vector, 16>
   getFuncVecMap() {
     return funcVecMap;

src/include/IR2Vec.h

@@ -25,6 +25,7 @@ class Embeddings {
                         float WA = 0.2, float WT = 0.5);
 
   llvm::SmallMapVector<const llvm::Instruction *, Vector, 128> instVecMap;
+  llvm::SmallMapVector<const llvm::BasicBlock *, IR2Vec::Vector, 16> bbVecMap;
   llvm::SmallMapVector<const llvm::Function *, Vector, 16> funcVecMap;
   Vector pgmVector;
   std::map<std::string, IR2Vec::Vector> vocabulary;
@@ -56,6 +57,14 @@ class Embeddings {
     return instVecMap;
   }
 
+  // Returns a map containing basic block and the corresponding vector
+  // representations for a given module corresponding to the IR2VecMode and
+  // other configurations that is set in constructor
+  llvm::SmallMapVector<const llvm::BasicBlock *, IR2Vec::Vector, 16>
+  getBBVecMap() {
+    return bbVecMap;
+  }
+
   // Returns a map containing functions and the corresponding vector
   // representations for a given module corresponding to the IR2VecMode and
   // other configurations that is set in constructor

src/include/Symbolic.h

@@ -34,6 +34,7 @@ class IR2Vec_Symbolic {
                           llvm::SmallVector<llvm::Function *, 15> &funcStack);
   std::string res;
   llvm::SmallMapVector<const llvm::Function *, IR2Vec::Vector, 16> funcVecMap;
+  llvm::SmallMapVector<const llvm::BasicBlock *, IR2Vec::Vector, 16> bbVecMap;
   llvm::SmallMapVector<const llvm::Instruction *, IR2Vec::Vector, 128>
       instVecMap;
 
@@ -52,6 +53,11 @@ class IR2Vec_Symbolic {
     return instVecMap;
   }
 
+  llvm::SmallMapVector<const llvm::BasicBlock *, IR2Vec::Vector, 16>
+  getBBVecMap() {
+    return bbVecMap;
+  }
+
   llvm::SmallMapVector<const llvm::Function *, IR2Vec::Vector, 16>
   getFuncVecMap() {
     return funcVecMap;

src/libIR2Vec.cpp

@@ -34,22 +34,26 @@ int IR2Vec::Embeddings::generateEncodings(llvm::Module &M,
     FA.generateFlowAwareEncodingsForFunction(o, funcName);
     instVecMap = FA.getInstVecMap();
     funcVecMap = FA.getFuncVecMap();
+    bbVecMap = FA.getBBVecMap();
   } else if (mode == IR2Vec::IR2VecMode::FlowAware) {
     IR2Vec_FA FA(M, vocabulary);
     FA.generateFlowAwareEncodings(o);
     instVecMap = FA.getInstVecMap();
     funcVecMap = FA.getFuncVecMap();
+    bbVecMap = FA.getBBVecMap();
     pgmVector = FA.getProgramVector();
   } else if (mode == IR2Vec::IR2VecMode::Symbolic && !funcName.empty()) {
     IR2Vec_Symbolic SYM(M, vocabulary);
     SYM.generateSymbolicEncodingsForFunction(0, funcName);
     instVecMap = SYM.getInstVecMap();
     funcVecMap = SYM.getFuncVecMap();
+    bbVecMap = SYM.getBBVecMap();
   } else if (mode == IR2Vec::IR2VecMode::Symbolic) {
     IR2Vec_Symbolic SYM(M, vocabulary);
     SYM.generateSymbolicEncodings(o);
     instVecMap = SYM.getInstVecMap();
     funcVecMap = SYM.getFuncVecMap();
+    bbVecMap = SYM.getBBVecMap();
     pgmVector = SYM.getProgramVector();
   }