[llvm-profdata] Use flattening sample profile in profile supplementation

[xur]

Summary

We need to flatten the SampleFDO profile in profile supplementation
because the InstrFDO profile does not have inlined callsite counters.
Without flattening profile, FDO optimizations are not stable:
we will not supplement the second generation profile when the modified
functions are all inlined.

This patch fixes this issue: we will flatten the profile for functions
that appears in FDO profile.

Note that we only need to find the hot/warm functions in SampleFDO
profile, so we will not perform a full flatten. We will use
a DFS traversal to compute the accumulated entry count and max bodycount.
This is much cheaper than full flattening.

[xur]

Simplified the patch a bit by refactoring the existing interface.

[davidxl]

llvm/include/llvm/ProfileData/SampleProf.h

928–929

The comment needs to be updated to reflect the new behavior.

llvm/test/tools/llvm-profdata/suppl-instr-with-sample-flatten.test

5

The test only sets bar.cc::bar's counter, not goo's?

llvm/tools/llvm-profdata/llvm-profdata.cpp

637

pre-inlined, do you mean early-inlined?

641

Perhaps describe with a short example.

663

There is no need for the nested lambda.

[xur]

llvm/include/llvm/ProfileData/SampleProf.h

928–929

Ack.

llvm/test/tools/llvm-profdata/suppl-instr-with-sample-flatten.test

5

Yes. That's a typo. I meant "bar.cc:bar". Function "goo" is considered preinlined in FDO -- as it's not appear in FDO profile. We should not supplement goo.

llvm/tools/llvm-profdata/llvm-profdata.cpp

637

I think we call "early-inline" pre-inliner in llvm.

641

I will use the test case as the example. To be updated.

663

I think I need to use two level lambda if we want to use lamda. I can do with a member function but I need to pass many things around.

[xur]

Addressed David's review comments.

[xur]

Removed lambda updateSampleMap() -- it was from the early version and now it was no longer needed as we just had one caller.

[davidxl]

llvm/tools/llvm-profdata/llvm-profdata.cpp

691

I am not sure if RootName parameter is actually needed here. In the code below, it is either be overwritten (with newroot), or used in cases we don't care -- i.e., function not found in instr profile.

[xur]

llvm/tools/llvm-profdata/llvm-profdata.cpp

691

It's used in line 696 (set as the default value) of NewRootName.
If we find the Name is in InstrProfile, it will be overwriteen.

If we don't find the Name, we don't care this function (we treat it as inline function), but we still interested in the max body count and we will update the RootName entry from the parameter.

In the example, when get the goo's FS, it's max count will used to update bar.c:bar() (this is the RootName).

[davidxl]

lgtm

llvm/include/llvm/ProfileData/SampleProf.h

public:
 
   /// Return all the samples collected in the body of the function.
   const BodySampleMap &getBodySamples() const { return BodySamples; }
 
   /// Return all the callsite samples collected in the body of the function.
   const CallsiteSampleMap &getCallsiteSamples() const {
     return CallsiteSamples;
   }
 
-  /// Return the maximum of sample counts in a function body including functions
-  /// inlined in it.
-  uint64_t getMaxCountInside() const {
+  /// Return the maximum of sample counts in a function body. When SkipCallSite
+  /// is false, which is the default, the return count includes samples in the
+  /// inlined functions. When SkipCallSite is true, the return count only
+  /// considers the body samples.
+  uint64_t getMaxCountInside(bool SkipCallSite = false) const {
     uint64_t MaxCount = 0;
     for (const auto &L : getBodySamples())
       MaxCount = std::max(MaxCount, L.second.getSamples());
+    if (SkipCallSite)
+      return MaxCount;
     for (const auto &C : getCallsiteSamples())
       for (const FunctionSamplesMap::value_type &F : C.second)
         MaxCount = std::max(MaxCount, F.second.getMaxCountInside());
     return MaxCount;
   }
 
   /// Merge the samples in \p Other into this one.
   /// Optionally scale samples by \p Weight.

llvm/test/tools/llvm-profdata/Inputs/flatten_instr.proftext

+# IR level Instrumentation Flag
+:ir
+# Always instrument the function entry block
+:entry_first
+foo
+# Func Hash:
+1111
+# Num Counters:
+5
+# Counter Values:
+10000
+50
+2000
+40
+6000
+
+bar.cc:bar
+# Func Hash:
+2222
+# Num Counters:
+10
+# Counter Values:
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0

llvm/test/tools/llvm-profdata/Inputs/flatten_sample.proftext

+foo:12345:1000
+ 1: 1000
+ 2.1: 1000
+ 15: 5000
+ 4: bar:1000
+  1: 1000
+  2: goo:3000
+   1: 3000
+ 8: bar:40000
+  1: 10000
+  2: goo:30000
+   1: 30000

llvm/test/tools/llvm-profdata/suppl-instr-with-sample-flatten.test

+Some basic tests for supplementing instrumentation profile with sample profile
+with flattening.
+
+Test bar.cc:bar's counters will be set to PseudoHot.
+RUN: llvm-profdata merge \
+RUN:     -supplement-instr-with-sample=%p/Inputs/flatten_sample.proftext \
+RUN:     %p/Inputs/flatten_instr.proftext -o %t
+RUN: llvm-profdata show %t -all-functions -counts | FileCheck %s --check-prefix=FLATTEN
+
+FLATTEN:  bar.cc:bar:
+FLATTEN-NEXT:    Hash: 0x00000000000008ae
+FLATTEN-NEXT:    Counters: 10    <PseudoHot>
+FLATTEN:  foo:
+FLATTEN-NEXT:    Hash: 0x0000000000000457
+FLATTEN-NEXT:    Counters: 5
+FLATTEN-NEXT:    Block counts: [10000, 50, 2000, 40, 6000]
+FLATTEN-NOT:  goo:

llvm/tools/llvm-profdata/llvm-profdata.cpp

auto buildStaticFuncMap = [&StaticFuncMap,
 
     if (StaticFuncMap.find(NewName) == StaticFuncMap.end()) {
       StaticFuncMap[NewName] = Name;
     } else {
       StaticFuncMap[NewName] = DuplicateNameStr;
     }
   };
 
+  // We need to flatten the SampleFDO profile as the InstrFDO
+  // profile does not have inlined callsite profiles.
+  // One caveat is the pre-inlined function -- their samples
+  // should be collapsed into the caller function.
+  // Here we do a DFS traversal to get the flatten profile
+  // info: the sum of entrycount and the max of maxcount.
+  // Here is the algorithm:
+  //   recursive (FS, root_name) {
+  //      name = FS->getName();
+  //      get samples for FS;
+  //      if (InstrProf.find(name) {
+  //        root_name = name;
+  //      } else {
+  //        if (name is in static_func map) {
+  //          root_name = static_name;
+  //        }
+  //      }
+  //      update the Map entry for root_name;
+  //      for (subfs: FS) {
+  //        recursive(subfs, root_name);
+  //      }
+  //   }
+  //
+  // Here is an example.
+  //
+  // SampleProfile:
+  // foo:12345:1000
+  // 1: 1000
+  // 2.1: 1000
+  // 15: 5000
+  // 4: bar:1000
+  //  1: 1000
+  //  2: goo:3000
+  //   1: 3000
+  // 8: bar:40000
+  //  1: 10000
+  //  2: goo:30000
+  //   1: 30000
+  //
+  // InstrProfile has two entries:
+  //  foo
+  //  bar.cc:bar
+  //
+  // After BuildMaxSampleMap, we should have the following in FlattenSampleMap:
+  // {"foo", {1000, 5000}}
+  // {"bar.cc:bar", {11000, 30000}}
+  //
+  // foo's has an entry count of 1000, and max body count of 5000.
+  // bar.cc:bar has an entry count of 11000 (sum two callsites of 1000 and
+  // 10000), and max count of 30000 (from the callsite in line 8).
+  //
+  // Note that goo's count will remain in bar.cc:bar() as it does not have an
+  // entry in InstrProfile.
+  DenseMap<StringRef, std::pair<uint64_t, uint64_t>> FlattenSampleMap;
+  auto BuildMaxSampleMap = [&FlattenSampleMap, &StaticFuncMap,
+                            &InstrProfileMap](const FunctionSamples &FS,
+                                              const StringRef &RootName) {
+    auto BuildMaxSampleMapImpl = [&](const FunctionSamples &FS,
+                                     const StringRef &RootName,
+                                     auto &BuildImpl) -> void {
+      const StringRef &Name = FS.getName();
+      const StringRef *NewRootName = &RootName;
+      uint64_t EntrySample = FS.getHeadSamplesEstimate();
+      uint64_t MaxBodySample = FS.getMaxCountInside(/* SkipCallSite*/ true);
+
+      auto It = InstrProfileMap.find(Name);
+      if (It != InstrProfileMap.end()) {
+        NewRootName = &Name;
+      } else {
+        auto NewName = StaticFuncMap.find(Name);
+        if (NewName != StaticFuncMap.end()) {
+          It = InstrProfileMap.find(NewName->second.str());
+          if (NewName->second != DuplicateNameStr) {
+            NewRootName = &NewName->second;
+          }
+        } else {
+          // Here the EntrySample is of an inlined function, so we should not
+          // update the EntrySample in the map.
+          EntrySample = 0;
+        }
+      }
+      EntrySample += FlattenSampleMap[*NewRootName].first;
+      MaxBodySample =
+          std::max(FlattenSampleMap[*NewRootName].second, MaxBodySample);
+      FlattenSampleMap[*NewRootName] =
+          std::make_pair(EntrySample, MaxBodySample);
+
+      for (const auto &C : FS.getCallsiteSamples())
+        for (const auto &F : C.second)
+          BuildImpl(F.second, *NewRootName, BuildImpl);
+    };
+    BuildMaxSampleMapImpl(FS, RootName, BuildMaxSampleMapImpl);
+  };
+
   for (auto &PD : WC->Writer.getProfileData()) {
     // Populate IPBuilder.
     for (const auto &PDV : PD.getValue()) {
       InstrProfRecord Record = PDV.second;
       IPBuilder.addRecord(Record);
     }
 
     // If a function has multiple entries in instr profile, skip it.
     if (PD.getValue().size() != 1)
       continue;
 
     // Initialize InstrProfileMap.
     InstrProfRecord *R = &PD.getValue().begin()->second;
     StringRef FullName = PD.getKey();
     InstrProfileMap[FullName] = InstrProfileEntry(R);
     buildStaticFuncMap(FullName);
   }
 
+  for (auto &PD : Reader->getProfiles()) {
+    sampleprof::FunctionSamples &FS = PD.second;
+    BuildMaxSampleMap(FS, FS.getName());
+  }
+
   ProfileSummary InstrPS = *IPBuilder.getSummary();
   ProfileSummary SamplePS = Reader->getSummary();
 
   // Compute cold thresholds for instr profile and sample profile.
   uint64_t HotSampleThreshold =
       ProfileSummaryBuilder::getEntryForPercentile(
           SamplePS.getDetailedSummary(),
           ProfileSummaryBuilder::DefaultCutoffs[HotPercentileIdx])
uint64_t ColdInstrThreshold =
           ? InstrProfColdThreshold
           : ProfileSummaryBuilder::getEntryForPercentile(
                 InstrPS.getDetailedSummary(),
                 ProfileSummaryBuilder::DefaultCutoffs[ColdPercentileIdx])
                 .MinCount;
 
   // Find hot/warm functions in sample profile which is cold in instr profile
   // and adjust the profiles of those functions in the instr profile.
-  for (const auto &PD : Reader->getProfiles()) {
-    const sampleprof::FunctionSamples &FS = PD.second;
-    uint64_t SampleMaxCount = FS.getMaxCountInside();
+  for (const auto &E : FlattenSampleMap) {
+    uint64_t SampleMaxCount = std::max(E.second.first, E.second.second);
     if (SampleMaxCount < ColdSampleThreshold)
       continue;
-    auto &FContext = PD.first;
-    auto It = InstrProfileMap.find(FContext.toString());
+    const StringRef &Name = E.first;
+    auto It = InstrProfileMap.find(Name);
     if (It == InstrProfileMap.end()) {
-      auto NewName = StaticFuncMap.find(FContext.toString());
+      auto NewName = StaticFuncMap.find(Name);
       if (NewName != StaticFuncMap.end()) {
         It = InstrProfileMap.find(NewName->second.str());
         if (NewName->second == DuplicateNameStr) {
           WithColor::warning()
-              << "Static function " << FContext.toString()
+              << "Static function " << Name
               << " has multiple promoted names, cannot adjust profile.\n";
         }
       }
     }
     if (It == InstrProfileMap.end() ||
         It->second.MaxCount > ColdInstrThreshold ||
         It->second.NumEdgeCounters < SupplMinSizeThreshold)
       continue;