[ProfileSummary] Add partial profile annotation on IR.

Profile and profile summary are usually read only once and then annotated
on IR. The profile summary metadata on IR should include the value of the
newly added partial profile flag, so that compilation phase like thinlto
postlink can get the full set of profile information.

Differential Revision: https://reviews.llvm.org/D78310
diff --git a/llvm/include/llvm/IR/ProfileSummary.h b/llvm/include/llvm/IR/ProfileSummary.h
index 7532b44..b105f8d 100644
--- a/llvm/include/llvm/IR/ProfileSummary.h
+++ b/llvm/include/llvm/IR/ProfileSummary.h
@@ -64,15 +64,16 @@
   ProfileSummary(Kind K, SummaryEntryVector DetailedSummary,
                  uint64_t TotalCount, uint64_t MaxCount,
                  uint64_t MaxInternalCount, uint64_t MaxFunctionCount,
-                 uint32_t NumCounts, uint32_t NumFunctions)
+                 uint32_t NumCounts, uint32_t NumFunctions,
+                 bool Partial = false)
       : PSK(K), DetailedSummary(std::move(DetailedSummary)),
         TotalCount(TotalCount), MaxCount(MaxCount),
         MaxInternalCount(MaxInternalCount), MaxFunctionCount(MaxFunctionCount),
-        NumCounts(NumCounts), NumFunctions(NumFunctions) {}
+        NumCounts(NumCounts), NumFunctions(NumFunctions), Partial(Partial) {}
 
   Kind getKind() const { return PSK; }
   /// Return summary information as metadata.
-  Metadata *getMD(LLVMContext &Context);
+  Metadata *getMD(LLVMContext &Context, bool AddPartialField = true);
   /// Construct profile summary from metdata.
   static ProfileSummary *getFromMD(Metadata *MD);
   SummaryEntryVector &getDetailedSummary() { return DetailedSummary; }
diff --git a/llvm/lib/IR/ProfileSummary.cpp b/llvm/lib/IR/ProfileSummary.cpp
index 11d95ac..d6c3c50 100644
--- a/llvm/lib/IR/ProfileSummary.cpp
+++ b/llvm/lib/IR/ProfileSummary.cpp
@@ -65,18 +65,24 @@
 // "ProfileFormat" and a string representing the format ("InstrProf" or
 // "SampleProfile"). The rest of the elements of the outer MDTuple are specific
 // to the kind of profile summary as returned by getFormatSpecificMD.
-Metadata *ProfileSummary::getMD(LLVMContext &Context) {
+// IsPartialProfile is an optional field and \p AddPartialField will decide
+// whether to add a field for it.
+Metadata *ProfileSummary::getMD(LLVMContext &Context, bool AddPartialField) {
   const char *KindStr[3] = {"InstrProf", "CSInstrProf", "SampleProfile"};
-  Metadata *Components[] = {
-    getKeyValMD(Context, "ProfileFormat", KindStr[PSK]),
-    getKeyValMD(Context, "TotalCount", getTotalCount()),
-    getKeyValMD(Context, "MaxCount", getMaxCount()),
-    getKeyValMD(Context, "MaxInternalCount", getMaxInternalCount()),
-    getKeyValMD(Context, "MaxFunctionCount", getMaxFunctionCount()),
-    getKeyValMD(Context, "NumCounts", getNumCounts()),
-    getKeyValMD(Context, "NumFunctions", getNumFunctions()),
-    getDetailedSummaryMD(Context),
-  };
+  SmallVector<Metadata *, 16> Components;
+  Components.push_back(getKeyValMD(Context, "ProfileFormat", KindStr[PSK]));
+  Components.push_back(getKeyValMD(Context, "TotalCount", getTotalCount()));
+  Components.push_back(getKeyValMD(Context, "MaxCount", getMaxCount()));
+  Components.push_back(
+      getKeyValMD(Context, "MaxInternalCount", getMaxInternalCount()));
+  Components.push_back(
+      getKeyValMD(Context, "MaxFunctionCount", getMaxFunctionCount()));
+  Components.push_back(getKeyValMD(Context, "NumCounts", getNumCounts()));
+  Components.push_back(getKeyValMD(Context, "NumFunctions", getNumFunctions()));
+  if (AddPartialField)
+    Components.push_back(
+        getKeyValMD(Context, "IsPartialProfile", isPartialProfile()));
+  Components.push_back(getDetailedSummaryMD(Context));
   return MDTuple::get(Context, Components);
 }
 
@@ -141,10 +147,11 @@
 
 ProfileSummary *ProfileSummary::getFromMD(Metadata *MD) {
   MDTuple *Tuple = dyn_cast_or_null<MDTuple>(MD);
-  if (!Tuple || Tuple->getNumOperands() != 8)
+  if (!Tuple && (Tuple->getNumOperands() < 8 || Tuple->getNumOperands() > 9))
     return nullptr;
 
-  auto &FormatMD = Tuple->getOperand(0);
+  int i = 0;
+  auto &FormatMD = Tuple->getOperand(i++);
   ProfileSummary::Kind SummaryKind;
   if (isKeyValuePair(dyn_cast_or_null<MDTuple>(FormatMD), "ProfileFormat",
                      "SampleProfile"))
@@ -160,27 +167,41 @@
 
   uint64_t NumCounts, TotalCount, NumFunctions, MaxFunctionCount, MaxCount,
       MaxInternalCount;
-  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(1)), "TotalCount",
+  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(i++)), "TotalCount",
               TotalCount))
     return nullptr;
-  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(2)), "MaxCount", MaxCount))
+  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(i++)), "MaxCount", MaxCount))
     return nullptr;
-  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(3)), "MaxInternalCount",
+  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(i++)), "MaxInternalCount",
               MaxInternalCount))
     return nullptr;
-  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(4)), "MaxFunctionCount",
+  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(i++)), "MaxFunctionCount",
               MaxFunctionCount))
     return nullptr;
-  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(5)), "NumCounts", NumCounts))
+  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(i++)), "NumCounts",
+              NumCounts))
     return nullptr;
-  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(6)), "NumFunctions",
+  if (!getVal(dyn_cast<MDTuple>(Tuple->getOperand(i++)), "NumFunctions",
               NumFunctions))
     return nullptr;
+  // Initialize IsPartialProfile because the field is optional.
+  uint64_t IsPartialProfile = 0;
+
+  // IsPartialProfile is optional so it doesn't matter even if the next val
+  // is not IsPartialProfile.
+  if (getVal(dyn_cast<MDTuple>(Tuple->getOperand(i)), "IsPartialProfile",
+             IsPartialProfile)) {
+    // Need to make sure when IsPartialProfile is presented, we won't step
+    // over the bound of Tuple operand array.
+    if (Tuple->getNumOperands() < 9)
+      return nullptr;
+    i++;
+  }
 
   SummaryEntryVector Summary;
-  if (!getSummaryFromMD(dyn_cast<MDTuple>(Tuple->getOperand(7)), Summary))
+  if (!getSummaryFromMD(dyn_cast<MDTuple>(Tuple->getOperand(i++)), Summary))
     return nullptr;
   return new ProfileSummary(SummaryKind, std::move(Summary), TotalCount,
                             MaxCount, MaxInternalCount, MaxFunctionCount,
-                            NumCounts, NumFunctions);
+                            NumCounts, NumFunctions, IsPartialProfile);
 }
diff --git a/llvm/test/Transforms/PGOProfile/cspgo_profile_summary.ll b/llvm/test/Transforms/PGOProfile/cspgo_profile_summary.ll
index b309833..8bd2fb9 100644
--- a/llvm/test/Transforms/PGOProfile/cspgo_profile_summary.ll
+++ b/llvm/test/Transforms/PGOProfile/cspgo_profile_summary.ll
@@ -142,7 +142,7 @@
 ; CSPGOSUMMARY: {{![0-9]+}} = !{!"MaxFunctionCount", i64 800000}
 ; CSPGOSUMMARY: {{![0-9]+}} = !{!"NumCounts", i64 14}
 ; CSPGOSUMMARY: {{![0-9]+}} = !{!"NumFunctions", i64 8}
-; CSPGOSUMMARY: {{![0-9]+}} = !{!"DetailedSummary", !10}
+; CSPGOSUMMARY: {{![0-9]+}} = !{!"DetailedSummary", !{{[0-9]+}}}
 ; CSPGOSUMMARY: {{![0-9]+}} = !{i32 1, !"CSProfileSummary", !{{[0-9]+}}}
 ; CSPGOSUMMARY: {{![0-9]+}} = !{!"ProfileFormat", !"CSInstrProf"}
 ; CSPGOSUMMARY: {{![0-9]+}} = !{!"TotalCount", i64 1299950}
diff --git a/llvm/test/Transforms/PGOProfile/unreachable_bb.ll b/llvm/test/Transforms/PGOProfile/unreachable_bb.ll
index 1b7fe4b..96aca89 100644
--- a/llvm/test/Transforms/PGOProfile/unreachable_bb.ll
+++ b/llvm/test/Transforms/PGOProfile/unreachable_bb.ll
@@ -16,7 +16,7 @@
 declare void @bar()
 
 ;USE: !0 = !{i32 1, !"ProfileSummary", !1}
-;USE: !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+;USE: !1 = !{!2, !3, !4, !5, !6, !7, !8, !9, !10}
 ;USE: !2 = !{!"ProfileFormat", !"InstrProf"}
 ;USE: !3 = !{!"TotalCount", i64 0}
 
diff --git a/llvm/unittests/ProfileData/SampleProfTest.cpp b/llvm/unittests/ProfileData/SampleProfTest.cpp
index 3aaba05..5421882 100644
--- a/llvm/unittests/ProfileData/SampleProfTest.cpp
+++ b/llvm/unittests/ProfileData/SampleProfTest.cpp
@@ -77,6 +77,62 @@
     OS->close();
   }
 
+  // Verify profile summary is consistent in the roundtrip to and from
+  // Metadata. \p AddPartialField is to choose whether the Metadata
+  // contains the IsPartialProfile field which is optional.
+  void verifyProfileSummary(ProfileSummary &Summary, Module &M,
+                            const bool AddPartialField) {
+    LLVMContext &Context = M.getContext();
+    const bool IsPartialProfile = Summary.isPartialProfile();
+    auto VerifySummary = [IsPartialProfile](ProfileSummary &Summary) mutable {
+      ASSERT_EQ(ProfileSummary::PSK_Sample, Summary.getKind());
+      ASSERT_EQ(137392u, Summary.getTotalCount());
+      ASSERT_EQ(8u, Summary.getNumCounts());
+      ASSERT_EQ(4u, Summary.getNumFunctions());
+      ASSERT_EQ(1437u, Summary.getMaxFunctionCount());
+      ASSERT_EQ(60351u, Summary.getMaxCount());
+      ASSERT_EQ(IsPartialProfile, Summary.isPartialProfile());
+
+      uint32_t Cutoff = 800000;
+      auto Predicate = [&Cutoff](const ProfileSummaryEntry &PE) {
+        return PE.Cutoff == Cutoff;
+      };
+      std::vector<ProfileSummaryEntry> &Details = Summary.getDetailedSummary();
+      auto EightyPerc = find_if(Details, Predicate);
+      Cutoff = 900000;
+      auto NinetyPerc = find_if(Details, Predicate);
+      Cutoff = 950000;
+      auto NinetyFivePerc = find_if(Details, Predicate);
+      Cutoff = 990000;
+      auto NinetyNinePerc = find_if(Details, Predicate);
+      ASSERT_EQ(60000u, EightyPerc->MinCount);
+      ASSERT_EQ(12557u, NinetyPerc->MinCount);
+      ASSERT_EQ(12557u, NinetyFivePerc->MinCount);
+      ASSERT_EQ(610u, NinetyNinePerc->MinCount);
+    };
+    VerifySummary(Summary);
+
+    // Test that conversion of summary to and from Metadata works.
+    Metadata *MD = Summary.getMD(Context, AddPartialField);
+    ASSERT_TRUE(MD);
+    ProfileSummary *PS = ProfileSummary::getFromMD(MD);
+    ASSERT_TRUE(PS);
+    VerifySummary(*PS);
+    delete PS;
+
+    // Test that summary can be attached to and read back from module.
+    PS = ProfileSummary::getFromMD(MD);
+
+    M.eraseNamedMetadata(M.getOrInsertModuleFlagsMetadata());
+    M.setProfileSummary(MD, ProfileSummary::PSK_Sample);
+    MD = M.getProfileSummary(/* IsCS */ false);
+    ASSERT_TRUE(MD);
+    PS = ProfileSummary::getFromMD(MD);
+    ASSERT_TRUE(PS);
+    VerifySummary(*PS);
+    delete PS;
+  }
+
   void testRoundTrip(SampleProfileFormat Format, bool Remap, bool UseMD5) {
     SmallVector<char, 128> ProfilePath;
     ASSERT_TRUE(NoError(llvm::sys::fs::createTemporaryFile("profile", "", ProfilePath)));
@@ -214,51 +270,16 @@
     ASSERT_EQ(1000u, CTMap.get()[MconstructRep]);
     ASSERT_EQ(437u, CTMap.get()[StringviewRep]);
 
-    auto VerifySummary = [](ProfileSummary &Summary) mutable {
-      ASSERT_EQ(ProfileSummary::PSK_Sample, Summary.getKind());
-      ASSERT_EQ(137392u, Summary.getTotalCount());
-      ASSERT_EQ(8u, Summary.getNumCounts());
-      ASSERT_EQ(4u, Summary.getNumFunctions());
-      ASSERT_EQ(1437u, Summary.getMaxFunctionCount());
-      ASSERT_EQ(60351u, Summary.getMaxCount());
-
-      uint32_t Cutoff = 800000;
-      auto Predicate = [&Cutoff](const ProfileSummaryEntry &PE) {
-        return PE.Cutoff == Cutoff;
-      };
-      std::vector<ProfileSummaryEntry> &Details = Summary.getDetailedSummary();
-      auto EightyPerc = find_if(Details, Predicate);
-      Cutoff = 900000;
-      auto NinetyPerc = find_if(Details, Predicate);
-      Cutoff = 950000;
-      auto NinetyFivePerc = find_if(Details, Predicate);
-      Cutoff = 990000;
-      auto NinetyNinePerc = find_if(Details, Predicate);
-      ASSERT_EQ(60000u, EightyPerc->MinCount);
-      ASSERT_EQ(12557u, NinetyPerc->MinCount);
-      ASSERT_EQ(12557u, NinetyFivePerc->MinCount);
-      ASSERT_EQ(610u, NinetyNinePerc->MinCount);
-    };
 
     ProfileSummary &Summary = Reader->getSummary();
-    VerifySummary(Summary);
+    Summary.setPartialProfile(true);
+    verifyProfileSummary(Summary, M, true);
 
-    // Test that conversion of summary to and from Metadata works.
-    Metadata *MD = Summary.getMD(Context);
-    ASSERT_TRUE(MD);
-    ProfileSummary *PS = ProfileSummary::getFromMD(MD);
-    ASSERT_TRUE(PS);
-    VerifySummary(*PS);
-    delete PS;
+    Summary.setPartialProfile(false);
+    verifyProfileSummary(Summary, M, true);
 
-    // Test that summary can be attached to and read back from module.
-    M.setProfileSummary(MD, ProfileSummary::PSK_Sample);
-    MD = M.getProfileSummary(/* IsCS */ false);
-    ASSERT_TRUE(MD);
-    PS = ProfileSummary::getFromMD(MD);
-    ASSERT_TRUE(PS);
-    VerifySummary(*PS);
-    delete PS;
+    Summary.setPartialProfile(false);
+    verifyProfileSummary(Summary, M, false);
   }
 
   void addFunctionSamples(StringMap<FunctionSamples> *Smap, const char *Fname,