983 lines · cpp
1//===- IR2VecTest.cpp - Unit tests for IR2Vec -----------------------------==//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/Analysis/IR2Vec.h"10#include "llvm/IR/Constants.h"11#include "llvm/IR/Instruction.h"12#include "llvm/IR/Instructions.h"13#include "llvm/IR/LLVMContext.h"14#include "llvm/IR/Module.h"15#include "llvm/IR/Type.h"16#include "llvm/Support/Error.h"17#include "llvm/Support/JSON.h"18 19#include "gmock/gmock.h"20#include "gtest/gtest.h"21#include <vector>22 23using namespace llvm;24using namespace ir2vec;25using namespace ::testing;26 27namespace {28 29class TestableEmbedder : public Embedder {30public:31 TestableEmbedder(const Function &F, const Vocabulary &V) : Embedder(F, V) {}32 Embedding computeEmbeddings(const Instruction &I) const override {33 return Embedding();34 }35};36 37TEST(EmbeddingTest, ConstructorsAndAccessors) {38 // Default constructor39 {40 Embedding E;41 EXPECT_TRUE(E.empty());42 EXPECT_EQ(E.size(), 0u);43 }44 45 // Constructor with const std::vector<double>&46 {47 std::vector<double> Data = {1.0, 2.0, 3.0};48 Embedding E(Data);49 EXPECT_FALSE(E.empty());50 ASSERT_THAT(E, SizeIs(3u));51 EXPECT_THAT(E.getData(), ElementsAre(1.0, 2.0, 3.0));52 EXPECT_EQ(E[0], 1.0);53 EXPECT_EQ(E[1], 2.0);54 EXPECT_EQ(E[2], 3.0);55 }56 57 // Constructor with std::vector<double>&&58 {59 Embedding E(std::vector<double>({4.0, 5.0}));60 ASSERT_THAT(E, SizeIs(2u));61 EXPECT_THAT(E.getData(), ElementsAre(4.0, 5.0));62 }63 64 // Constructor with std::initializer_list<double>65 {66 Embedding E({6.0, 7.0, 8.0, 9.0});67 ASSERT_THAT(E, SizeIs(4u));68 EXPECT_THAT(E.getData(), ElementsAre(6.0, 7.0, 8.0, 9.0));69 EXPECT_EQ(E[0], 6.0);70 E[0] = 6.5;71 EXPECT_EQ(E[0], 6.5);72 }73 74 // Constructor with size_t75 {76 Embedding E(5);77 ASSERT_THAT(E, SizeIs(5u));78 EXPECT_THAT(E.getData(), ElementsAre(0.0, 0.0, 0.0, 0.0, 0.0));79 }80 81 // Constructor with size_t and double82 {83 Embedding E(5, 1.5);84 ASSERT_THAT(E, SizeIs(5u));85 EXPECT_THAT(E.getData(), ElementsAre(1.5, 1.5, 1.5, 1.5, 1.5));86 }87 88 // Test iterators89 {90 Embedding E({6.5, 7.0, 8.0, 9.0});91 std::vector<double> VecE;92 for (double Val : E) {93 VecE.push_back(Val);94 }95 EXPECT_THAT(VecE, ElementsAre(6.5, 7.0, 8.0, 9.0));96 97 const Embedding CE = E;98 std::vector<double> VecCE;99 for (const double &Val : CE) {100 VecCE.push_back(Val);101 }102 EXPECT_THAT(VecCE, ElementsAre(6.5, 7.0, 8.0, 9.0));103 104 EXPECT_EQ(*E.begin(), 6.5);105 EXPECT_EQ(*(E.end() - 1), 9.0);106 EXPECT_EQ(*CE.cbegin(), 6.5);107 EXPECT_EQ(*(CE.cend() - 1), 9.0);108 }109}110 111TEST(EmbeddingTest, AddVectorsOutOfPlace) {112 Embedding E1 = {1.0, 2.0, 3.0};113 Embedding E2 = {0.5, 1.5, -1.0};114 115 Embedding E3 = E1 + E2;116 EXPECT_THAT(E3, ElementsAre(1.5, 3.5, 2.0));117 118 // Check that E1 and E2 are unchanged119 EXPECT_THAT(E1, ElementsAre(1.0, 2.0, 3.0));120 EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));121}122 123TEST(EmbeddingTest, AddVectors) {124 Embedding E1 = {1.0, 2.0, 3.0};125 Embedding E2 = {0.5, 1.5, -1.0};126 127 E1 += E2;128 EXPECT_THAT(E1, ElementsAre(1.5, 3.5, 2.0));129 130 // Check that E2 is unchanged131 EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));132}133 134TEST(EmbeddingTest, SubtractVectorsOutOfPlace) {135 Embedding E1 = {1.0, 2.0, 3.0};136 Embedding E2 = {0.5, 1.5, -1.0};137 138 Embedding E3 = E1 - E2;139 EXPECT_THAT(E3, ElementsAre(0.5, 0.5, 4.0));140 141 // Check that E1 and E2 are unchanged142 EXPECT_THAT(E1, ElementsAre(1.0, 2.0, 3.0));143 EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));144}145 146TEST(EmbeddingTest, SubtractVectors) {147 Embedding E1 = {1.0, 2.0, 3.0};148 Embedding E2 = {0.5, 1.5, -1.0};149 150 E1 -= E2;151 EXPECT_THAT(E1, ElementsAre(0.5, 0.5, 4.0));152 153 // Check that E2 is unchanged154 EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));155}156 157TEST(EmbeddingTest, ScaleVector) {158 Embedding E1 = {1.0, 2.0, 3.0};159 E1 *= 0.5f;160 EXPECT_THAT(E1, ElementsAre(0.5, 1.0, 1.5));161}162 163TEST(EmbeddingTest, ScaleVectorOutOfPlace) {164 Embedding E1 = {1.0, 2.0, 3.0};165 Embedding E2 = E1 * 0.5f;166 EXPECT_THAT(E2, ElementsAre(0.5, 1.0, 1.5));167 168 // Check that E1 is unchanged169 EXPECT_THAT(E1, ElementsAre(1.0, 2.0, 3.0));170}171 172TEST(EmbeddingTest, AddScaledVector) {173 Embedding E1 = {1.0, 2.0, 3.0};174 Embedding E2 = {2.0, 0.5, -1.0};175 176 E1.scaleAndAdd(E2, 0.5f);177 EXPECT_THAT(E1, ElementsAre(2.0, 2.25, 2.5));178 179 // Check that E2 is unchanged180 EXPECT_THAT(E2, ElementsAre(2.0, 0.5, -1.0));181}182 183TEST(EmbeddingTest, ApproximatelyEqual) {184 Embedding E1 = {1.0, 2.0, 3.0};185 Embedding E2 = {1.0000001, 2.0000001, 3.0000001};186 EXPECT_TRUE(E1.approximatelyEquals(E2)); // Diff = 1e-7187 188 Embedding E3 = {1.00002, 2.00002, 3.00002}; // Diff = 2e-5189 EXPECT_FALSE(E1.approximatelyEquals(E3, 1e-6));190 EXPECT_TRUE(E1.approximatelyEquals(E3, 3e-5));191 192 Embedding E_clearly_within = {1.0000005, 2.0000005, 3.0000005}; // Diff = 5e-7193 EXPECT_TRUE(E1.approximatelyEquals(E_clearly_within));194 195 Embedding E_clearly_outside = {1.00001, 2.00001, 3.00001}; // Diff = 1e-5196 EXPECT_FALSE(E1.approximatelyEquals(E_clearly_outside, 1e-6));197 198 Embedding E4 = {1.0, 2.0, 3.5}; // Large diff199 EXPECT_FALSE(E1.approximatelyEquals(E4, 0.01));200 201 Embedding E5 = {1.0, 2.0, 3.0};202 EXPECT_TRUE(E1.approximatelyEquals(E5, 0.0));203 EXPECT_TRUE(E1.approximatelyEquals(E5));204}205 206#if GTEST_HAS_DEATH_TEST207#ifndef NDEBUG208TEST(EmbeddingTest, AccessOutOfBounds) {209 Embedding E = {1.0, 2.0, 3.0};210 EXPECT_DEATH(E[3], "Index out of bounds");211 EXPECT_DEATH(E[-1], "Index out of bounds");212 EXPECT_DEATH(E[4] = 4.0, "Index out of bounds");213}214 215TEST(EmbeddingTest, MismatchedDimensionsAddVectorsOutOfPlace) {216 Embedding E1 = {1.0, 2.0};217 Embedding E2 = {1.0};218 EXPECT_DEATH(E1 + E2, "Vectors must have the same dimension");219}220 221TEST(EmbeddingTest, MismatchedDimensionsAddVectors) {222 Embedding E1 = {1.0, 2.0};223 Embedding E2 = {1.0};224 EXPECT_DEATH(E1 += E2, "Vectors must have the same dimension");225}226 227TEST(EmbeddingTest, MismatchedDimensionsSubtractVectors) {228 Embedding E1 = {1.0, 2.0};229 Embedding E2 = {1.0};230 EXPECT_DEATH(E1 -= E2, "Vectors must have the same dimension");231}232 233TEST(EmbeddingTest, MismatchedDimensionsAddScaledVector) {234 Embedding E1 = {1.0, 2.0};235 Embedding E2 = {1.0};236 EXPECT_DEATH(E1.scaleAndAdd(E2, 1.0f),237 "Vectors must have the same dimension");238}239 240TEST(EmbeddingTest, MismatchedDimensionsApproximatelyEqual) {241 Embedding E1 = {1.0, 2.0};242 Embedding E2 = {1.010};243 EXPECT_DEATH(E1.approximatelyEquals(E2),244 "Vectors must have the same dimension");245}246#endif // NDEBUG247#endif // GTEST_HAS_DEATH_TEST248 249TEST(IR2VecTest, CreateSymbolicEmbedder) {250 Vocabulary V = Vocabulary(Vocabulary::createDummyVocabForTest());251 252 LLVMContext Ctx;253 Module M("M", Ctx);254 FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), false);255 Function *F = Function::Create(FTy, Function::ExternalLinkage, "f", M);256 257 auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, V);258 EXPECT_NE(Emb, nullptr);259}260 261TEST(IR2VecTest, CreateFlowAwareEmbedder) {262 Vocabulary V = Vocabulary(Vocabulary::createDummyVocabForTest());263 264 LLVMContext Ctx;265 Module M("M", Ctx);266 FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), false);267 Function *F = Function::Create(FTy, Function::ExternalLinkage, "f", M);268 269 auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, V);270 EXPECT_NE(Emb, nullptr);271}272 273TEST(IR2VecTest, CreateInvalidMode) {274 Vocabulary V = Vocabulary(Vocabulary::createDummyVocabForTest());275 276 LLVMContext Ctx;277 Module M("M", Ctx);278 FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), false);279 Function *F = Function::Create(FTy, Function::ExternalLinkage, "f", M);280 281 // static_cast an invalid int to IR2VecKind282 auto Result = Embedder::create(static_cast<IR2VecKind>(-1), *F, V);283 EXPECT_FALSE(static_cast<bool>(Result));284}285 286TEST(IR2VecTest, ZeroDimensionEmbedding) {287 Embedding E1;288 Embedding E2;289 // Should be no-op, but not crash290 E1 += E2;291 E1 -= E2;292 E1.scaleAndAdd(E2, 1.0f);293 EXPECT_TRUE(E1.empty());294}295 296// Fixture for IR2Vec tests requiring IR setup.297class IR2VecTestFixture : public ::testing::Test {298protected:299 std::unique_ptr<Vocabulary> V;300 LLVMContext Ctx;301 std::unique_ptr<Module> M;302 Function *F = nullptr;303 BasicBlock *BB = nullptr;304 Instruction *AddInst = nullptr;305 Instruction *RetInst = nullptr;306 307 void SetUp() override {308 V = std::make_unique<Vocabulary>(Vocabulary::createDummyVocabForTest(2));309 310 // Setup IR311 M = std::make_unique<Module>("TestM", Ctx);312 FunctionType *FTy = FunctionType::get(313 Type::getInt32Ty(Ctx), {Type::getInt32Ty(Ctx), Type::getInt32Ty(Ctx)},314 false);315 F = Function::Create(FTy, Function::ExternalLinkage, "f", M.get());316 BB = BasicBlock::Create(Ctx, "entry", F);317 Argument *Arg = F->getArg(0);318 llvm::Value *Const = ConstantInt::get(Type::getInt32Ty(Ctx), 42);319 320 AddInst = BinaryOperator::CreateAdd(Arg, Const, "add", BB);321 RetInst = ReturnInst::Create(Ctx, AddInst, BB);322 }323};324 325TEST_F(IR2VecTestFixture, GetInstVec_Symbolic) {326 auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);327 ASSERT_TRUE(static_cast<bool>(Emb));328 329 const auto &AddEmb = Emb->getInstVector(*AddInst);330 const auto &RetEmb = Emb->getInstVector(*RetInst);331 EXPECT_EQ(AddEmb.size(), 2u);332 EXPECT_EQ(RetEmb.size(), 2u);333 334 EXPECT_TRUE(AddEmb.approximatelyEquals(Embedding(2, 25.5)));335 EXPECT_TRUE(RetEmb.approximatelyEquals(Embedding(2, 15.5)));336}337 338TEST_F(IR2VecTestFixture, GetInstVec_FlowAware) {339 auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, *V);340 ASSERT_TRUE(static_cast<bool>(Emb));341 342 const auto &AddEmb = Emb->getInstVector(*AddInst);343 const auto &RetEmb = Emb->getInstVector(*RetInst);344 EXPECT_EQ(AddEmb.size(), 2u);345 EXPECT_EQ(RetEmb.size(), 2u);346 347 EXPECT_TRUE(AddEmb.approximatelyEquals(Embedding(2, 25.5)));348 EXPECT_TRUE(RetEmb.approximatelyEquals(Embedding(2, 32.6)));349}350 351TEST_F(IR2VecTestFixture, GetBBVector_Symbolic) {352 auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);353 ASSERT_TRUE(static_cast<bool>(Emb));354 355 const auto &BBVec = Emb->getBBVector(*BB);356 357 EXPECT_EQ(BBVec.size(), 2u);358 // BB vector should be sum of add and ret: {25.5, 25.5} + {15.5, 15.5} =359 // {41.0, 41.0}360 EXPECT_TRUE(BBVec.approximatelyEquals(Embedding(2, 41.0)));361}362 363TEST_F(IR2VecTestFixture, GetBBVector_FlowAware) {364 auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, *V);365 ASSERT_TRUE(static_cast<bool>(Emb));366 367 const auto &BBVec = Emb->getBBVector(*BB);368 369 EXPECT_EQ(BBVec.size(), 2u);370 // BB vector should be sum of add and ret: {25.5, 25.5} + {32.6, 32.6} =371 // {58.1, 58.1}372 EXPECT_TRUE(BBVec.approximatelyEquals(Embedding(2, 58.1)));373}374 375TEST_F(IR2VecTestFixture, GetFunctionVector_Symbolic) {376 auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);377 ASSERT_TRUE(static_cast<bool>(Emb));378 379 const auto &FuncVec = Emb->getFunctionVector();380 381 EXPECT_EQ(FuncVec.size(), 2u);382 383 // Function vector should match BB vector (only one BB): {41.0, 41.0}384 EXPECT_TRUE(FuncVec.approximatelyEquals(Embedding(2, 41.0)));385}386 387TEST_F(IR2VecTestFixture, GetFunctionVector_FlowAware) {388 auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, *V);389 ASSERT_TRUE(static_cast<bool>(Emb));390 391 const auto &FuncVec = Emb->getFunctionVector();392 393 EXPECT_EQ(FuncVec.size(), 2u);394 // Function vector should match BB vector (only one BB): {58.1, 58.1}395 EXPECT_TRUE(FuncVec.approximatelyEquals(Embedding(2, 58.1)));396}397 398TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_Symbolic) {399 auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);400 ASSERT_TRUE(static_cast<bool>(Emb));401 402 // Get initial function vector403 const auto &FuncVec1 = Emb->getFunctionVector();404 405 // Compute embeddings again by calling getFunctionVector multiple times406 const auto &FuncVec2 = Emb->getFunctionVector();407 const auto &FuncVec3 = Emb->getFunctionVector();408 409 // All function vectors should be identical410 EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec2));411 EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec3));412 EXPECT_TRUE(FuncVec2.approximatelyEquals(FuncVec3));413 414 Emb->invalidateEmbeddings();415 const auto &FuncVec4 = Emb->getFunctionVector();416 EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec4));417}418 419TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_FlowAware) {420 auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, *V);421 ASSERT_TRUE(static_cast<bool>(Emb));422 423 // Get initial function vector424 const auto &FuncVec1 = Emb->getFunctionVector();425 426 // Compute embeddings again by calling getFunctionVector multiple times427 const auto &FuncVec2 = Emb->getFunctionVector();428 const auto &FuncVec3 = Emb->getFunctionVector();429 430 // All function vectors should be identical431 EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec2));432 EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec3));433 EXPECT_TRUE(FuncVec2.approximatelyEquals(FuncVec3));434 435 Emb->invalidateEmbeddings();436 const auto &FuncVec4 = Emb->getFunctionVector();437 EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec4));438}439 440static constexpr unsigned MaxOpcodes = Vocabulary::MaxOpcodes;441[[maybe_unused]]442static constexpr unsigned MaxTypeIDs = Vocabulary::MaxTypeIDs;443static constexpr unsigned MaxCanonicalTypeIDs = Vocabulary::MaxCanonicalTypeIDs;444static constexpr unsigned MaxOperands = Vocabulary::MaxOperandKinds;445static constexpr unsigned MaxPredicateKinds = Vocabulary::MaxPredicateKinds;446 447// Mapping between LLVM Type::TypeID tokens and Vocabulary::CanonicalTypeID448// names and their canonical string keys.449#define IR2VEC_HANDLE_TYPE_BIMAP(X) \450 X(VoidTyID, VoidTy, "VoidTy") \451 X(IntegerTyID, IntegerTy, "IntegerTy") \452 X(FloatTyID, FloatTy, "FloatTy") \453 X(PointerTyID, PointerTy, "PointerTy") \454 X(FunctionTyID, FunctionTy, "FunctionTy") \455 X(StructTyID, StructTy, "StructTy") \456 X(ArrayTyID, ArrayTy, "ArrayTy") \457 X(FixedVectorTyID, VectorTy, "VectorTy") \458 X(LabelTyID, LabelTy, "LabelTy") \459 X(TokenTyID, TokenTy, "TokenTy") \460 X(MetadataTyID, MetadataTy, "MetadataTy")461 462TEST(IR2VecVocabularyTest, DummyVocabTest) {463 for (unsigned Dim = 1; Dim <= 10; ++Dim) {464 auto VocabVec = Vocabulary::createDummyVocabForTest(Dim);465 auto VocabVecSize = VocabVec.size();466 // All embeddings should have the same dimension467 for (const auto &Emb : VocabVec)468 EXPECT_EQ(Emb.size(), Dim);469 470 // Should have the correct total number of embeddings471 EXPECT_EQ(VocabVecSize, MaxOpcodes + MaxCanonicalTypeIDs + MaxOperands +472 MaxPredicateKinds);473 474 // Collect embeddings for later comparison before moving VocabVec475 std::vector<Embedding> ExpectedVocab;476 for (const auto &Emb : VocabVec)477 ExpectedVocab.push_back(Emb);478 479 IR2VecVocabAnalysis VocabAnalysis(std::move(VocabVec));480 LLVMContext TestCtx;481 Module TestMod("TestModuleForVocabAnalysis", TestCtx);482 ModuleAnalysisManager MAM;483 Vocabulary Result = VocabAnalysis.run(TestMod, MAM);484 EXPECT_TRUE(Result.isValid());485 EXPECT_EQ(Result.getDimension(), Dim);486 EXPECT_EQ(Result.getCanonicalSize(), VocabVecSize);487 488 unsigned CurPos = 0;489 for (const auto &Entry : Result)490 EXPECT_TRUE(Entry.approximatelyEquals(ExpectedVocab[CurPos++], 0.01));491 }492}493 494TEST(IR2VecVocabularyTest, SlotIdxMapping) {495 // Test getIndex for Opcodes496#define EXPECT_OPCODE_SLOT(NUM, OPCODE, CLASS) \497 EXPECT_EQ(Vocabulary::getIndex(NUM), static_cast<unsigned>(NUM - 1));498#define HANDLE_INST(NUM, OPCODE, CLASS) EXPECT_OPCODE_SLOT(NUM, OPCODE, CLASS)499#include "llvm/IR/Instruction.def"500#undef HANDLE_INST501#undef EXPECT_OPCODE_SLOT502 503 // Test getIndex for Types504#define EXPECT_TYPE_SLOT(TypeIDTok, CanonEnum, CanonStr) \505 EXPECT_EQ(Vocabulary::getIndex(Type::TypeIDTok), \506 MaxOpcodes + static_cast<unsigned>( \507 Vocabulary::CanonicalTypeID::CanonEnum));508 509 IR2VEC_HANDLE_TYPE_BIMAP(EXPECT_TYPE_SLOT)510 511#undef EXPECT_TYPE_SLOT512 513 // Test getIndex for Value operands514 LLVMContext Ctx;515 Module M("TestM", Ctx);516 FunctionType *FTy =517 FunctionType::get(Type::getVoidTy(Ctx), {Type::getInt32Ty(Ctx)}, false);518 Function *F = Function::Create(FTy, Function::ExternalLinkage, "testFunc", M);519 520#define EXPECTED_VOCAB_OPERAND_SLOT(X) \521 MaxOpcodes + MaxCanonicalTypeIDs + static_cast<unsigned>(X)522 // Test Function operand523 EXPECT_EQ(Vocabulary::getIndex(*F),524 EXPECTED_VOCAB_OPERAND_SLOT(Vocabulary::OperandKind::FunctionID));525 526 // Test Constant operand527 Constant *C = ConstantInt::get(Type::getInt32Ty(Ctx), 42);528 EXPECT_EQ(Vocabulary::getIndex(*C),529 EXPECTED_VOCAB_OPERAND_SLOT(Vocabulary::OperandKind::ConstantID));530 531 // Test Pointer operand532 BasicBlock *BB = BasicBlock::Create(Ctx, "entry", F);533 AllocaInst *PtrVal = new AllocaInst(Type::getInt32Ty(Ctx), 0, "ptr", BB);534 EXPECT_EQ(Vocabulary::getIndex(*PtrVal),535 EXPECTED_VOCAB_OPERAND_SLOT(Vocabulary::OperandKind::PointerID));536 537 // Test Variable operand (function argument)538 Argument *Arg = F->getArg(0);539 EXPECT_EQ(Vocabulary::getIndex(*Arg),540 EXPECTED_VOCAB_OPERAND_SLOT(Vocabulary::OperandKind::VariableID));541#undef EXPECTED_VOCAB_OPERAND_SLOT542 543 // Test getIndex for predicates544#define EXPECTED_VOCAB_PREDICATE_SLOT(X) \545 MaxOpcodes + MaxCanonicalTypeIDs + MaxOperands + static_cast<unsigned>(X)546 for (unsigned P = CmpInst::FIRST_FCMP_PREDICATE;547 P <= CmpInst::LAST_FCMP_PREDICATE; ++P) {548 CmpInst::Predicate Pred = static_cast<CmpInst::Predicate>(P);549 unsigned ExpectedIdx =550 EXPECTED_VOCAB_PREDICATE_SLOT((P - CmpInst::FIRST_FCMP_PREDICATE));551 EXPECT_EQ(Vocabulary::getIndex(Pred), ExpectedIdx);552 }553 auto ICMP_Start = CmpInst::LAST_FCMP_PREDICATE + 1;554 for (unsigned P = CmpInst::FIRST_ICMP_PREDICATE;555 P <= CmpInst::LAST_ICMP_PREDICATE; ++P) {556 CmpInst::Predicate Pred = static_cast<CmpInst::Predicate>(P);557 unsigned ExpectedIdx = EXPECTED_VOCAB_PREDICATE_SLOT(558 ICMP_Start + P - CmpInst::FIRST_ICMP_PREDICATE);559 EXPECT_EQ(Vocabulary::getIndex(Pred), ExpectedIdx);560 }561#undef EXPECTED_VOCAB_PREDICATE_SLOT562}563 564#if GTEST_HAS_DEATH_TEST565#ifndef NDEBUG566TEST(IR2VecVocabularyTest, NumericIDMapInvalidInputs) {567 // Test invalid opcode IDs568 EXPECT_DEATH(Vocabulary::getIndex(0u), "Invalid opcode");569 EXPECT_DEATH(Vocabulary::getIndex(MaxOpcodes + 1), "Invalid opcode");570 571 // Test invalid type IDs572 EXPECT_DEATH(Vocabulary::getIndex(static_cast<Type::TypeID>(MaxTypeIDs)),573 "Invalid type ID");574 EXPECT_DEATH(Vocabulary::getIndex(static_cast<Type::TypeID>(MaxTypeIDs + 10)),575 "Invalid type ID");576}577#endif // NDEBUG578#endif // GTEST_HAS_DEATH_TEST579 580TEST(IR2VecVocabularyTest, StringKeyGeneration) {581 EXPECT_EQ(Vocabulary::getStringKey(0), "Ret");582 EXPECT_EQ(Vocabulary::getStringKey(12), "Add");583 584#define EXPECT_OPCODE(NUM, OPCODE, CLASS) \585 EXPECT_EQ(Vocabulary::getStringKey(Vocabulary::getIndex(NUM)), \586 Vocabulary::getVocabKeyForOpcode(NUM));587#define HANDLE_INST(NUM, OPCODE, CLASS) EXPECT_OPCODE(NUM, OPCODE, CLASS)588#include "llvm/IR/Instruction.def"589#undef HANDLE_INST590#undef EXPECT_OPCODE591 592 // Verify CanonicalTypeID -> string mapping593#define EXPECT_CANONICAL_TYPE_NAME(TypeIDTok, CanonEnum, CanonStr) \594 EXPECT_EQ(Vocabulary::getStringKey( \595 MaxOpcodes + static_cast<unsigned>( \596 Vocabulary::CanonicalTypeID::CanonEnum)), \597 CanonStr);598 599 IR2VEC_HANDLE_TYPE_BIMAP(EXPECT_CANONICAL_TYPE_NAME)600 601#undef EXPECT_CANONICAL_TYPE_NAME602 603 // Verify OperandKind -> string mapping604#define HANDLE_OPERAND_KINDS(X) \605 X(FunctionID, "Function") \606 X(PointerID, "Pointer") \607 X(ConstantID, "Constant") \608 X(VariableID, "Variable")609 610#define EXPECT_OPERAND_KIND(EnumName, Str) \611 EXPECT_EQ(Vocabulary::getStringKey( \612 MaxOpcodes + MaxCanonicalTypeIDs + \613 static_cast<unsigned>(Vocabulary::OperandKind::EnumName)), \614 Str);615 616 HANDLE_OPERAND_KINDS(EXPECT_OPERAND_KIND)617 618#undef EXPECT_OPERAND_KIND619#undef HANDLE_OPERAND_KINDS620 621 StringRef FuncArgKey =622 Vocabulary::getStringKey(MaxOpcodes + MaxCanonicalTypeIDs + 0);623 StringRef PtrArgKey =624 Vocabulary::getStringKey(MaxOpcodes + MaxCanonicalTypeIDs + 1);625 EXPECT_EQ(FuncArgKey, "Function");626 EXPECT_EQ(PtrArgKey, "Pointer");627 628// Verify PredicateKind -> string mapping629#define EXPECT_PREDICATE_KIND(PredNum, PredPos, PredKind) \630 do { \631 std::string PredStr = \632 std::string(PredKind) + "_" + \633 CmpInst::getPredicateName(static_cast<CmpInst::Predicate>(PredNum)) \634 .str(); \635 unsigned Pos = MaxOpcodes + MaxCanonicalTypeIDs + MaxOperands + PredPos; \636 EXPECT_EQ(Vocabulary::getStringKey(Pos), PredStr); \637 } while (0)638 639 for (unsigned P = CmpInst::FIRST_FCMP_PREDICATE;640 P <= CmpInst::LAST_FCMP_PREDICATE; ++P)641 EXPECT_PREDICATE_KIND(P, P - CmpInst::FIRST_FCMP_PREDICATE, "FCMP");642 643 auto ICMP_Pos = CmpInst::LAST_FCMP_PREDICATE + 1;644 for (unsigned P = CmpInst::FIRST_ICMP_PREDICATE;645 P <= CmpInst::LAST_ICMP_PREDICATE; ++P)646 EXPECT_PREDICATE_KIND(P, ICMP_Pos++, "ICMP");647 648#undef EXPECT_PREDICATE_KIND649}650 651TEST(IR2VecVocabularyTest, VocabularyDimensions) {652 {653 Vocabulary V(Vocabulary::createDummyVocabForTest(1));654 EXPECT_TRUE(V.isValid());655 EXPECT_EQ(V.getDimension(), 1u);656 }657 658 {659 Vocabulary V(Vocabulary::createDummyVocabForTest(5));660 EXPECT_TRUE(V.isValid());661 EXPECT_EQ(V.getDimension(), 5u);662 }663 664 {665 Vocabulary V(Vocabulary::createDummyVocabForTest(10));666 EXPECT_TRUE(V.isValid());667 EXPECT_EQ(V.getDimension(), 10u);668 }669}670 671#if GTEST_HAS_DEATH_TEST672#ifndef NDEBUG673TEST(IR2VecVocabularyTest, InvalidAccess) {674 Vocabulary V(Vocabulary::createDummyVocabForTest(2));675 676 EXPECT_DEATH(V[0u], "Invalid opcode");677 678 EXPECT_DEATH(V[100u], "Invalid opcode");679}680#endif // NDEBUG681#endif // GTEST_HAS_DEATH_TEST682 683TEST(IR2VecVocabularyTest, TypeIDStringKeyMapping) {684 Vocabulary V = Vocabulary(Vocabulary::createDummyVocabForTest());685#define EXPECT_TYPE_TO_CANONICAL(TypeIDTok, CanonEnum, CanonStr) \686 do { \687 unsigned FlatIdx = V.getIndex(Type::TypeIDTok); \688 EXPECT_EQ(Vocabulary::getStringKey(FlatIdx), CanonStr); \689 } while (0);690 691 IR2VEC_HANDLE_TYPE_BIMAP(EXPECT_TYPE_TO_CANONICAL)692 693#undef EXPECT_TYPE_TO_CANONICAL694}695 696TEST(IR2VecVocabularyTest, InvalidVocabularyConstruction) {697 // Test 1: Create invalid VocabStorage with insufficient sections698 std::vector<std::vector<Embedding>> InvalidSectionData;699 // Only add one section with 2 embeddings, but the vocabulary needs 4 sections700 std::vector<Embedding> Section1;701 Section1.push_back(Embedding(2, 1.0));702 Section1.push_back(Embedding(2, 2.0));703 InvalidSectionData.push_back(std::move(Section1));704 705 VocabStorage InvalidStorage(std::move(InvalidSectionData));706 Vocabulary V(std::move(InvalidStorage));707 EXPECT_FALSE(V.isValid());708 709 {710 // Test 2: Default-constructed vocabulary should be invalid711 Vocabulary InvalidResult;712 EXPECT_FALSE(InvalidResult.isValid());713#if GTEST_HAS_DEATH_TEST714#ifndef NDEBUG715 EXPECT_DEATH(InvalidResult.getDimension(), "IR2Vec Vocabulary is invalid");716#endif // NDEBUG717#endif // GTEST_HAS_DEATH_TEST718 }719}720 721TEST(VocabStorageTest, DefaultConstructor) {722 VocabStorage storage;723 724 EXPECT_EQ(storage.size(), 0u);725 EXPECT_EQ(storage.getNumSections(), 0u);726 EXPECT_EQ(storage.getDimension(), 0u);727 EXPECT_FALSE(storage.isValid());728 729 // Test iterators on empty storage730 EXPECT_EQ(storage.begin(), storage.end());731}732 733TEST(VocabStorageTest, BasicConstruction) {734 // Create test data with 3 sections735 std::vector<std::vector<Embedding>> sectionData;736 737 // Section 0: 2 embeddings of dimension 3738 std::vector<Embedding> section0;739 section0.emplace_back(std::vector<double>{1.0, 2.0, 3.0});740 section0.emplace_back(std::vector<double>{4.0, 5.0, 6.0});741 sectionData.push_back(std::move(section0));742 743 // Section 1: 1 embedding of dimension 3744 std::vector<Embedding> section1;745 section1.emplace_back(std::vector<double>{7.0, 8.0, 9.0});746 sectionData.push_back(std::move(section1));747 748 // Section 2: 3 embeddings of dimension 3749 std::vector<Embedding> section2;750 section2.emplace_back(std::vector<double>{10.0, 11.0, 12.0});751 section2.emplace_back(std::vector<double>{13.0, 14.0, 15.0});752 section2.emplace_back(std::vector<double>{16.0, 17.0, 18.0});753 sectionData.push_back(std::move(section2));754 755 VocabStorage storage(std::move(sectionData));756 757 EXPECT_EQ(storage.size(), 6u); // Total: 2 + 1 + 3 = 6758 EXPECT_EQ(storage.getNumSections(), 3u);759 EXPECT_EQ(storage.getDimension(), 3u);760 EXPECT_TRUE(storage.isValid());761}762 763TEST(VocabStorageTest, SectionAccess) {764 // Create test data765 std::vector<std::vector<Embedding>> sectionData;766 767 std::vector<Embedding> section0;768 section0.emplace_back(std::vector<double>{1.0, 2.0});769 section0.emplace_back(std::vector<double>{3.0, 4.0});770 sectionData.push_back(std::move(section0));771 772 std::vector<Embedding> section1;773 section1.emplace_back(std::vector<double>{5.0, 6.0});774 sectionData.push_back(std::move(section1));775 776 VocabStorage storage(std::move(sectionData));777 778 // Test section access779 EXPECT_EQ(storage[0].size(), 2u);780 EXPECT_EQ(storage[1].size(), 1u);781 782 // Test embedding values783 EXPECT_THAT(storage[0][0].getData(), ElementsAre(1.0, 2.0));784 EXPECT_THAT(storage[0][1].getData(), ElementsAre(3.0, 4.0));785 EXPECT_THAT(storage[1][0].getData(), ElementsAre(5.0, 6.0));786}787 788#if GTEST_HAS_DEATH_TEST789#ifndef NDEBUG790TEST(VocabStorageTest, InvalidSectionAccess) {791 std::vector<std::vector<Embedding>> sectionData;792 std::vector<Embedding> section0;793 section0.emplace_back(std::vector<double>{1.0, 2.0});794 sectionData.push_back(std::move(section0));795 796 VocabStorage storage(std::move(sectionData));797 798 EXPECT_DEATH(storage[1], "Invalid section ID");799 EXPECT_DEATH(storage[10], "Invalid section ID");800}801 802TEST(VocabStorageTest, EmptySection) {803 std::vector<std::vector<Embedding>> sectionData;804 std::vector<Embedding> emptySection; // Empty section805 sectionData.push_back(std::move(emptySection));806 807 std::vector<Embedding> validSection;808 validSection.emplace_back(std::vector<double>{1.0});809 sectionData.push_back(std::move(validSection));810 811 EXPECT_DEATH(VocabStorage(std::move(sectionData)),812 "Vocabulary section is empty");813}814 815TEST(VocabStorageTest, EmptyMiddleSection) {816 std::vector<std::vector<Embedding>> sectionData;817 818 // Valid first section819 std::vector<Embedding> validSection1;820 validSection1.emplace_back(std::vector<double>{1.0});821 sectionData.push_back(std::move(validSection1));822 823 // Empty middle section824 std::vector<Embedding> emptySection;825 sectionData.push_back(std::move(emptySection));826 827 // Valid last section828 std::vector<Embedding> validSection2;829 validSection2.emplace_back(std::vector<double>{2.0});830 sectionData.push_back(std::move(validSection2));831 832 EXPECT_DEATH(VocabStorage(std::move(sectionData)),833 "Vocabulary section is empty");834}835 836TEST(VocabStorageTest, NoSections) {837 std::vector<std::vector<Embedding>> sectionData; // No sections838 839 EXPECT_DEATH(VocabStorage(std::move(sectionData)),840 "Vocabulary has no sections");841}842 843TEST(VocabStorageTest, MismatchedDimensionsAcrossSections) {844 std::vector<std::vector<Embedding>> sectionData;845 846 // Section 0: embeddings with dimension 2847 std::vector<Embedding> section0;848 section0.emplace_back(std::vector<double>{1.0, 2.0});849 section0.emplace_back(std::vector<double>{3.0, 4.0});850 sectionData.push_back(std::move(section0));851 852 // Section 1: embedding with dimension 3 (mismatch!)853 std::vector<Embedding> section1;854 section1.emplace_back(std::vector<double>{5.0, 6.0, 7.0});855 sectionData.push_back(std::move(section1));856 857 EXPECT_DEATH(VocabStorage(std::move(sectionData)),858 "All embeddings must have the same dimension");859}860 861TEST(VocabStorageTest, MismatchedDimensionsWithinSection) {862 std::vector<std::vector<Embedding>> sectionData;863 864 // Section 0: first embedding with dimension 2, second with dimension 3865 std::vector<Embedding> section0;866 section0.emplace_back(std::vector<double>{1.0, 2.0});867 section0.emplace_back(std::vector<double>{3.0, 4.0, 5.0}); // Mismatch!868 sectionData.push_back(std::move(section0));869 870 EXPECT_DEATH(VocabStorage(std::move(sectionData)),871 "All embeddings must have the same dimension");872}873#endif // NDEBUG874#endif // GTEST_HAS_DEATH_TEST875 876TEST(VocabStorageTest, IteratorBasics) {877 std::vector<std::vector<Embedding>> sectionData;878 879 std::vector<Embedding> section0;880 section0.emplace_back(std::vector<double>{1.0, 2.0});881 section0.emplace_back(std::vector<double>{3.0, 4.0});882 sectionData.push_back(std::move(section0));883 884 std::vector<Embedding> section1;885 section1.emplace_back(std::vector<double>{5.0, 6.0});886 sectionData.push_back(std::move(section1));887 888 VocabStorage storage(std::move(sectionData));889 890 // Test iterator basics891 auto it = storage.begin();892 auto end = storage.end();893 894 EXPECT_NE(it, end);895 896 // Check first embedding897 EXPECT_THAT((*it).getData(), ElementsAre(1.0, 2.0));898 899 // Advance to second embedding900 ++it;901 EXPECT_NE(it, end);902 EXPECT_THAT((*it).getData(), ElementsAre(3.0, 4.0));903 904 // Advance to third embedding (in section 1)905 ++it;906 EXPECT_NE(it, end);907 EXPECT_THAT((*it).getData(), ElementsAre(5.0, 6.0));908 909 // Advance past the end910 ++it;911 EXPECT_EQ(it, end);912}913 914TEST(VocabStorageTest, IteratorTraversal) {915 std::vector<std::vector<Embedding>> sectionData;916 917 // Section 0: 2 embeddings918 std::vector<Embedding> section0;919 section0.emplace_back(std::vector<double>{10.0});920 section0.emplace_back(std::vector<double>{20.0});921 sectionData.push_back(std::move(section0));922 923 // Section 1: 1 embedding924 std::vector<Embedding> section1;925 section1.emplace_back(std::vector<double>{25.0});926 sectionData.push_back(std::move(section1));927 928 // Section 2: 3 embeddings929 std::vector<Embedding> section2;930 section2.emplace_back(std::vector<double>{30.0});931 section2.emplace_back(std::vector<double>{40.0});932 section2.emplace_back(std::vector<double>{50.0});933 sectionData.push_back(std::move(section2));934 935 VocabStorage storage(std::move(sectionData));936 937 // Collect all values using iterator938 std::vector<double> values;939 for (const auto &emb : storage) {940 EXPECT_EQ(emb.size(), 1u);941 values.push_back(emb[0]);942 }943 944 // Should get all embeddings from all sections945 EXPECT_THAT(values, ElementsAre(10.0, 20.0, 25.0, 30.0, 40.0, 50.0));946}947 948TEST(VocabStorageTest, IteratorComparison) {949 std::vector<std::vector<Embedding>> sectionData;950 std::vector<Embedding> section0;951 section0.emplace_back(std::vector<double>{1.0});952 section0.emplace_back(std::vector<double>{2.0});953 sectionData.push_back(std::move(section0));954 955 VocabStorage storage(std::move(sectionData));956 957 auto it1 = storage.begin();958 auto it2 = storage.begin();959 auto end = storage.end();960 961 // Test equality962 EXPECT_EQ(it1, it2);963 EXPECT_NE(it1, end);964 965 // Advance one iterator966 ++it1;967 EXPECT_NE(it1, it2);968 EXPECT_NE(it1, end);969 970 // Advance second iterator to match971 ++it2;972 EXPECT_EQ(it1, it2);973 974 // Advance both to end975 ++it1;976 ++it2;977 EXPECT_EQ(it1, end);978 EXPECT_EQ(it2, end);979 EXPECT_EQ(it1, it2);980}981 982} // end anonymous namespace983