Enforce skipflatprof for text format sample profile

Author: huangjd

llvm/include/llvm/ProfileData/SampleProfReader.h

@@ -495,9 +495,9 @@ class SampleProfileReader {
   /// are present.
   virtual void setProfileUseMD5() { ProfileIsMD5 = true; }
 
-  /// Don't read profile without context if the flag is set. This is only meaningful
-  /// for ExtBinary format.
-  virtual void setSkipFlatProf(bool Skip) {}
+  /// Don't read profile without context if the flag is set.
+  void setSkipFlatProf(bool Skip) { SkipFlatProf = Skip; }
+
   /// Return whether any name in the profile contains ".__uniq." suffix.
   virtual bool hasUniqSuffix() { return false; }
 
@@ -581,6 +581,10 @@ class SampleProfileReader {
   /// Whether the profile uses MD5 for Sample Contexts and function names. This
   /// can be one-way overriden by the user to force use MD5.
   bool ProfileIsMD5 = false;
+
+  /// If SkipFlatProf is true, skip functions marked with !Flat in text mode or
+  /// sections with SecFlagFlat flag in ExtBinary mode.
+  bool SkipFlatProf = false;
 };
 
 class SampleProfileReaderText : public SampleProfileReader {
@@ -789,10 +793,6 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
 
-  /// If SkipFlatProf is true, skip the sections with
-  /// SecFlagFlat flag.
-  bool SkipFlatProf = false;
-
 public:
   SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C, SampleProfileFormat Format)
@@ -815,8 +815,6 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
     return std::move(ProfSymList);
   };
 
-  void setSkipFlatProf(bool Skip) override { SkipFlatProf = Skip; }
-
 private:
   /// Read the profiles on-demand for the given functions. This is used after
   /// stale call graph matching finds new functions whose profiles aren't loaded

llvm/lib/ProfileData/SampleProfReader.cpp

@@ -215,7 +215,8 @@ static bool ParseLine(const StringRef &Input, LineType &LineTy, uint32_t &Depth,
                       uint64_t &NumSamples, uint32_t &LineOffset,
                       uint32_t &Discriminator, StringRef &CalleeName,
                       DenseMap<StringRef, uint64_t> &TargetCountMap,
-                      uint64_t &FunctionHash, uint32_t &Attributes) {
+                      uint64_t &FunctionHash, uint32_t &Attributes,
+                      bool &IsFlat) {
   for (Depth = 0; Input[Depth] == ' '; Depth++)
     ;
   if (Depth == 0)
@@ -226,8 +227,10 @@ static bool ParseLine(const StringRef &Input, LineType &LineTy, uint32_t &Depth,
     // This metadata is only for manual inspection only. We already created a
     // FunctionSamples and put it in the profile map, so there is no point
     // to skip profiles even they have no use for ThinLTO.
-    if (Input == StringRef(" !Flat"))
+    if (Input == StringRef(" !Flat")) {
+      IsFlat = true;
       return true;
+    }
     return parseMetadata(Input.substr(Depth), FunctionHash, Attributes);
   }
 
@@ -330,6 +333,8 @@ std::error_code SampleProfileReaderText::readImpl() {
   // top-level or nested function profile.
   uint32_t DepthMetadata = 0;
 
+  std::vector<SampleContext *> FlatSamples;
+
   ProfileIsFS = ProfileIsFSDisciminator;
   FunctionSamples::ProfileIsFS = ProfileIsFS;
   for (; !LineIt.is_at_eof(); ++LineIt) {
@@ -373,9 +378,10 @@ std::error_code SampleProfileReaderText::readImpl() {
       LineType LineTy;
       uint64_t FunctionHash = 0;
       uint32_t Attributes = 0;
+      bool IsFlat = false;
       if (!ParseLine(*LineIt, LineTy, Depth, NumSamples, LineOffset,
                      Discriminator, FName, TargetCountMap, FunctionHash,
-                     Attributes)) {
+                     Attributes, IsFlat)) {
         reportError(LineIt.line_number(),
                     "Expected 'NUM[.NUM]: NUM[ mangled_name:NUM]*', found " +
                         *LineIt);
@@ -431,12 +437,26 @@ std::error_code SampleProfileReaderText::readImpl() {
         if (Attributes & (uint32_t)ContextShouldBeInlined)
           ProfileIsPreInlined = true;
         DepthMetadata = Depth;
+        if (IsFlat) {
+          if (Depth == 1)
+            FlatSamples.push_back(&FProfile.getContext());
+          else
+            Ctx.diagnose(DiagnosticInfoSampleProfile(
+                Buffer->getBufferIdentifier(), LineIt.line_number(),
+                "!Flat may only be used at top level function.", DS_Warning));
+        }
         break;
       }
       }
     }
   }
 
+  // Honor the option to skip flat functions. Since they are already added to
+  // the profile map, remove them all here.
+  if (SkipFlatProf)
+    for (SampleContext *FlatSample : FlatSamples)
+      Profiles.erase(*FlatSample);
+
   assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
          "Cannot have both context-sensitive and regular profile");
   ProfileIsCS = (CSProfileCount > 0);