[PGO] Improve the working set size heuristics under the partial sample PGO.

[hjyamauchi]

Summary

The working set size heuristics (ProfileSummaryInfo::hasHugeWorkingSetSize)
under the partial sample PGO may not be accurate because the profile is partial
and the number of hot profile counters in the ProfileSummary may not reflect the
actual working set size of the program being compiled.

To improve this, the (approximated) ratio of the the number of profile counters
of the program being compiled to the number of profile counters in the partial
sample profile is computed (which is called the partial profile ratio) and the
working set size of the profile is scaled by this ratio to reflect the working
set size of the program being compiled and used for the working set size
heuristics.

The partial profile ratio is approximated based on the number of the basic
blocks in the program and the NumCounts field in the ProfileSummary and computed
through the thin LTO indexing. This means that there is the limitation that the
scaled working set size is available to the thin LTO post link passes only.

[hjyamauchi]

clang-format.

[hjyamauchi]

I'm going to split this patch into parts:

• Use Module::setModuleFlag rather than addModuleFlag for the ProfileSummary metadata.
• Add the PartialProfileRatio field in ProfileSummary metadata.
• Compute the block count through thin LTO.
• The ProfileSummaryInfo change.

[hjyamauchi]

Rebase.

[mtrofin]

If the profile parser part were in a separate patch, would the size of the resulting patches be more balanced out? I'm wondering if splitting this into smaller patches would be meaningful.

llvm/include/llvm/IR/ProfileSummary.h

63

Can this be initialized, for maintainability (even if it's currently initialized in the ctor - code evolution may omit that later. Initializing at declaration reduces surprise bugs later)

llvm/lib/Analysis/ProfileSummaryInfo.cpp

82

How did we come by the 0.008 value?

llvm/lib/Bitcode/Reader/BitcodeAnalyzer.cpp

987 ↗

(On Diff #263766)

unrelated change (the last line)

llvm/lib/IR/ProfileSummary.cpp

118 ↗

(On Diff #263766)

Most of the code looks similar to the getVal above. Could some factoring reduce that duplication?

194 ↗

(On Diff #263766)

While at it - could we rename lower-case i to upper case something?

[hjyamauchi]

In D79831#2034601, @mtrofin wrote:

If the profile parser part were in a separate patch, would the size of the resulting patches be more balanced out? I'm wondering if splitting this into smaller patches would be meaningful.

Yup, splitting is the plan (see the earlier comment.)

[hjyamauchi]

llvm/include/llvm/IR/ProfileSummary.h

63

Done in D79951.

llvm/lib/Analysis/ProfileSummaryInfo.cpp

82

It came from tuning of large app / benchmarks so large working set size won't trigger for benchmarks.

llvm/lib/IR/ProfileSummary.cpp

118 ↗

(On Diff #263766)

Done in D79951.

194 ↗

(On Diff #263766)

Done in D79951.

[hjyamauchi]

Please review the split patches (child revisions) D79902 and D79951 first.

[hjyamauchi]

D79902 and D79951 have been checked in.
Split the part that computes the block count into D80403.

llvm/lib/Bitcode/Reader/BitcodeAnalyzer.cpp

987 ↗

(On Diff #263766)

Done in D80403.

[wmi]

llvm/include/llvm/IR/ProfileSummary.h

63

PartialProfileRatio is only available in ThinLTO mode. Better make it explicit in the comment.

PartialProfileRatio is set to "BlockCount / NumCounts". It doesn't match the ratio of "profile counters of the program being built" / "the number of profile counters in the partial sample profile", because for AFDO every block may have multiple counters inside of it. I think that is why you put "approximiately" in front of the comment.

Can we redefine PartialProfileRatio to "BlockCount / NumCounts * PartialSampleProfileWorkingSetSizeScaleFactor", so PartialProfileRatio can use the definition in the comment exactly instead of approximately? We may rename PartialSampleProfileWorkingSetSizeScaleFactor using something like NumofCounterInBlockScaleFactor. That will also give the flag PartialSampleProfileWorkingSetSizeScaleFactor more concrete meaning.

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

396 ↗

(On Diff #263766)

Can we use Index.addBlockCount(F.size())?

[hjyamauchi]

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

396 ↗

(On Diff #263766)

Done in D80403.

[hjyamauchi]

llvm/include/llvm/IR/ProfileSummary.h

63

Will add a comment about "only available in thin LTO mode."

It'd be nice to give precise meaning, but I wonder if those counts are that precise.

If NumofCounterInBlockScaleFactor == the number of counters per block, it'd have the value of 0.008, which is set by tuning, would the value/name make sense?

[wmi]

llvm/include/llvm/IR/ProfileSummary.h

63

Ok, the scale factor is not just the "factor of counters per block", but "factors of counters per block * (number of block s in working set / number of all the blocks of current program)". How about keep using PartialSampleProfileWorkingSetSizeScaleFactor but explain its meaning more clearly?

[hjyamauchi]

Rebase past D79902, D79951 and D80403.
Address comments.
Some cleanup.

[hjyamauchi]

llvm/include/llvm/IR/ProfileSummary.h

63

I thought about this more. I think that by and large, the small scale factor value is due to the fact that the hot working set sizes (NumCounts at the 99 percentile) in the PGO/PartialSamplePGO profiles are looking like ~1:100, respectively and the ProfileSummaryHugeWorkingSetSizeThreshold value was originally tuned for PGO. I guess the 0.008 value resulted from tuning to scale it down to the PGO level because the threshold values are shared.

I clarified the description of PartialSampleProfileWorkingSetSizeScaleFactor to include this and the "factor of counters per block".

Thoughts?

[hjyamauchi]

Disable the flag to be enabled later.

[wmi]

LGTM.

llvm/include/llvm/IR/ProfileSummary.h

63

It sounds good to me.

llvm/include/llvm/IR/Module.h

 
 namespace llvm {
 
 class Error;
 class FunctionType;
 class GVMaterializer;
 class LLVMContext;
 class MemoryBuffer;
+class ModuleSummaryIndex;
 class Pass;
 class RandomNumberGenerator;
 template <class PtrType> class SmallPtrSetImpl;
 class StructType;
 class VersionTuple;
 
 /// A Module instance is used to store all the information related to an
 /// LLVM module. Modules are the top level container of all other LLVM
/// @}
   /// Get the build SDK version metadata.
   ///
   /// An empty version is returned if no such metadata is attached.
   VersionTuple getSDKVersion() const;
   /// @}
 
   /// Take ownership of the given memory buffer.
   void setOwnedMemoryBuffer(std::unique_ptr<MemoryBuffer> MB);
+
+  /// Set the partial sample profile ratio in the profile summary module flag,
+  /// if applicable.
+  void setPartialSampleProfileRatio(const ModuleSummaryIndex &Index);
 };
 
 /// Given "llvm.used" or "llvm.compiler.used" as a global name, collect
 /// the initializer elements of that global in Set and return the global itself.
 GlobalVariable *collectUsedGlobalVariables(const Module &M,
                                            SmallPtrSetImpl<GlobalValue *> &Set,
                                            bool CompilerUsed);
 

llvm/include/llvm/IR/ProfileSummary.h

private:
   /// If 'Partial' is false, it means the profile being used to optimize
   /// a target is collected from the same target.
   /// If 'Partial' is true, it means the profile is for common/shared
   /// code. The common profile is usually merged from profiles collected
   /// from running other targets.
   bool Partial = false;
   /// This approximately represents the ratio of the number of profile counters
   /// of the program being built to the number of profile counters in the
-  /// partial sample profile. When 'Partial' is false, it is undefined.
+  /// partial sample profile. When 'Partial' is false, it is undefined. This is
+  /// currently only available under thin LTO mode.
   double PartialProfileRatio = 0;
   /// Return detailed summary as metadata.
   Metadata *getDetailedSummaryMD(LLVMContext &Context);
 
 public:
   static const int Scale = 1000000;
 
   ProfileSummary(Kind K, SummaryEntryVector DetailedSummary,

llvm/lib/Analysis/ProfileSummaryInfo.cpp

static cl::opt<int> ProfileSummaryColdCount(
     "profile-summary-cold-count", cl::ReallyHidden, cl::ZeroOrMore,
     cl::desc("A fixed cold count that overrides the count derived from"
              " profile-summary-cutoff-cold"));
 
 static cl::opt<bool> PartialProfile(
     "partial-profile", cl::Hidden, cl::init(false),
     cl::desc("Specify the current profile is used as a partial profile."));
 
+static cl::opt<double> ScalePartialSampleProfileWorkingSetSize(
+    "scale-partial-sample-profile-working-set-size", cl::Hidden, cl::init(true),
+    cl::desc(
+        "If true, scale the working set size of the partial sample profile "
+        "by the partial profile ratio to reflect the size of the program "
+        "being compiled."));
+
+static cl::opt<double> PartialSampleProfileWorkingSetSizeScaleFactor(
+    "partial-sample-profile-working-set-size-scale-factor", cl::Hidden,
+    cl::init(0.008),
+    cl::desc("The scale factor used to scale the working set size of the "
+             "partial sample profile along with the partial profile ratio. "
+             "This includes the factor of the profile counter per block "
+             "and the factor to scale the working set size to use the same "
+             "shared thresholds as PGO."));
+
 // Find the summary entry for a desired percentile of counts.
 static const ProfileSummaryEntry &getEntryForPercentile(SummaryEntryVector &DS,
                                                         uint64_t Percentile) {
   auto It = partition_point(DS, [=](const ProfileSummaryEntry &Entry) {
     return Entry.Cutoff < Percentile;
   });
   // The required percentile has to be <= one of the percentiles in the
   // detailed summary.
if (ProfileSummaryHotCount.getNumOccurrences() > 0)
     HotCountThreshold = ProfileSummaryHotCount;
   auto &ColdEntry =
       getEntryForPercentile(DetailedSummary, ProfileSummaryCutoffCold);
   ColdCountThreshold = ColdEntry.MinCount;
   if (ProfileSummaryColdCount.getNumOccurrences() > 0)
     ColdCountThreshold = ProfileSummaryColdCount;
   assert(ColdCountThreshold <= HotCountThreshold &&
          "Cold count threshold cannot exceed hot count threshold!");
+  if (!hasPartialSampleProfile() || !ScalePartialSampleProfileWorkingSetSize) {
     HasHugeWorkingSetSize =
         HotEntry.NumCounts > ProfileSummaryHugeWorkingSetSizeThreshold;
     HasLargeWorkingSetSize =
         HotEntry.NumCounts > ProfileSummaryLargeWorkingSetSizeThreshold;
+  } else {
+    // Scale the working set size of the partial sample profile to reflect the
+    // size of the program being compiled.
+    double PartialProfileRatio = Summary->getPartialProfileRatio();
+    uint64_t ScaledHotEntryNumCounts =
+        static_cast<uint64_t>(HotEntry.NumCounts * PartialProfileRatio *
+                              PartialSampleProfileWorkingSetSizeScaleFactor);
+    HasHugeWorkingSetSize =
+        ScaledHotEntryNumCounts > ProfileSummaryHugeWorkingSetSizeThreshold;
+    HasLargeWorkingSetSize =
+        ScaledHotEntryNumCounts > ProfileSummaryLargeWorkingSetSizeThreshold;
+  }
 }
 
 Optional<uint64_t>
 ProfileSummaryInfo::computeThreshold(int PercentileCutoff) const {
   if (!hasProfileSummary())
     return None;
   auto iter = ThresholdCache.find(PercentileCutoff);
   if (iter != ThresholdCache.end()) {

llvm/lib/IR/Module.cpp

 #include "llvm/IR/Function.h"
 #include "llvm/IR/GVMaterializer.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalIFunc.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
+#include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/IR/SymbolTableListTraits.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/TypeFinder.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueSymbolTable.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CodeGen.h"
GlobalVariable *llvm::collectUsedGlobalVariables(
 
   const ConstantArray *Init = cast<ConstantArray>(GV->getInitializer());
   for (Value *Op : Init->operands()) {
     GlobalValue *G = cast<GlobalValue>(Op->stripPointerCasts());
     Set.insert(G);
   }
   return GV;
 }
+
+void Module::setPartialSampleProfileRatio(const ModuleSummaryIndex &Index) {
+  if (auto *SummaryMD = getProfileSummary(/*IsCS*/ false)) {
+    std::unique_ptr<ProfileSummary> ProfileSummary(
+        ProfileSummary::getFromMD(SummaryMD));
+    if (ProfileSummary) {
+      if (ProfileSummary->getKind() != ProfileSummary::PSK_Sample ||
+          !ProfileSummary->isPartialProfile())
+        return;
+      uint64_t BlockCount = Index.getBlockCount();
+      uint32_t NumCounts = ProfileSummary->getNumCounts();
+      if (!NumCounts)
+        return;
+      double Ratio = (double)BlockCount / NumCounts;
+      ProfileSummary->setPartialProfileRatio(Ratio);
+      setProfileSummary(ProfileSummary->getMD(getContext()),
+                        ProfileSummary::PSK_Sample);
+    }
+  }
+}

llvm/lib/LTO/LTOBackend.cpp

Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   // Setup optimization remarks.
   auto DiagFileOrErr = lto::setupLLVMOptimizationRemarks(
       Mod.getContext(), Conf.RemarksFilename, Conf.RemarksPasses,
       Conf.RemarksFormat, Conf.RemarksWithHotness, Task);
   if (!DiagFileOrErr)
     return DiagFileOrErr.takeError();
   auto DiagnosticOutputFile = std::move(*DiagFileOrErr);
 
+  // Set the partial sample profile ratio in the profile summary module flag of
+  // the module, if applicable.
+  Mod.setPartialSampleProfileRatio(CombinedIndex);
+
   if (Conf.CodeGenOnly) {
     codegen(Conf, TM.get(), AddStream, Task, Mod);
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
   }
 
   if (Conf.PreOptModuleHook && !Conf.PreOptModuleHook(Task, Mod))
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 

llvm/lib/Transforms/IPO/FunctionImport.cpp

for (GlobalAlias &GA : SrcModule->aliases()) {
         GlobalsToImport.insert(Fn);
       }
     }
 
     // Upgrade debug info after we're done materializing all the globals and we
     // have loaded all the required metadata!
     UpgradeDebugInfo(*SrcModule);
 
+    // Set the partial sample profile ratio in the profile summary module flag
+    // of the imported source module, if applicable, so that the profile summary
+    // module flag will match with that of the destination module when it's
+    // imported.
+    SrcModule->setPartialSampleProfileRatio(Index);
+
     // Link in the specified functions.
     if (renameModuleForThinLTO(*SrcModule, Index, ClearDSOLocalOnDeclarations,
                                &GlobalsToImport))
       return true;
 
     if (PrintImports) {
       for (const auto *GV : GlobalsToImport)
         dbgs() << DestModule.getSourceFileName() << ": Import " << GV->getName()

llvm/unittests/Analysis/ProfileSummaryInfoTest.cpp

ProfileSummaryInfo buildPSI(Module *M) {
     return ProfileSummaryInfo(*M);
   }
   BlockFrequencyInfo buildBFI(Function &F) {
     DT.reset(new DominatorTree(F));
     LI.reset(new LoopInfo(*DT));
     BPI.reset(new BranchProbabilityInfo(F, *LI));
     return BlockFrequencyInfo(F, *BPI, *LI);
   }
-  std::unique_ptr<Module> makeLLVMModule(const char *ProfKind = nullptr) {
+  std::unique_ptr<Module> makeLLVMModule(const char *ProfKind = nullptr,
+                                         uint64_t NumCounts = 3,
+                                         uint64_t IsPartialProfile = 0,
+                                         double PartialProfileRatio = 0.0,
+                                         uint64_t HotNumCounts = 3,
+                                         uint64_t ColdNumCounts = 10) {
     const char *ModuleString =
         "define i32 @g(i32 %x) !prof !21 {{\n"
         "  ret i32 0\n"
         "}\n"
         "define i32 @h(i32 %x) !prof !22 {{\n"
         "  ret i32 0\n"
         "}\n"
         "define i32 @f(i32 %x) !prof !20 {{\n"
const char *ModuleString =
         "  %y2 = phi i32 [0, %bb1], [1, %bb2] \n"
         "  ret i32 %y2\n"
         "}\n"
         "!20 = !{{!\"function_entry_count\", i64 400}\n"
         "!21 = !{{!\"function_entry_count\", i64 1}\n"
         "!22 = !{{!\"function_entry_count\", i64 100}\n"
         "!23 = !{{!\"branch_weights\", i32 64, i32 4}\n"
         "{0}";
-    const char *SummaryString = "!llvm.module.flags = !{{!1}"
-                                "!1 = !{{i32 1, !\"ProfileSummary\", !2}"
-                                "!2 = !{{!3, !4, !5, !6, !7, !8, !9, !10}"
-                                "!3 = !{{!\"ProfileFormat\", !\"{0}\"}"
-                                "!4 = !{{!\"TotalCount\", i64 10000}"
-                                "!5 = !{{!\"MaxCount\", i64 10}"
-                                "!6 = !{{!\"MaxInternalCount\", i64 1}"
-                                "!7 = !{{!\"MaxFunctionCount\", i64 1000}"
-                                "!8 = !{{!\"NumCounts\", i64 3}"
-                                "!9 = !{{!\"NumFunctions\", i64 3}"
-                                "!10 = !{{!\"DetailedSummary\", !11}"
-                                "!11 = !{{!12, !13, !14}"
-                                "!12 = !{{i32 10000, i64 1000, i32 1}"
-                                "!13 = !{{i32 999000, i64 300, i32 3}"
-                                "!14 = !{{i32 999999, i64 5, i32 10}";
+    const char *SummaryString =
+        "!llvm.module.flags = !{{!1}\n"
+        "!1 = !{{i32 1, !\"ProfileSummary\", !2}\n"
+        "!2 = !{{!3, !4, !5, !6, !7, !8, !9, !10, !11, !12}\n"
+        "!3 = !{{!\"ProfileFormat\", !\"{0}\"}\n"
+        "!4 = !{{!\"TotalCount\", i64 10000}\n"
+        "!5 = !{{!\"MaxCount\", i64 10}\n"
+        "!6 = !{{!\"MaxInternalCount\", i64 1}\n"
+        "!7 = !{{!\"MaxFunctionCount\", i64 1000}\n"
+        "!8 = !{{!\"NumCounts\", i64 {1}}\n"
+        "!9 = !{{!\"NumFunctions\", i64 3}\n"
+        "!10 = !{{!\"IsPartialProfile\", i64 {2}}\n"
+        "!11 = !{{!\"PartialProfileRatio\", double {3}}\n"
+        "!12 = !{{!\"DetailedSummary\", !13}\n"
+        "!13 = !{{!14, !15, !16}\n"
+        "!14 = !{{i32 10000, i64 1000, i32 1}\n"
+        "!15 = !{{i32 990000, i64 300, i32 {4}}\n"
+        "!16 = !{{i32 999999, i64 5, i32 {5}}\n";
     SMDiagnostic Err;
-    if (ProfKind)
-      return parseAssemblyString(
-          formatv(ModuleString, formatv(SummaryString, ProfKind).str()).str(),
-          Err, C);
-    else
+    if (ProfKind) {
+      auto Summary =
+          formatv(SummaryString, ProfKind, NumCounts, IsPartialProfile,
+                  PartialProfileRatio, HotNumCounts, ColdNumCounts)
+              .str();
+      return parseAssemblyString(formatv(ModuleString, Summary).str(), Err, C);
+    } else
       return parseAssemblyString(formatv(ModuleString, "").str(), Err, C);
   }
 };
 
 TEST_F(ProfileSummaryInfoTest, TestNoProfile) {
   auto M = makeLLVMModule(/*ProfKind=*/nullptr);
   Function *F = M->getFunction("f");
 

 }
 
 TEST_F(ProfileSummaryInfoTest, SampleProf) {
   auto M = makeLLVMModule("SampleProfile");
   Function *F = M->getFunction("f");
   ProfileSummaryInfo PSI = buildPSI(M.get());
   EXPECT_TRUE(PSI.hasProfileSummary());
   EXPECT_TRUE(PSI.hasSampleProfile());
+  EXPECT_FALSE(PSI.hasPartialSampleProfile());
 
   BasicBlock &BB0 = F->getEntryBlock();
   BasicBlock *BB1 = BB0.getTerminator()->getSuccessor(0);
   BasicBlock *BB2 = BB0.getTerminator()->getSuccessor(1);
   BasicBlock *BB3 = BB1->getSingleSuccessor();
 
   BlockFrequencyInfo BFI = buildBFI(*F);
   EXPECT_TRUE(PSI.isHotBlock(&BB0, &BFI));
TEST_F(ProfileSummaryInfoTest, SampleProfNoFuncEntryCount) {
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB1, &BFI));
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB2, &BFI));
   EXPECT_FALSE(PSI.isColdBlockNthPercentile(990000, BB3, &BFI));
 
   EXPECT_FALSE(PSI.isFunctionHotInCallGraphNthPercentile(990000, F, BFI));
   EXPECT_FALSE(PSI.isFunctionColdInCallGraphNthPercentile(990000, F, BFI));
 }
 
+TEST_F(ProfileSummaryInfoTest, PartialSampleProfWorkingSetSize) {
+  // With PartialProfileRatio unset (zero.)
+  auto M1 = makeLLVMModule("SampleProfile", /*NumCounts*/ 3,
+                           /*IsPartialProfile*/ 1,
+                           /*PartialProfileRatio*/ 0.0,
+                           /*HotNumCounts*/ 3, /*ColdNumCounts*/ 10);
+  ProfileSummaryInfo PSI1 = buildPSI(M1.get());
+  EXPECT_TRUE(PSI1.hasProfileSummary());
+  EXPECT_TRUE(PSI1.hasSampleProfile());
+  EXPECT_TRUE(PSI1.hasPartialSampleProfile());
+  EXPECT_FALSE(PSI1.hasHugeWorkingSetSize());
+  EXPECT_FALSE(PSI1.hasLargeWorkingSetSize());
+
+  // With PartialProfileRatio set (non-zero) and a small working set size.
+  auto M2 = makeLLVMModule("SampleProfile", /*NumCounts*/ 27493235,
+                           /*IsPartialProfile*/ 1,
+                           /*PartialProfileRatio*/ 0.00000012,
+                           /*HotNumCounts*/ 3102082,
+                           /*ColdNumCounts*/ 18306149);
+  ProfileSummaryInfo PSI2 = buildPSI(M2.get());
+  EXPECT_TRUE(PSI2.hasProfileSummary());
+  EXPECT_TRUE(PSI2.hasSampleProfile());
+  EXPECT_TRUE(PSI2.hasPartialSampleProfile());
+  EXPECT_FALSE(PSI2.hasHugeWorkingSetSize());
+  EXPECT_FALSE(PSI2.hasLargeWorkingSetSize());
+
+  // With PartialProfileRatio is set (non-zero) and a large working set size.
+  auto M3 = makeLLVMModule("SampleProfile", /*NumCounts*/ 27493235,
+                           /*IsPartialProfile*/ 1,
+                           /*PartialProfileRatio*/ 0.9,
+                           /*HotNumCounts*/ 3102082,
+                           /*ColdNumCounts*/ 18306149);
+  ProfileSummaryInfo PSI3 = buildPSI(M3.get());
+  EXPECT_TRUE(PSI3.hasProfileSummary());
+  EXPECT_TRUE(PSI3.hasSampleProfile());
+  EXPECT_TRUE(PSI3.hasPartialSampleProfile());
+  EXPECT_TRUE(PSI3.hasHugeWorkingSetSize());
+  EXPECT_TRUE(PSI3.hasLargeWorkingSetSize());
+}
+
 } // end anonymous namespace
 } // end namespace llvm

llvm/unittests/IR/ModuleTest.cpp

 //===- unittests/IR/ModuleTest.cpp - Module unit tests --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Module.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/RandomNumberGenerator.h"
 #include "gtest/gtest.h"
 
 #include <random>
 
 using namespace llvm;
 
TEST(ModuleTest, setProfileSummary) {
   M->setProfileSummary(PS->getMD(Context), ProfileSummary::PSK_Sample);
   delete PS;
   PS = ProfileSummary::getFromMD(M->getProfileSummary(/*IsCS*/ false));
   EXPECT_NE(nullptr, PS);
   EXPECT_EQ(true, PS->isPartialProfile());
   delete PS;
 }
 
+
+TEST(ModuleTest, setPartialSampleProfileRatio) {
+  const char *IRString = R"IR(
+  !llvm.module.flags = !{!0}
+
+  !0 = !{i32 1, !"ProfileSummary", !1}
+  !1 = !{!2, !3, !4, !5, !6, !7, !8, !9, !10, !11}
+  !2 = !{!"ProfileFormat", !"SampleProfile"}
+  !3 = !{!"TotalCount", i64 10000}
+  !4 = !{!"MaxCount", i64 10}
+  !5 = !{!"MaxInternalCount", i64 1}
+  !6 = !{!"MaxFunctionCount", i64 1000}
+  !7 = !{!"NumCounts", i64 200}
+  !8 = !{!"NumFunctions", i64 3}
+  !9 = !{!"IsPartialProfile", i64 1}
+  !10 = !{!"PartialProfileRatio", double 0.0}
+  !11 = !{!"DetailedSummary", !12}
+  !12 = !{!13, !14, !15}
+  !13 = !{i32 10000, i64 1000, i32 1}
+  !14 = !{i32 990000, i64 300, i32 10}
+  !15 = !{i32 999999, i64 5, i32 100}
+  )IR";
+
+  SMDiagnostic Err;
+  LLVMContext Context;
+  std::unique_ptr<Module> M = parseAssemblyString(IRString, Err, Context);
+  ModuleSummaryIndex Index(/*HaveGVs*/ false);
+  const unsigned BlockCount = 100;
+  const unsigned NumCounts = 200;
+  Index.setBlockCount(BlockCount);
+  M->setPartialSampleProfileRatio(Index);
+  double Ratio = (double)BlockCount / NumCounts;
+  std::unique_ptr<ProfileSummary> ProfileSummary(
+      ProfileSummary::getFromMD(M->getProfileSummary(/*IsCS*/ false)));
+  EXPECT_EQ(Ratio, ProfileSummary->getPartialProfileRatio());
+}
+
 } // end namespace