84 lines · cpp
1//===----- ProfDataUtils.cpp - Unit tests for ProfDataUtils ---------------===//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/AsmParser/Parser.h"10#include "llvm/IR/BasicBlock.h"11#include "llvm/IR/Function.h"12#include "llvm/IR/Instructions.h"13#include "llvm/IR/LLVMContext.h"14#include "llvm/IR/Module.h"15#include "llvm/IR/ProfDataUtils.h"16#include "llvm/Support/SourceMgr.h"17#include "gtest/gtest.h"18 19using namespace llvm;20 21static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {22 SMDiagnostic Err;23 std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);24 if (!Mod)25 Err.print("ProfDataUtilsTests", errs());26 return Mod;27}28 29TEST(ProfDataUtils, extractWeights) {30 LLVMContext C;31 std::unique_ptr<Module> M = parseIR(C, R"IR(32define void @foo(i1 %cond0) {33entry:34 br i1 %cond0, label %bb0, label %bb1, !prof !135bb0:36 %0 = mul i32 1, 237 br label %bb138bb1:39 ret void40}41 42!1 = !{!"branch_weights", i32 1, i32 100000}43)IR");44 Function *F = M->getFunction("foo");45 auto &Entry = F->getEntryBlock();46 auto &I = Entry.front();47 auto *Branch = dyn_cast<BranchInst>(&I);48 EXPECT_NE(nullptr, Branch);49 auto *ProfileData = Branch->getMetadata(LLVMContext::MD_prof);50 EXPECT_NE(ProfileData, nullptr);51 EXPECT_TRUE(hasProfMD(I));52 SmallVector<uint32_t> Weights;53 EXPECT_TRUE(extractBranchWeights(ProfileData, Weights));54 EXPECT_EQ(Weights[0], 1U);55 EXPECT_EQ(Weights[1], 100000U);56 EXPECT_EQ(Weights.size(), 2U);57}58 59TEST(ProfDataUtils, NoWeights) {60 LLVMContext C;61 std::unique_ptr<Module> M = parseIR(C, R"IR(62define void @foo(i1 %cond0) {63entry:64 br i1 %cond0, label %bb0, label %bb165bb0:66 %0 = mul i32 1, 267 br label %bb168bb1:69 ret void70}71)IR");72 Function *F = M->getFunction("foo");73 auto &Entry = F->getEntryBlock();74 auto &I = Entry.front();75 auto *Branch = dyn_cast<BranchInst>(&I);76 EXPECT_NE(nullptr, Branch);77 auto *ProfileData = Branch->getMetadata(LLVMContext::MD_prof);78 EXPECT_EQ(ProfileData, nullptr);79 EXPECT_FALSE(hasProfMD(I));80 SmallVector<uint32_t> Weights;81 EXPECT_FALSE(extractBranchWeights(ProfileData, Weights));82 EXPECT_EQ(Weights.size(), 0U);83}84