[MLGO][IR2Vec] Integrating IR2Vec with MLInliner (#143479)

Changes to use Symbolic embeddings in MLInliner. 

(Fixes #141836, Tracking issue - #141817)

1 files changed, 131 insertions, 14 deletions

diff --git a/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp b/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp
index 0720d93..e50486bc 100644
--- a/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp
+++ b/llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp
@@ -8,6 +8,7 @@
 
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/IR2Vec.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
@@ -20,15 +21,20 @@
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <cstring>
 
 using namespace llvm;
+using namespace testing;
 
 namespace llvm {
 LLVM_ABI extern cl::opt<bool> EnableDetailedFunctionProperties;
 LLVM_ABI extern cl::opt<bool> BigBasicBlockInstructionThreshold;
 LLVM_ABI extern cl::opt<bool> MediumBasicBlockInstrutionThreshold;
+LLVM_ABI extern cl::opt<float> ir2vec::OpcWeight;
+LLVM_ABI extern cl::opt<float> ir2vec::TypeWeight;
+LLVM_ABI extern cl::opt<float> ir2vec::ArgWeight;
 } // namespace llvm
 
 namespace {
@@ -36,17 +42,65 @@ namespace {
 class FunctionPropertiesAnalysisTest : public testing::Test {
 public:
   FunctionPropertiesAnalysisTest() {
+    createTestVocabulary(1);
+    MAM.registerPass([&] { return IR2VecVocabAnalysis(Vocabulary); });
+    MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+    FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
     FAM.registerPass([&] { return DominatorTreeAnalysis(); });
     FAM.registerPass([&] { return LoopAnalysis(); });
     FAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+
+    ir2vec::OpcWeight = 1.0;
+    ir2vec::TypeWeight = 1.0;
+    ir2vec::ArgWeight = 1.0;
+  }
+
+private:
+  float OriginalOpcWeight = ir2vec::OpcWeight;
+  float OriginalTypeWeight = ir2vec::TypeWeight;
+  float OriginalArgWeight = ir2vec::ArgWeight;
+
+  void createTestVocabulary(unsigned Dim) {
+    llvm::SmallVector<std::string, 32> SampleEntities = {
+        "add",        "sub",      "mul",        "icmp",          "br",
+        "ret",        "store",    "load",       "alloca",        "phi",
+        "call",       "voidTy",   "floatTy",    "integerTy",     "functionTy",
+        "structTy",   "arrayTy",  "pointerTy",  "vectorTy",      "emptyTy",
+        "labelTy",    "tokenTy",  "metadataTy", "unknownTy",     "function",
+        "pointer",    "constant", "variable",   "getelementptr", "invoke",
+        "landingpad", "resume",   "catch",      "cleanup"};
+    float EmbVal = 0.1;
+
+    // Helper lambda to add entries to the vocabulary
+    auto addEntry = [&](std::string key) {
+      Vocabulary[key] = ir2vec::Embedding(Dim, EmbVal);
+      EmbVal += 0.1;
+    };
+
+    for (auto &Name : SampleEntities)
+      addEntry(Name);
+    return;
   }
 
 protected:
   std::unique_ptr<DominatorTree> DT;
   std::unique_ptr<LoopInfo> LI;
   FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
+  ir2vec::Vocab Vocabulary;
+
+  void TearDown() override {
+    // Restore original IR2Vec weights
+    ir2vec::OpcWeight = OriginalOpcWeight;
+    ir2vec::TypeWeight = OriginalTypeWeight;
+    ir2vec::ArgWeight = OriginalArgWeight;
+  }
 
   FunctionPropertiesInfo buildFPI(Function &F) {
+    // FunctionPropertiesInfo assumes IR2VecVocabAnalysis has been run to
+    // use IR2Vec.
+    auto VocabResult = MAM.getResult<IR2VecVocabAnalysis>(*F.getParent());
+    (void)VocabResult;
     return FunctionPropertiesInfo::getFunctionPropertiesInfo(F, FAM);
   }
 
@@ -62,15 +116,22 @@ protected:
       Err.print("MLAnalysisTests", errs());
     return Mod;
   }
-  
-  CallBase* findCall(Function& F, const char* Name = nullptr) {
+
+  CallBase *findCall(Function &F, const char *Name = nullptr) {
     for (auto &BB : F)
-      for (auto &I : BB )
+      for (auto &I : BB)
         if (auto *CB = dyn_cast<CallBase>(&I))
           if (!Name || CB->getName() == Name)
             return CB;
     return nullptr;
   }
+
+  std::unique_ptr<ir2vec::Embedder> createEmbedder(const Function &F) {
+    auto EmbResult =
+        ir2vec::Embedder::create(IR2VecKind::Symbolic, F, Vocabulary);
+    EXPECT_TRUE(static_cast<bool>(EmbResult));
+    return std::move(*EmbResult);
+  }
 };
 
 TEST_F(FunctionPropertiesAnalysisTest, BasicTest) {
@@ -113,6 +174,8 @@ define internal i32 @top() {
   EXPECT_EQ(BranchesFeatures.StoreInstCount, 0);
   EXPECT_EQ(BranchesFeatures.MaxLoopDepth, 0);
   EXPECT_EQ(BranchesFeatures.TopLevelLoopCount, 0);
+  EXPECT_TRUE(BranchesFeatures.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*BranchesFunction)->getFunctionVector()));
 
   Function *TopFunction = M->getFunction("top");
   FunctionPropertiesInfo TopFeatures = buildFPI(*TopFunction);
@@ -120,6 +183,8 @@ define internal i32 @top() {
   EXPECT_EQ(TopFeatures.BlocksReachedFromConditionalInstruction, 0);
   EXPECT_EQ(TopFeatures.Uses, 0);
   EXPECT_EQ(TopFeatures.DirectCallsToDefinedFunctions, 1);
+  EXPECT_TRUE(TopFeatures.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*TopFunction)->getFunctionVector()));
   EXPECT_EQ(BranchesFeatures.LoadInstCount, 0);
   EXPECT_EQ(BranchesFeatures.StoreInstCount, 0);
   EXPECT_EQ(BranchesFeatures.MaxLoopDepth, 0);
@@ -159,6 +224,9 @@ define internal i32 @top() {
   EXPECT_EQ(DetailedBranchesFeatures.CallReturnsPointerCount, 0);
   EXPECT_EQ(DetailedBranchesFeatures.CallWithManyArgumentsCount, 0);
   EXPECT_EQ(DetailedBranchesFeatures.CallWithPointerArgumentCount, 0);
+  EXPECT_TRUE(
+      DetailedBranchesFeatures.getFunctionEmbedding().approximatelyEquals(
+          createEmbedder(*BranchesFunction)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }
 
@@ -210,6 +278,8 @@ finally:
   EXPECT_EQ(DetailedF1Properties.CallReturnsPointerCount, 0);
   EXPECT_EQ(DetailedF1Properties.CallWithManyArgumentsCount, 0);
   EXPECT_EQ(DetailedF1Properties.CallWithPointerArgumentCount, 0);
+  EXPECT_TRUE(DetailedF1Properties.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }
 
@@ -232,7 +302,7 @@ define i32 @f2(i32 %a) {
 )IR");
 
   Function *F1 = M->getFunction("f1");
-  CallBase* CB = findCall(*F1, "b");
+  CallBase *CB = findCall(*F1, "b");
   EXPECT_NE(CB, nullptr);
 
   FunctionPropertiesInfo ExpectedInitial;
@@ -240,6 +310,8 @@ define i32 @f2(i32 %a) {
   ExpectedInitial.TotalInstructionCount = 3;
   ExpectedInitial.Uses = 1;
   ExpectedInitial.DirectCallsToDefinedFunctions = 1;
+  ExpectedInitial.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
 
   FunctionPropertiesInfo ExpectedFinal = ExpectedInitial;
   ExpectedFinal.DirectCallsToDefinedFunctions = 0;
@@ -252,6 +324,9 @@ define i32 @f2(i32 %a) {
   auto IR = llvm::InlineFunction(*CB, IFI);
   EXPECT_TRUE(IR.isSuccess());
   invalidate(*F1);
+  ExpectedFinal.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
+
   EXPECT_TRUE(FPU.finishAndTest(FAM));
   EXPECT_EQ(FPI, ExpectedFinal);
 }
@@ -285,7 +360,7 @@ define i32 @f2(i32 %a) {
 )IR");
 
   Function *F1 = M->getFunction("f1");
-  CallBase* CB = findCall(*F1, "b");
+  CallBase *CB = findCall(*F1, "b");
   EXPECT_NE(CB, nullptr);
 
   FunctionPropertiesInfo ExpectedInitial;
@@ -294,6 +369,8 @@ define i32 @f2(i32 %a) {
   ExpectedInitial.TotalInstructionCount = 9;
   ExpectedInitial.Uses = 1;
   ExpectedInitial.DirectCallsToDefinedFunctions = 1;
+  ExpectedInitial.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
 
   FunctionPropertiesInfo ExpectedFinal = ExpectedInitial;
   ExpectedFinal.DirectCallsToDefinedFunctions = 0;
@@ -307,6 +384,9 @@ define i32 @f2(i32 %a) {
   EXPECT_TRUE(IR.isSuccess());
   invalidate(*F1);
   EXPECT_TRUE(FPU.finishAndTest(FAM));
+
+  ExpectedFinal.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
   EXPECT_EQ(FPI, ExpectedFinal);
 }
 
@@ -347,7 +427,7 @@ exit:
 )IR");
 
   Function *F1 = M->getFunction("f1");
-  CallBase* CB = findCall(*F1, "b");
+  CallBase *CB = findCall(*F1, "b");
   EXPECT_NE(CB, nullptr);
 
   FunctionPropertiesInfo ExpectedInitial;
@@ -356,6 +436,8 @@ exit:
   ExpectedInitial.TotalInstructionCount = 9;
   ExpectedInitial.Uses = 1;
   ExpectedInitial.DirectCallsToDefinedFunctions = 1;
+  ExpectedInitial.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
 
   FunctionPropertiesInfo ExpectedFinal;
   ExpectedFinal.BasicBlockCount = 6;
@@ -374,6 +456,9 @@ exit:
   EXPECT_TRUE(IR.isSuccess());
   invalidate(*F1);
   EXPECT_TRUE(FPU.finishAndTest(FAM));
+
+  ExpectedFinal.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
   EXPECT_EQ(FPI, ExpectedFinal);
 }
 
@@ -409,7 +494,7 @@ declare i32 @__gxx_personality_v0(...)
 )IR");
 
   Function *F1 = M->getFunction("caller");
-  CallBase* CB = findCall(*F1);
+  CallBase *CB = findCall(*F1);
   EXPECT_NE(CB, nullptr);
 
   auto FPI = buildFPI(*F1);
@@ -422,6 +507,8 @@ declare i32 @__gxx_personality_v0(...)
   EXPECT_EQ(static_cast<size_t>(FPI.BasicBlockCount), F1->size());
   EXPECT_EQ(static_cast<size_t>(FPI.TotalInstructionCount),
             F1->getInstructionCount());
+  EXPECT_TRUE(FPI.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
 }
 
 TEST_F(FunctionPropertiesAnalysisTest, InvokeUnreachableHandler) {
@@ -462,7 +549,7 @@ declare i32 @__gxx_personality_v0(...)
 )IR");
 
   Function *F1 = M->getFunction("caller");
-  CallBase* CB = findCall(*F1);
+  CallBase *CB = findCall(*F1);
   EXPECT_NE(CB, nullptr);
 
   auto FPI = buildFPI(*F1);
@@ -475,6 +562,8 @@ declare i32 @__gxx_personality_v0(...)
   EXPECT_EQ(static_cast<size_t>(FPI.BasicBlockCount), F1->size() - 1);
   EXPECT_EQ(static_cast<size_t>(FPI.TotalInstructionCount),
             F1->getInstructionCount() - 2);
+  EXPECT_TRUE(FPI.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EXPECT_EQ(FPI, FunctionPropertiesInfo::getFunctionPropertiesInfo(*F1, FAM));
 }
 
@@ -516,7 +605,7 @@ declare i32 @__gxx_personality_v0(...)
 )IR");
 
   Function *F1 = M->getFunction("caller");
-  CallBase* CB = findCall(*F1);
+  CallBase *CB = findCall(*F1);
   EXPECT_NE(CB, nullptr);
 
   auto FPI = buildFPI(*F1);
@@ -568,7 +657,7 @@ lpad:
 )IR");
 
   Function *F1 = M->getFunction("outer");
-  CallBase* CB = findCall(*F1);
+  CallBase *CB = findCall(*F1);
   EXPECT_NE(CB, nullptr);
 
   auto FPI = buildFPI(*F1);
@@ -581,6 +670,8 @@ lpad:
   EXPECT_EQ(static_cast<size_t>(FPI.BasicBlockCount), F1->size() - 1);
   EXPECT_EQ(static_cast<size_t>(FPI.TotalInstructionCount),
             F1->getInstructionCount() - 2);
+  EXPECT_TRUE(FPI.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EXPECT_EQ(FPI, FunctionPropertiesInfo::getFunctionPropertiesInfo(*F1, FAM));
 }
 
@@ -624,7 +715,7 @@ lpad:
 )IR");
 
   Function *F1 = M->getFunction("outer");
-  CallBase* CB = findCall(*F1);
+  CallBase *CB = findCall(*F1);
   EXPECT_NE(CB, nullptr);
 
   auto FPI = buildFPI(*F1);
@@ -637,6 +728,8 @@ lpad:
   EXPECT_EQ(static_cast<size_t>(FPI.BasicBlockCount), F1->size() - 1);
   EXPECT_EQ(static_cast<size_t>(FPI.TotalInstructionCount),
             F1->getInstructionCount() - 2);
+  EXPECT_TRUE(FPI.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EXPECT_EQ(FPI, FunctionPropertiesInfo::getFunctionPropertiesInfo(*F1, FAM));
 }
 
@@ -689,6 +782,8 @@ end:
   ExpectedInitial.DirectCallsToDefinedFunctions = 1;
   ExpectedInitial.MaxLoopDepth = 1;
   ExpectedInitial.TopLevelLoopCount = 1;
+  ExpectedInitial.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
 
   FunctionPropertiesInfo ExpectedFinal = ExpectedInitial;
   ExpectedFinal.BasicBlockCount = 6;
@@ -705,6 +800,9 @@ end:
   EXPECT_TRUE(IR.isSuccess());
   invalidate(*F1);
   EXPECT_TRUE(FPU.finishAndTest(FAM));
+
+  ExpectedFinal.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
   EXPECT_EQ(FPI, ExpectedFinal);
 }
 
@@ -733,7 +831,7 @@ extra:
 extra2:
   br label %cond.end
 
-cond.end:                                         ; preds = %cond.false, %cond.true
+cond.end:                                         ; preds = %extra2, %cond.true
   %cond = phi i64 [ %conv2, %cond.true ], [ %call3, %extra ]
   ret i64 %cond
 }
@@ -757,7 +855,9 @@ declare void @llvm.trap()
   ExpectedInitial.BlocksReachedFromConditionalInstruction = 2;
   ExpectedInitial.Uses = 1;
   ExpectedInitial.DirectCallsToDefinedFunctions = 1;
-  
+  ExpectedInitial.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
+
   FunctionPropertiesInfo ExpectedFinal = ExpectedInitial;
   ExpectedFinal.BasicBlockCount = 4;
   ExpectedFinal.DirectCallsToDefinedFunctions = 0;
@@ -772,6 +872,9 @@ declare void @llvm.trap()
   EXPECT_TRUE(IR.isSuccess());
   invalidate(*F1);
   EXPECT_TRUE(FPU.finishAndTest(FAM));
+
+  ExpectedFinal.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
   EXPECT_EQ(FPI, ExpectedFinal);
 }
 
@@ -817,6 +920,8 @@ declare void @f3()
   ExpectedInitial.BlocksReachedFromConditionalInstruction = 0;
   ExpectedInitial.Uses = 1;
   ExpectedInitial.DirectCallsToDefinedFunctions = 1;
+  ExpectedInitial.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
 
   FunctionPropertiesInfo ExpectedFinal = ExpectedInitial;
   ExpectedFinal.BasicBlockCount = 6;
@@ -832,6 +937,9 @@ declare void @f3()
   EXPECT_TRUE(IR.isSuccess());
   invalidate(*F1);
   EXPECT_TRUE(FPU.finishAndTest(FAM));
+
+  ExpectedFinal.setFunctionEmbeddingForTest(
+      createEmbedder(*F1)->getFunctionVector());
   EXPECT_EQ(FPI, ExpectedFinal);
 }
 
@@ -885,6 +993,8 @@ define i64 @f1(i64 %e) {
   EXPECT_EQ(DetailedF1Properties.CallReturnsPointerCount, 0);
   EXPECT_EQ(DetailedF1Properties.CallWithManyArgumentsCount, 0);
   EXPECT_EQ(DetailedF1Properties.CallWithPointerArgumentCount, 0);
+  EXPECT_TRUE(DetailedF1Properties.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }
 
@@ -910,6 +1020,8 @@ declare float @llvm.cos.f32(float)
   EXPECT_EQ(DetailedF1Properties.CallReturnsPointerCount, 0);
   EXPECT_EQ(DetailedF1Properties.CallWithManyArgumentsCount, 0);
   EXPECT_EQ(DetailedF1Properties.CallWithPointerArgumentCount, 0);
+  EXPECT_TRUE(DetailedF1Properties.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }
 
@@ -943,6 +1055,8 @@ declare float @f5()
   EXPECT_EQ(DetailedF1Properties.CallReturnsPointerCount, 1);
   EXPECT_EQ(DetailedF1Properties.CallWithManyArgumentsCount, 1);
   EXPECT_EQ(DetailedF1Properties.CallWithPointerArgumentCount, 1);
+  EXPECT_TRUE(DetailedF1Properties.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }
 
@@ -972,10 +1086,11 @@ BottomBlock2:
   EnableDetailedFunctionProperties.setValue(true);
   FunctionPropertiesInfo DetailedF1Properties = buildFPI(*F1);
   EXPECT_EQ(DetailedF1Properties.CriticalEdgeCount, 1);
+  EXPECT_TRUE(DetailedF1Properties.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }
 
-
 TEST_F(FunctionPropertiesAnalysisTest, FunctionReturnVectors) {
   LLVMContext C;
   std::unique_ptr<Module> M = makeLLVMModule(C,
@@ -998,6 +1113,8 @@ declare <4 x ptr> @f4()
   EXPECT_EQ(DetailedF1Properties.CallReturnsVectorIntCount, 1);
   EXPECT_EQ(DetailedF1Properties.CallReturnsVectorFloatCount, 1);
   EXPECT_EQ(DetailedF1Properties.CallReturnsVectorPointerCount, 1);
+  EXPECT_TRUE(DetailedF1Properties.getFunctionEmbedding().approximatelyEquals(
+      createEmbedder(*F1)->getFunctionVector()));
   EnableDetailedFunctionProperties.setValue(false);
 }