771 lines · cpp
1//===- MIR2VecTest.cpp ---------------------------------------------------===//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/CodeGen/MIR2Vec.h"10#include "llvm/CodeGen/MachineBasicBlock.h"11#include "llvm/CodeGen/MachineFunction.h"12#include "llvm/CodeGen/MachineInstr.h"13#include "llvm/CodeGen/MachineModuleInfo.h"14#include "llvm/CodeGen/TargetInstrInfo.h"15#include "llvm/CodeGen/TargetSubtargetInfo.h"16#include "llvm/IR/IRBuilder.h"17#include "llvm/IR/Module.h"18#include "llvm/MC/TargetRegistry.h"19#include "llvm/Support/TargetSelect.h"20#include "llvm/Support/raw_ostream.h"21#include "llvm/Target/TargetMachine.h"22#include "llvm/Target/TargetOptions.h"23#include "llvm/TargetParser/Triple.h"24#include "gtest/gtest.h"25 26using namespace llvm;27using namespace mir2vec;28using VocabMap = std::map<std::string, ir2vec::Embedding>;29 30namespace {31 32TEST(MIR2VecTest, RegexExtraction) {33 // Test simple instruction names34 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("NOP"), "NOP");35 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("RET"), "RET");36 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("ADD16ri"), "ADD");37 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("ADD32rr"), "ADD");38 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("ADD64rm"), "ADD");39 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("MOV8ri"), "MOV");40 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("MOV32mr"), "MOV");41 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("PUSH64r"), "PUSH");42 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("POP64r"), "POP");43 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("JMP_4"), "JMP");44 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("CALL64pcrel32"), "CALL");45 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("SOME_INSTR_123"),46 "SOME_INSTR");47 EXPECT_EQ(MIRVocabulary::extractBaseOpcodeName("123ADD"), "ADD");48 EXPECT_FALSE(MIRVocabulary::extractBaseOpcodeName("123").empty());49}50 51class MIR2VecVocabTestFixture : public ::testing::Test {52protected:53 std::unique_ptr<LLVMContext> Ctx;54 std::unique_ptr<Module> M;55 std::unique_ptr<TargetMachine> TM;56 const TargetInstrInfo *TII = nullptr;57 const TargetRegisterInfo *TRI = nullptr;58 std::unique_ptr<MachineModuleInfo> MMI;59 MachineFunction *MF = nullptr;60 61 static void SetUpTestCase() {62 InitializeAllTargets();63 InitializeAllTargetMCs();64 }65 66 void SetUp() override {67 Triple TargetTriple("x86_64-unknown-linux-gnu");68 std::string Error;69 const Target *T = TargetRegistry::lookupTarget("", TargetTriple, Error);70 if (!T) {71 GTEST_SKIP() << "x86_64-unknown-linux-gnu target triple not available; "72 "Skipping test";73 return;74 }75 76 Ctx = std::make_unique<LLVMContext>();77 M = std::make_unique<Module>("test", *Ctx);78 M->setTargetTriple(TargetTriple);79 80 TargetOptions Options;81 TM = std::unique_ptr<TargetMachine>(82 T->createTargetMachine(TargetTriple, "", "", Options, std::nullopt));83 if (!TM) {84 GTEST_SKIP() << "Failed to create X86 target machine; Skipping test";85 return;86 }87 88 // Set the data layout to match the target machine89 M->setDataLayout(TM->createDataLayout());90 91 // Create a dummy function to get subtarget info92 FunctionType *FT = FunctionType::get(Type::getVoidTy(*Ctx), false);93 Function *F =94 Function::Create(FT, Function::ExternalLinkage, "test", M.get());95 96 // Create MMI and MF to get TRI and MRI97 MMI = std::make_unique<MachineModuleInfo>(TM.get());98 MF = &MMI->getOrCreateMachineFunction(*F);99 100 // Get the target instruction info and register info101 TII = TM->getSubtargetImpl(*F)->getInstrInfo();102 TRI = TM->getSubtargetImpl(*F)->getRegisterInfo();103 if (!TII || !TRI) {104 GTEST_SKIP()105 << "Failed to get target instruction/register info; Skipping test";106 return;107 }108 }109 110 void TearDown() override {111 TII = nullptr;112 TRI = nullptr;113 }114 115 // Find an opcode by name116 int findOpcodeByName(StringRef Name) {117 for (unsigned Opcode = 1; Opcode < TII->getNumOpcodes(); ++Opcode) {118 if (TII->getName(Opcode) == Name)119 return Opcode;120 }121 return -1; // Not found122 }123 124 // Create a vocabulary with specific opcodes and embeddings125 // This might cause errors in future when the validation in126 // MIRVocabulary::generateStorage() enforces hard checks on the vocabulary127 // entries.128 Expected<MIRVocabulary> createTestVocab(129 std::initializer_list<std::pair<const char *, float>> Opcodes,130 std::initializer_list<std::pair<const char *, float>> CommonOperands,131 std::initializer_list<std::pair<const char *, float>> PhyRegs,132 std::initializer_list<std::pair<const char *, float>> VirtRegs,133 unsigned Dimension = 2) {134 assert(TII && TRI && MF && "Target info not initialized");135 VocabMap OpcodeMap, CommonOperandMap, PhyRegMap, VirtRegMap;136 for (const auto &[Name, Value] : Opcodes)137 OpcodeMap[Name] = Embedding(Dimension, Value);138 139 for (const auto &[Name, Value] : CommonOperands)140 CommonOperandMap[Name] = Embedding(Dimension, Value);141 142 for (const auto &[Name, Value] : PhyRegs)143 PhyRegMap[Name] = Embedding(Dimension, Value);144 145 for (const auto &[Name, Value] : VirtRegs)146 VirtRegMap[Name] = Embedding(Dimension, Value);147 148 // If any section is empty, create minimal maps for other vocabulary149 // sections to satisfy validation150 if (Opcodes.size() == 0)151 OpcodeMap["NOOP"] = Embedding(Dimension, 0.0f);152 if (CommonOperands.size() == 0)153 CommonOperandMap["Immediate"] = Embedding(Dimension, 0.0f);154 if (PhyRegs.size() == 0)155 PhyRegMap["GR32"] = Embedding(Dimension, 0.0f);156 if (VirtRegs.size() == 0)157 VirtRegMap["GR32"] = Embedding(Dimension, 0.0f);158 159 return MIRVocabulary::create(160 std::move(OpcodeMap), std::move(CommonOperandMap), std::move(PhyRegMap),161 std::move(VirtRegMap), *TII, *TRI, MF->getRegInfo());162 }163};164 165// Parameterized test for empty vocab sections166class MIR2VecVocabEmptySectionTestFixture167 : public MIR2VecVocabTestFixture,168 public ::testing::WithParamInterface<int> {169protected:170 void SetUp() override {171 MIR2VecVocabTestFixture::SetUp();172 // If base class setup was skipped (TII not initialized), skip derived setup173 if (!TII)174 GTEST_SKIP() << "Failed to get target instruction info in "175 "the base class setup; Skipping test";176 }177};178 179TEST_P(MIR2VecVocabEmptySectionTestFixture, EmptySectionFailsValidation) {180 int EmptySection = GetParam();181 VocabMap OpcodeMap, CommonOperandMap, PhyRegMap, VirtRegMap;182 183 if (EmptySection != 0)184 OpcodeMap["ADD"] = Embedding(2, 1.0f);185 if (EmptySection != 1)186 CommonOperandMap["Immediate"] = Embedding(2, 0.0f);187 if (EmptySection != 2)188 PhyRegMap["GR32"] = Embedding(2, 0.0f);189 if (EmptySection != 3)190 VirtRegMap["GR32"] = Embedding(2, 0.0f);191 192 ASSERT_TRUE(TII != nullptr);193 ASSERT_TRUE(TRI != nullptr);194 ASSERT_TRUE(MF != nullptr);195 196 auto VocabOrErr = MIRVocabulary::create(197 std::move(OpcodeMap), std::move(CommonOperandMap), std::move(PhyRegMap),198 std::move(VirtRegMap), *TII, *TRI, MF->getRegInfo());199 EXPECT_FALSE(static_cast<bool>(VocabOrErr))200 << "Factory method should fail when section " << EmptySection201 << " is empty";202 203 if (!VocabOrErr) {204 auto Err = VocabOrErr.takeError();205 std::string ErrorMsg = toString(std::move(Err));206 EXPECT_FALSE(ErrorMsg.empty());207 }208}209 210INSTANTIATE_TEST_SUITE_P(EmptySection, MIR2VecVocabEmptySectionTestFixture,211 ::testing::Values(0, 1, 2, 3));212 213TEST_F(MIR2VecVocabTestFixture, CanonicalOpcodeMappingTest) {214 // Test that same base opcodes get same canonical indices215 std::string BaseName1 = MIRVocabulary::extractBaseOpcodeName("ADD16ri");216 std::string BaseName2 = MIRVocabulary::extractBaseOpcodeName("ADD32rr");217 std::string BaseName3 = MIRVocabulary::extractBaseOpcodeName("ADD64rm");218 219 EXPECT_EQ(BaseName1, BaseName2);220 EXPECT_EQ(BaseName2, BaseName3);221 222 // Create a MIRVocabulary instance to test the mapping223 // Use a minimal MIRVocabulary to trigger canonical mapping construction224 Embedding Val = Embedding(64, 1.0f);225 auto TestVocabOrErr = createTestVocab({{"ADD", 1.0f}}, {}, {}, {}, 64);226 ASSERT_TRUE(static_cast<bool>(TestVocabOrErr))227 << "Failed to create vocabulary: "228 << toString(TestVocabOrErr.takeError());229 auto &TestVocab = *TestVocabOrErr;230 231 unsigned Index1 = TestVocab.getCanonicalIndexForBaseName(BaseName1);232 unsigned Index2 = TestVocab.getCanonicalIndexForBaseName(BaseName2);233 unsigned Index3 = TestVocab.getCanonicalIndexForBaseName(BaseName3);234 EXPECT_EQ(Index1, Index2);235 EXPECT_EQ(Index2, Index3);236 237 // Test that different base opcodes get different canonical indices238 std::string AddBase = MIRVocabulary::extractBaseOpcodeName("ADD32rr");239 std::string SubBase = MIRVocabulary::extractBaseOpcodeName("SUB32rr");240 std::string MovBase = MIRVocabulary::extractBaseOpcodeName("MOV32rr");241 242 unsigned AddIndex = TestVocab.getCanonicalIndexForBaseName(AddBase);243 unsigned SubIndex = TestVocab.getCanonicalIndexForBaseName(SubBase);244 unsigned MovIndex = TestVocab.getCanonicalIndexForBaseName(MovBase);245 246 EXPECT_NE(AddIndex, SubIndex);247 EXPECT_NE(SubIndex, MovIndex);248 EXPECT_NE(AddIndex, MovIndex);249 250 // Even though we only added "ADD" to the vocab, the canonical mapping251 // should assign unique indices to all the base opcodes of the target252 // Ideally, we would check against the exact number of unique base opcodes253 // for X86, but that would make the test brittle. So we just check that254 // the number is reasonably closer to the expected number (>6880) and not just255 // opcodes that we added.256 EXPECT_GT(TestVocab.getCanonicalSize(),257 6880u); // X86 has >6880 unique base opcodes258 259 // Check that the embeddings for opcodes not in the vocab are zero vectors260 int Add32rrOpcode = findOpcodeByName("ADD32rr");261 ASSERT_NE(Add32rrOpcode, -1) << "ADD32rr opcode not found";262 EXPECT_TRUE(TestVocab[Add32rrOpcode].approximatelyEquals(Val));263 264 int Sub32rrOpcode = findOpcodeByName("SUB32rr");265 ASSERT_NE(Sub32rrOpcode, -1) << "SUB32rr opcode not found";266 EXPECT_TRUE(267 TestVocab[Sub32rrOpcode].approximatelyEquals(Embedding(64, 0.0f)));268 269 int Mov32rrOpcode = findOpcodeByName("MOV32rr");270 ASSERT_NE(Mov32rrOpcode, -1) << "MOV32rr opcode not found";271 EXPECT_TRUE(272 TestVocab[Mov32rrOpcode].approximatelyEquals(Embedding(64, 0.0f)));273}274 275// Test deterministic mapping276TEST_F(MIR2VecVocabTestFixture, DeterministicMapping) {277 // Test that the same base name always maps to the same canonical index278 std::string BaseName = "ADD";279 280 // Create a MIRVocabulary instance to test deterministic mapping281 // Use a minimal MIRVocabulary to trigger canonical mapping construction282 auto TestVocabOrErr = createTestVocab({{"ADD", 1.0f}}, {}, {}, {}, 64);283 ASSERT_TRUE(static_cast<bool>(TestVocabOrErr))284 << "Failed to create vocabulary: "285 << toString(TestVocabOrErr.takeError());286 auto &TestVocab = *TestVocabOrErr;287 288 unsigned Index1 = TestVocab.getCanonicalIndexForBaseName(BaseName);289 unsigned Index2 = TestVocab.getCanonicalIndexForBaseName(BaseName);290 unsigned Index3 = TestVocab.getCanonicalIndexForBaseName(BaseName);291 EXPECT_EQ(Index2, Index3);292 293 // Test across multiple runs294 for (int Pos = 0; Pos < 100; ++Pos) {295 unsigned Index = TestVocab.getCanonicalIndexForBaseName(BaseName);296 EXPECT_EQ(Index, Index1);297 }298}299 300// Test MIRVocabulary construction301TEST_F(MIR2VecVocabTestFixture, VocabularyConstruction) {302 auto VocabOrErr =303 createTestVocab({{"ADD", 1.0f}, {"SUB", 2.0f}}, {}, {}, {}, 128);304 ASSERT_TRUE(static_cast<bool>(VocabOrErr))305 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());306 auto &Vocab = *VocabOrErr;307 EXPECT_EQ(Vocab.getDimension(), 128u);308 309 // Test iterator - iterates over individual embeddings310 auto IT = Vocab.begin();311 EXPECT_NE(IT, Vocab.end());312 313 // Check first embedding exists and has correct dimension314 EXPECT_EQ((*IT).size(), 128u);315 316 size_t Count = 0;317 for (auto IT = Vocab.begin(); IT != Vocab.end(); ++IT) {318 EXPECT_EQ((*IT).size(), 128u);319 ++Count;320 }321 EXPECT_GT(Count, 0u);322}323 324// Fixture for embedding related tests325class MIR2VecEmbeddingTestFixture : public MIR2VecVocabTestFixture {326protected:327 void SetUp() override {328 MIR2VecVocabTestFixture::SetUp();329 // If base class setup was skipped (TII not initialized), skip derived setup330 if (!TII)331 GTEST_SKIP() << "Failed to get target instruction info in "332 "the base class setup; Skipping test";333 }334 335 void TearDown() override { MIR2VecVocabTestFixture::TearDown(); }336 337 // Create a machine instruction338 MachineInstr *createMachineInstr(MachineBasicBlock &MBB, unsigned Opcode) {339 const MCInstrDesc &Desc = TII->get(Opcode);340 // Create instruction - operands don't affect opcode-based embeddings341 MachineInstr *MI = BuildMI(MBB, MBB.end(), DebugLoc(), Desc);342 return MI;343 }344 345 MachineInstr *createMachineInstr(MachineBasicBlock &MBB,346 const char *OpcodeName) {347 int Opcode = findOpcodeByName(OpcodeName);348 if (Opcode == -1)349 return nullptr;350 return createMachineInstr(MBB, Opcode);351 }352 353 void createMachineInstrs(MachineBasicBlock &MBB,354 std::initializer_list<const char *> Opcodes) {355 for (const char *OpcodeName : Opcodes) {356 MachineInstr *MI = createMachineInstr(MBB, OpcodeName);357 ASSERT_TRUE(MI != nullptr);358 }359 }360};361 362// Test factory method for creating embedder363TEST_F(MIR2VecEmbeddingTestFixture, CreateSymbolicEmbedder) {364 auto VocabOrErr =365 MIRVocabulary::createDummyVocabForTest(*TII, *TRI, MF->getRegInfo(), 1);366 ASSERT_TRUE(static_cast<bool>(VocabOrErr))367 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());368 auto &V = *VocabOrErr;369 auto Emb = MIREmbedder::create(MIR2VecKind::Symbolic, *MF, V);370 EXPECT_NE(Emb, nullptr);371}372 373TEST_F(MIR2VecEmbeddingTestFixture, CreateInvalidMode) {374 auto VocabOrErr =375 MIRVocabulary::createDummyVocabForTest(*TII, *TRI, MF->getRegInfo(), 1);376 ASSERT_TRUE(static_cast<bool>(VocabOrErr))377 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());378 auto &V = *VocabOrErr;379 auto Result = MIREmbedder::create(static_cast<MIR2VecKind>(-1), *MF, V);380 EXPECT_FALSE(static_cast<bool>(Result));381}382 383// Test SymbolicMIREmbedder with simple target opcodes384TEST_F(MIR2VecEmbeddingTestFixture, TestSymbolicEmbedder) {385 // Create a test vocabulary with specific values386 auto VocabOrErr = createTestVocab(387 {388 {"NOOP", 1.0f}, // [1.0, 1.0, 1.0, 1.0]389 {"RET", 2.0f}, // [2.0, 2.0, 2.0, 2.0]390 {"TRAP", 3.0f} // [3.0, 3.0, 3.0, 3.0]391 },392 {}, {}, {}, 4);393 ASSERT_TRUE(static_cast<bool>(VocabOrErr))394 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());395 auto &Vocab = *VocabOrErr;396 // Create a basic block using fixture's MF397 MachineBasicBlock *MBB = MF->CreateMachineBasicBlock();398 MF->push_back(MBB);399 400 // Use real X86 opcodes that should exist and not be pseudo401 auto NoopInst = createMachineInstr(*MBB, "NOOP");402 ASSERT_TRUE(NoopInst != nullptr);403 404 auto RetInst = createMachineInstr(*MBB, "RET64");405 ASSERT_TRUE(RetInst != nullptr);406 407 auto TrapInst = createMachineInstr(*MBB, "TRAP");408 ASSERT_TRUE(TrapInst != nullptr);409 410 // Verify these are not pseudo instructions411 ASSERT_FALSE(NoopInst->isPseudo()) << "NOOP is marked as pseudo instruction";412 ASSERT_FALSE(RetInst->isPseudo()) << "RET is marked as pseudo instruction";413 ASSERT_FALSE(TrapInst->isPseudo()) << "TRAP is marked as pseudo instruction";414 415 // Create embedder416 auto Embedder = SymbolicMIREmbedder::create(*MF, Vocab);417 ASSERT_TRUE(Embedder != nullptr);418 419 // Test instruction embeddings420 auto NoopEmb = Embedder->getMInstVector(*NoopInst);421 auto RetEmb = Embedder->getMInstVector(*RetInst);422 auto TrapEmb = Embedder->getMInstVector(*TrapInst);423 424 // Verify embeddings match expected values (accounting for weight scaling)425 float ExpectedWeight = mir2vec::OpcWeight; // Global weight from command line426 EXPECT_TRUE(NoopEmb.approximatelyEquals(Embedding(4, 1.0f * ExpectedWeight)));427 EXPECT_TRUE(RetEmb.approximatelyEquals(Embedding(4, 2.0f * ExpectedWeight)));428 EXPECT_TRUE(TrapEmb.approximatelyEquals(Embedding(4, 3.0f * ExpectedWeight)));429 430 // Test basic block embedding (should be sum of instruction embeddings)431 auto MBBVector = Embedder->getMBBVector(*MBB);432 433 // Expected BB vector: NOOP + RET + TRAP = [1+2+3, 1+2+3, 1+2+3, 1+2+3] *434 // weight = [6, 6, 6, 6] * weight435 Embedding ExpectedMBBVector(4, 6.0f * ExpectedWeight);436 EXPECT_TRUE(MBBVector.approximatelyEquals(ExpectedMBBVector));437 438 // Test function embedding (should equal MBB embedding since we have one MBB)439 auto MFuncVector = Embedder->getMFunctionVector();440 EXPECT_TRUE(MFuncVector.approximatelyEquals(ExpectedMBBVector));441}442 443// Test embedder with multiple basic blocks444TEST_F(MIR2VecEmbeddingTestFixture, MultipleBasicBlocks) {445 // Create a test vocabulary446 auto VocabOrErr =447 createTestVocab({{"NOOP", 1.0f}, {"TRAP", 2.0f}}, {}, {}, {});448 ASSERT_TRUE(static_cast<bool>(VocabOrErr))449 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());450 auto &Vocab = *VocabOrErr;451 452 // Create two basic blocks using fixture's MF453 MachineBasicBlock *MBB1 = MF->CreateMachineBasicBlock();454 MachineBasicBlock *MBB2 = MF->CreateMachineBasicBlock();455 MF->push_back(MBB1);456 MF->push_back(MBB2);457 458 createMachineInstrs(*MBB1, {"NOOP", "NOOP"});459 createMachineInstr(*MBB2, "TRAP");460 461 // Create embedder462 auto Embedder = SymbolicMIREmbedder::create(*MF, Vocab);463 ASSERT_TRUE(Embedder != nullptr);464 465 // Test basic block embeddings466 auto MBB1Vector = Embedder->getMBBVector(*MBB1);467 auto MBB2Vector = Embedder->getMBBVector(*MBB2);468 469 float ExpectedWeight = mir2vec::OpcWeight;470 // BB1: NOOP + NOOP = 2 * ([1, 1] * weight)471 Embedding ExpectedMBB1Vector(2, 2.0f * ExpectedWeight);472 EXPECT_TRUE(MBB1Vector.approximatelyEquals(ExpectedMBB1Vector));473 474 // BB2: TRAP = [2, 2] * weight475 Embedding ExpectedMBB2Vector(2, 2.0f * ExpectedWeight);476 EXPECT_TRUE(MBB2Vector.approximatelyEquals(ExpectedMBB2Vector));477 478 // Function embedding: BB1 + BB2 = [2+2, 2+2] * weight = [4, 4] * weight479 // Function embedding should be just the first BB embedding as the second BB480 // is unreachable481 auto MFuncVector = Embedder->getMFunctionVector();482 EXPECT_TRUE(MFuncVector.approximatelyEquals(ExpectedMBB1Vector));483 484 // Add a branch from BB1 to BB2 to make both reachable; now function embedding485 // should be MBB1 + MBB2486 MBB1->addSuccessor(MBB2);487 auto NewMFuncVector = Embedder->getMFunctionVector(); // Recompute embeddings488 Embedding ExpectedFuncVector = MBB1Vector + MBB2Vector;489 EXPECT_TRUE(NewMFuncVector.approximatelyEquals(ExpectedFuncVector));490}491 492// Test embedder with empty basic block493TEST_F(MIR2VecEmbeddingTestFixture, EmptyBasicBlock) {494 495 // Create an empty basic block496 MachineBasicBlock *MBB = MF->CreateMachineBasicBlock();497 MF->push_back(MBB);498 499 // Create embedder500 auto VocabOrErr =501 MIRVocabulary::createDummyVocabForTest(*TII, *TRI, MF->getRegInfo(), 2);502 ASSERT_TRUE(static_cast<bool>(VocabOrErr))503 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());504 auto &V = *VocabOrErr;505 auto Embedder = SymbolicMIREmbedder::create(*MF, V);506 ASSERT_TRUE(Embedder != nullptr);507 508 // Test that empty BB has zero embedding509 auto MBBVector = Embedder->getMBBVector(*MBB);510 Embedding ExpectedBBVector(2, 0.0f);511 EXPECT_TRUE(MBBVector.approximatelyEquals(ExpectedBBVector));512 513 // Function embedding should also be zero514 auto MFuncVector = Embedder->getMFunctionVector();515 EXPECT_TRUE(MFuncVector.approximatelyEquals(ExpectedBBVector));516}517 518// Test embedder with opcodes not in vocabulary519TEST_F(MIR2VecEmbeddingTestFixture, UnknownOpcodes) {520 // Create a test vocabulary with limited entries521 // SUB is intentionally not included522 auto VocabOrErr = createTestVocab({{"ADD", 1.0f}}, {}, {}, {});523 ASSERT_TRUE(static_cast<bool>(VocabOrErr))524 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());525 auto &Vocab = *VocabOrErr;526 527 // Create a basic block528 MachineBasicBlock *MBB = MF->CreateMachineBasicBlock();529 MF->push_back(MBB);530 531 // Find opcodes532 int AddOpcode = findOpcodeByName("ADD32rr");533 int SubOpcode = findOpcodeByName("SUB32rr");534 535 ASSERT_NE(AddOpcode, -1) << "ADD32rr opcode not found";536 ASSERT_NE(SubOpcode, -1) << "SUB32rr opcode not found";537 538 // Create instructions539 MachineInstr *AddInstr = createMachineInstr(*MBB, AddOpcode);540 MachineInstr *SubInstr = createMachineInstr(*MBB, SubOpcode);541 542 // Create embedder543 auto Embedder = SymbolicMIREmbedder::create(*MF, Vocab);544 ASSERT_TRUE(Embedder != nullptr);545 546 // Test instruction embeddings547 auto AddVector = Embedder->getMInstVector(*AddInstr);548 auto SubVector = Embedder->getMInstVector(*SubInstr);549 550 float ExpectedWeight = mir2vec::OpcWeight;551 // ADD should have the embedding from vocabulary552 EXPECT_TRUE(553 AddVector.approximatelyEquals(Embedding(2, 1.0f * ExpectedWeight)));554 555 // SUB should have zero embedding (not in vocabulary)556 EXPECT_TRUE(SubVector.approximatelyEquals(Embedding(2, 0.0f)));557 558 // Basic block embedding should be ADD + SUB = [1.0, 1.0] * weight + [0.0,559 // 0.0] = [1.0, 1.0] * weight560 const auto &MBBVector = Embedder->getMBBVector(*MBB);561 Embedding ExpectedBBVector(2, 1.0f * ExpectedWeight);562 EXPECT_TRUE(MBBVector.approximatelyEquals(ExpectedBBVector));563}564 565// Test vocabulary string key generation566TEST_F(MIR2VecEmbeddingTestFixture, VocabularyStringKeys) {567 auto VocabOrErr =568 createTestVocab({{"ADD", 1.0f}, {"SUB", 2.0f}}, {}, {}, {}, 2);569 ASSERT_TRUE(static_cast<bool>(VocabOrErr))570 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());571 auto &Vocab = *VocabOrErr;572 573 // Test that we can get string keys for all positions574 for (size_t Pos = 0; Pos < Vocab.getCanonicalSize(); ++Pos) {575 std::string Key = Vocab.getStringKey(Pos);576 EXPECT_FALSE(Key.empty()) << "Empty key at position " << Pos;577 }578 579 // Test specific known positions if we can identify them580 unsigned AddIndex = Vocab.getCanonicalIndexForBaseName("ADD");581 std::string AddKey = Vocab.getStringKey(AddIndex);582 EXPECT_EQ(AddKey, "ADD");583 584 unsigned SubIndex = Vocab.getCanonicalIndexForBaseName("SUB");585 std::string SubKey = Vocab.getStringKey(SubIndex);586 EXPECT_EQ(SubKey, "SUB");587 588 unsigned ImmIndex = Vocab.getCanonicalIndexForOperandName("Immediate");589 std::string ImmKey = Vocab.getStringKey(ImmIndex);590 EXPECT_EQ(ImmKey, "Immediate");591 592 unsigned PhyRegIndex = Vocab.getCanonicalIndexForRegisterClass("GR32", true);593 std::string PhyRegKey = Vocab.getStringKey(PhyRegIndex);594 EXPECT_EQ(PhyRegKey, "PhyReg_GR32");595 596 unsigned VirtRegIndex =597 Vocab.getCanonicalIndexForRegisterClass("GR32", false);598 std::string VirtRegKey = Vocab.getStringKey(VirtRegIndex);599 EXPECT_EQ(VirtRegKey, "VirtReg_GR32");600}601 602// Test vocabulary dimension consistency603TEST_F(MIR2VecEmbeddingTestFixture, DimensionConsistency) {604 auto VocabOrErr = createTestVocab({{"TEST", 1.0f}}, {}, {}, {}, 5);605 ASSERT_TRUE(static_cast<bool>(VocabOrErr))606 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());607 auto &Vocab = *VocabOrErr;608 609 EXPECT_EQ(Vocab.getDimension(), 5u);610 611 // All embeddings should have the same dimension612 for (auto IT = Vocab.begin(); IT != Vocab.end(); ++IT)613 EXPECT_EQ((*IT).size(), 5u);614}615 616// Test invalid register handling through machine instruction creation617TEST_F(MIR2VecEmbeddingTestFixture, InvalidRegisterHandling) {618 float MOVValue = 1.5f;619 float ImmValue = 0.5f;620 float PhyRegValue = 0.2f;621 auto VocabOrErr = createTestVocab(622 {{"MOV", MOVValue}}, {{"Immediate", ImmValue}},623 {{"GR8_ABCD_H", PhyRegValue}, {"GR8_ABCD_L", PhyRegValue + 0.1f}}, {}, 3);624 ASSERT_TRUE(static_cast<bool>(VocabOrErr))625 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());626 auto &Vocab = *VocabOrErr;627 628 MachineBasicBlock *MBB = MF->CreateMachineBasicBlock();629 MF->push_back(MBB);630 631 // Create a MOV instruction with actual operands including potential $noreg632 // This tests the actual scenario where invalid registers are encountered633 auto MovOpcode = findOpcodeByName("MOV32mr");634 ASSERT_NE(MovOpcode, -1) << "MOV32mr opcode not found";635 const MCInstrDesc &Desc = TII->get(MovOpcode);636 637 // Use available physical registers from the target638 unsigned BaseReg =639 TRI->getNumRegs() > 1 ? 1 : 0; // First available physical register640 unsigned ValueReg = TRI->getNumRegs() > 2 ? 2 : BaseReg;641 642 // MOV32mr typically has: base, scale, index, displacement, segment, value643 // Use the MachineInstrBuilder API properly644 auto MovInst = BuildMI(*MBB, MBB->end(), DebugLoc(), Desc)645 .addReg(BaseReg) // base646 .addImm(1) // scale647 .addReg(0) // index ($noreg)648 .addImm(-4) // displacement649 .addReg(0) // segment ($noreg)650 .addReg(ValueReg); // value651 652 auto Embedder = SymbolicMIREmbedder::create(*MF, Vocab);653 ASSERT_TRUE(Embedder != nullptr);654 655 // This should not crash even if the instruction has $noreg operands656 auto InstEmb = Embedder->getMInstVector(*MovInst);657 EXPECT_EQ(InstEmb.size(), 3u);658 659 // Test the expected embedding value660 Embedding ExpectedOpcodeContribution(3, MOVValue * mir2vec::OpcWeight);661 auto ExpectedOperandContribution =662 Embedding(3, PhyRegValue * mir2vec::RegOperandWeight) // Base663 + Embedding(3, ImmValue * mir2vec::CommonOperandWeight) // Scale664 + Embedding(3, 0.0f) // noreg665 + Embedding(3, ImmValue * mir2vec::CommonOperandWeight) // displacement666 + Embedding(3, 0.0f) // noreg667 + Embedding(3, (PhyRegValue + 0.1f) * mir2vec::RegOperandWeight); // Value668 auto ExpectedEmb = ExpectedOpcodeContribution + ExpectedOperandContribution;669 EXPECT_TRUE(InstEmb.approximatelyEquals(ExpectedEmb))670 << "MOV instruction embedding should match expected embedding";671}672 673// Test handling of both physical and virtual registers in an instruction674TEST_F(MIR2VecEmbeddingTestFixture, PhysicalAndVirtualRegisterHandling) {675 float MOVValue = 2.0f;676 float ImmValue = 0.7f;677 float PhyRegValue = 0.3f;678 float VirtRegValue = 0.9f;679 680 // Find GR32 register class681 const TargetRegisterClass *GR32RC = nullptr;682 for (unsigned i = 0; i < TRI->getNumRegClasses(); ++i) {683 const TargetRegisterClass *RC = TRI->getRegClass(i);684 if (std::string(TRI->getRegClassName(RC)) == "GR32") {685 GR32RC = RC;686 break;687 }688 }689 ASSERT_TRUE(GR32RC != nullptr && GR32RC->isAllocatable())690 << "No allocatable GR32 register class found";691 692 // Get first available physical register from GR32693 unsigned PhyReg = *GR32RC->begin();694 // Create a virtual register of class GR32695 unsigned VirtReg = MF->getRegInfo().createVirtualRegister(GR32RC);696 697 // Create vocabulary with register class based keys698 auto VocabOrErr =699 createTestVocab({{"MOV", MOVValue}}, {{"Immediate", ImmValue}},700 {{"GR32_AD", PhyRegValue}}, // GR32_AD is the minimal key701 {{"GR32", VirtRegValue}}, 4);702 ASSERT_TRUE(static_cast<bool>(VocabOrErr))703 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());704 auto &Vocab = *VocabOrErr;705 706 MachineBasicBlock *MBB = MF->CreateMachineBasicBlock();707 MF->push_back(MBB);708 709 // Create a MOV32rr instruction: MOV32rr dst, src710 auto MovOpcode = findOpcodeByName("MOV32rr");711 ASSERT_NE(MovOpcode, -1) << "MOV32rr opcode not found";712 const MCInstrDesc &Desc = TII->get(MovOpcode);713 714 // MOV32rr: dst (physical), src (virtual)715 auto MovInst = BuildMI(*MBB, MBB->end(), DebugLoc(), Desc)716 .addReg(PhyReg) // physical register destination717 .addReg(VirtReg); // virtual register source718 719 // Create embedder with virtual register support720 auto Embedder = SymbolicMIREmbedder::create(*MF, Vocab);721 ASSERT_TRUE(Embedder != nullptr);722 723 // This should not crash and should produce a valid embedding724 auto InstEmb = Embedder->getMInstVector(*MovInst);725 EXPECT_EQ(InstEmb.size(), 4u);726 727 // Test the expected embedding value728 Embedding ExpectedOpcodeContribution(4, MOVValue * mir2vec::OpcWeight);729 auto ExpectedOperandContribution =730 Embedding(4, PhyRegValue * mir2vec::RegOperandWeight) // dst (physical)731 + Embedding(4, VirtRegValue * mir2vec::RegOperandWeight); // src (virtual)732 auto ExpectedEmb = ExpectedOpcodeContribution + ExpectedOperandContribution;733 EXPECT_TRUE(InstEmb.approximatelyEquals(ExpectedEmb))734 << "MOV32rr instruction embedding should match expected embedding";735}736 737// Test precise embedding calculation with known operands738TEST_F(MIR2VecEmbeddingTestFixture, EmbeddingCalculation) {739 auto VocabOrErr = createTestVocab({{"NOOP", 2.0f}}, {}, {}, {}, 2);740 ASSERT_TRUE(static_cast<bool>(VocabOrErr))741 << "Failed to create vocabulary: " << toString(VocabOrErr.takeError());742 auto &Vocab = *VocabOrErr;743 744 MachineBasicBlock *MBB = MF->CreateMachineBasicBlock();745 MF->push_back(MBB);746 747 // Create a simple NOOP instruction (no operands)748 auto NoopInst = createMachineInstr(*MBB, "NOOP");749 ASSERT_TRUE(NoopInst != nullptr);750 751 auto Embedder = SymbolicMIREmbedder::create(*MF, Vocab);752 ASSERT_TRUE(Embedder != nullptr);753 754 // Get the instruction embedding755 auto InstEmb = Embedder->getMInstVector(*NoopInst);756 EXPECT_EQ(InstEmb.size(), 2u);757 758 // For NOOP with no operands, the embedding should be exactly the opcode759 // embedding760 float ExpectedWeight = mir2vec::OpcWeight;761 Embedding ExpectedEmb(2, 2.0f * ExpectedWeight);762 763 EXPECT_TRUE(InstEmb.approximatelyEquals(ExpectedEmb))764 << "NOOP instruction embedding should match opcode embedding";765 766 // Verify individual components767 EXPECT_FLOAT_EQ(InstEmb[0], 2.0f * ExpectedWeight);768 EXPECT_FLOAT_EQ(InstEmb[1], 2.0f * ExpectedWeight);769}770} // namespace771