[PGO] Add interfaces to annotate instr with VP data

  Add interfaces to do value profile data IR annnotation
  and read. Needed by both FE and IR based PGO.

llvm-svn: 259813

3 files changed, 167 insertions, 7 deletions

diff --git a/llvm/include/llvm/ProfileData/InstrProf.h b/llvm/include/llvm/ProfileData/InstrProf.h
index a577936ff00..a4b2775886b 100644
--- a/llvm/include/llvm/ProfileData/InstrProf.h
+++ b/llvm/include/llvm/ProfileData/InstrProf.h
@@ -178,8 +178,8 @@ StringRef getFuncNameWithoutPrefix(StringRef PGOFuncName, StringRef FileName);
 /// The first field is the legnth of the uncompressed strings, and the
 /// the second field is the length of the zlib-compressed string.
 /// Both fields are encoded in ULEB128.  If \c doCompress is false, the
-///  third field is the uncompressed strings; otherwise it is the 
-/// compressed string. When the string compression is off, the 
+///  third field is the uncompressed strings; otherwise it is the
+/// compressed string. When the string compression is off, the
 /// second field will have value zero.
 int collectPGOFuncNameStrings(const std::vector<std::string> &NameStrs,
                               bool doCompression, std::string &Result);
@@ -193,6 +193,29 @@ class InstrProfSymtab;
 /// bytes. This method decodes the string and populates the \c Symtab.
 int readPGOFuncNameStrings(StringRef NameStrings, InstrProfSymtab &Symtab);
 
+enum InstrProfValueKind : uint32_t {
+#define VALUE_PROF_KIND(Enumerator, Value) Enumerator = Value,
+#include "llvm/ProfileData/InstrProfData.inc"
+};
+
+struct InstrProfRecord;
+
+/// Extract value profile data for value site \p SiteIdx from \p InstrProfR
+/// and annotate the instruction \p Inst with the value profile meta data.
+void annotateValueSite(Module &M, Instruction &Inst,
+                       const InstrProfRecord &InstrProfR,
+                       InstrProfValueKind ValueKind, uint32_t SiteIndx);
+/// Extract the value profile data from the \p Inst which is annotated with
+/// value
+/// profile meta data. Return false if there is no value data annotated,
+/// otherwise
+/// return true.
+bool getValueProfDataFromInst(const Instruction &Inst,
+                              InstrProfValueKind ValueKind,
+                              uint32_t MaxNumValueData,
+                              InstrProfValueData ValueData[],
+                              uint32_t &ActualNumValueData, uint64_t &TotalC);
+
 const std::error_category &instrprof_category();
 
 enum class instrprof_error {
@@ -227,11 +250,6 @@ inline instrprof_error MergeResult(instrprof_error &Accumulator,
   return Accumulator;
 }
 
-enum InstrProfValueKind : uint32_t {
-#define VALUE_PROF_KIND(Enumerator, Value) Enumerator = Value,
-#include "llvm/ProfileData/InstrProfData.inc"
-};
-
 namespace object {
 class SectionRef;
 }
diff --git a/llvm/lib/ProfileData/InstrProf.cpp b/llvm/lib/ProfileData/InstrProf.cpp
index 6bf52436811..ae97ace4633 100644
--- a/llvm/lib/ProfileData/InstrProf.cpp
+++ b/llvm/lib/ProfileData/InstrProf.cpp
@@ -17,6 +17,7 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -606,6 +607,92 @@ void ValueProfData::swapBytesFromHost(support::endianness Endianness) {
   sys::swapByteOrder<uint32_t>(NumValueKinds);
 }
 
+void annotateValueSite(Module &M, Instruction &Inst,
+                       const InstrProfRecord &InstrProfR,
+                       InstrProfValueKind ValueKind, uint32_t SiteIdx) {
+  uint32_t NV = InstrProfR.getNumValueDataForSite(ValueKind, SiteIdx);
+
+  uint64_t Sum = 0;
+  std::unique_ptr<InstrProfValueData[]> VD =
+      InstrProfR.getValueForSite(ValueKind, SiteIdx, &Sum);
+
+  LLVMContext &Ctx = M.getContext();
+  MDBuilder MDHelper(Ctx);
+  SmallVector<Metadata *, 3> Vals;
+  // Tag
+  Vals.push_back(MDHelper.createString("VP"));
+  // Value Kind
+  Vals.push_back(MDHelper.createConstant(
+      ConstantInt::get(Type::getInt32Ty(Ctx), ValueKind)));
+  // Total Count
+  Vals.push_back(
+      MDHelper.createConstant(ConstantInt::get(Type::getInt64Ty(Ctx), Sum)));
+
+  // Value Profile Data
+  uint32_t MDCount = 3;
+  for (uint32_t I = 0; I < NV; ++I) {
+    Vals.push_back(MDHelper.createConstant(
+        ConstantInt::get(Type::getInt64Ty(Ctx), VD[I].Value)));
+    Vals.push_back(MDHelper.createConstant(
+        ConstantInt::get(Type::getInt64Ty(Ctx), VD[I].Count)));
+    if (--MDCount == 0)
+      break;
+  }
+  Inst.setMetadata(LLVMContext::MD_prof, MDNode::get(Ctx, Vals));
+}
+
+bool getValueProfDataFromInst(const Instruction &Inst,
+                              InstrProfValueKind ValueKind,
+                              uint32_t MaxNumValueData,
+                              InstrProfValueData ValueData[],
+                              uint32_t &ActualNumValueData, uint64_t &TotalC) {
+  MDNode *MD = Inst.getMetadata(LLVMContext::MD_prof);
+  if (!MD)
+    return false;
+
+  unsigned NOps = MD->getNumOperands();
+
+  if (NOps < 5)
+    return false;
+
+  // Operand 0 is a string tag "VP":
+  MDString *Tag = cast<MDString>(MD->getOperand(0));
+  if (!Tag)
+    return false;
+
+  if (!Tag->getString().equals("VP"))
+    return false;
+
+  // Now check kind:
+  ConstantInt *KindInt = mdconst::dyn_extract<ConstantInt>(MD->getOperand(1));
+  if (!KindInt)
+    return false;
+  if (KindInt->getZExtValue() != ValueKind)
+    return false;
+
+  // Get total count
+  ConstantInt *TotalCInt = mdconst::dyn_extract<ConstantInt>(MD->getOperand(2));
+  if (!TotalCInt)
+    return false;
+  TotalC = TotalCInt->getZExtValue();
+
+  ActualNumValueData = 0;
+
+  for (unsigned I = 3; I < NOps; I += 2) {
+    if (ActualNumValueData >= MaxNumValueData)
+      break;
+    ConstantInt *Value = mdconst::dyn_extract<ConstantInt>(MD->getOperand(I));
+    ConstantInt *Count =
+        mdconst::dyn_extract<ConstantInt>(MD->getOperand(I + 1));
+    if (!Value || !Count)
+      return false;
+    ValueData[ActualNumValueData].Value = Value->getZExtValue();
+    ValueData[ActualNumValueData].Count = Count->getZExtValue();
+    ActualNumValueData++;
+  }
+  return true;
+}
+
 // The argument to this method is a vector of cutoff percentages and the return
 // value is a vector of (Cutoff, MinBlockCount, NumBlocks) triplets.
 void ProfileSummary::computeDetailedSummary() {
diff --git a/llvm/unittests/ProfileData/InstrProfTest.cpp b/llvm/unittests/ProfileData/InstrProfTest.cpp
index be3ed427636..e8e24bda60d 100644
--- a/llvm/unittests/ProfileData/InstrProfTest.cpp
+++ b/llvm/unittests/ProfileData/InstrProfTest.cpp
@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/ProfileData/InstrProfReader.h"
@@ -225,6 +226,60 @@ TEST_P(MaybeSparseInstrProfTest, get_icall_data_read_write) {
   ASSERT_EQ(StringRef((const char *)VD[2].Value, 7), StringRef("callee1"));
 }
 
+TEST_P(MaybeSparseInstrProfTest, annotate_vp_data) {
+  InstrProfRecord Record("caller", 0x1234, {1, 2});
+  Record.reserveSites(IPVK_IndirectCallTarget, 1);
+  InstrProfValueData VD0[] = {{1000, 1}, {2000, 2}, {3000, 3}};
+  Record.addValueData(IPVK_IndirectCallTarget, 0, VD0, 3, nullptr);
+  Writer.addRecord(std::move(Record));
+  auto Profile = Writer.writeBuffer();
+  readProfile(std::move(Profile));
+  ErrorOr<InstrProfRecord> R = Reader->getInstrProfRecord("caller", 0x1234);
+  ASSERT_TRUE(NoError(R.getError()));
+
+  LLVMContext Ctx;
+  std::unique_ptr<Module> M(new Module("MyModule", Ctx));
+  FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx),
+                                        /*isVarArg=*/false);
+  Function *F =
+      Function::Create(FTy, Function::ExternalLinkage, "caller", M.get());
+  BasicBlock *BB = BasicBlock::Create(Ctx, "", F);
+
+  IRBuilder<> Builder(BB);
+  BasicBlock *TBB = BasicBlock::Create(Ctx, "", F);
+  BasicBlock *FBB = BasicBlock::Create(Ctx, "", F);
+
+  // Use branch instruction to annotate with value profile data for simplicity
+  Instruction *Inst = Builder.CreateCondBr(Builder.getTrue(), TBB, FBB);
+  Instruction *Inst2 = Builder.CreateCondBr(Builder.getTrue(), TBB, FBB);
+  annotateValueSite(*M.get(), *Inst, R.get(), IPVK_IndirectCallTarget, 0);
+
+  InstrProfValueData ValueData[5];
+  uint32_t N;
+  uint64_t T;
+  bool Res = getValueProfDataFromInst(*Inst, IPVK_IndirectCallTarget, 5,
+                                      ValueData, N, T);
+  ASSERT_TRUE(Res);
+  ASSERT_EQ(3U, N);
+  ASSERT_EQ(6U, T);
+  // The result should be sorted already:
+  ASSERT_EQ(3000U, ValueData[0].Value);
+  ASSERT_EQ(3U, ValueData[0].Count);
+  ASSERT_EQ(2000U, ValueData[1].Value);
+  ASSERT_EQ(2U, ValueData[1].Count);
+  ASSERT_EQ(1000U, ValueData[2].Value);
+  ASSERT_EQ(1U, ValueData[2].Count);
+  Res = getValueProfDataFromInst(*Inst, IPVK_IndirectCallTarget, 1, ValueData,
+                                 N, T);
+  ASSERT_TRUE(Res);
+  ASSERT_EQ(1U, N);
+  ASSERT_EQ(6U, T);
+
+  Res = getValueProfDataFromInst(*Inst2, IPVK_IndirectCallTarget, 5, ValueData,
+                                 N, T);
+  ASSERT_FALSE(Res);
+}
+
 TEST_P(MaybeSparseInstrProfTest, get_icall_data_read_write_with_weight) {
   InstrProfRecord Record1("caller", 0x1234, {1, 2});
   InstrProfRecord Record2("callee1", 0x1235, {3, 4});