1 files changed, 53 insertions, 0 deletions

diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
index 7baeb3f..fbcd614 100644
--- a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -76,6 +76,21 @@ BasicBlockSectionsProfileReader::getClonePathsForFunction(
   return ProgramPathAndClusterInfo.lookup(getAliasName(FuncName)).ClonePaths;
 }
 
+uint64_t BasicBlockSectionsProfileReader::getEdgeCount(
+    StringRef FuncName, const UniqueBBID &SrcBBID,
+    const UniqueBBID &SinkBBID) const {
+  auto It = ProgramPathAndClusterInfo.find(getAliasName(FuncName));
+  if (It == ProgramPathAndClusterInfo.end())
+    return 0;
+  auto NodeIt = It->second.EdgeCounts.find(SrcBBID);
+  if (NodeIt == It->second.EdgeCounts.end())
+    return 0;
+  auto EdgeIt = NodeIt->second.find(SinkBBID);
+  if (EdgeIt == NodeIt->second.end())
+    return 0;
+  return EdgeIt->second;
+}
+
 // Reads the version 1 basic block sections profile. Profile for each function
 // is encoded as follows:
 //   m <module_name>
@@ -240,6 +255,38 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       }
       continue;
     }
+    case 'g': { // CFG profile specifier.
+      // Skip the profile when we the profile iterator (FI) refers to the
+      // past-the-end element.
+      if (FI == ProgramPathAndClusterInfo.end())
+        continue;
+      // For each node, its CFG profile is encoded as
+      // <src>:<count>,<sink_1>:<count_1>,<sink_2>:<count_2>,...
+      for (auto BasicBlockEdgeProfile : Values) {
+        if (BasicBlockEdgeProfile.empty())
+          continue;
+        SmallVector<StringRef, 4> NodeEdgeCounts;
+        BasicBlockEdgeProfile.split(NodeEdgeCounts, ',');
+        UniqueBBID SrcBBID;
+        for (size_t i = 0; i < NodeEdgeCounts.size(); ++i) {
+          auto [BBIDStr, CountStr] = NodeEdgeCounts[i].split(':');
+          auto BBID = parseUniqueBBID(BBIDStr);
+          if (!BBID)
+            return BBID.takeError();
+          unsigned long long Count = 0;
+          if (getAsUnsignedInteger(CountStr, 10, Count))
+            return createProfileParseError(
+                Twine("unsigned integer expected: '") + CountStr + "'");
+          if (i == 0) {
+            // The first element represents the source and its total count.
+            FI->second.NodeCounts[SrcBBID = *BBID] = Count;
+            continue;
+          }
+          FI->second.EdgeCounts[SrcBBID][*BBID] = Count;
+        }
+      }
+      continue;
+    }
     default:
       return createProfileParseError(Twine("invalid specifier: '") +
                                      Twine(Specifier) + "'");
@@ -440,6 +487,12 @@ BasicBlockSectionsProfileReaderWrapperPass::getClonePathsForFunction(
   return BBSPR.getClonePathsForFunction(FuncName);
 }
 
+uint64_t BasicBlockSectionsProfileReaderWrapperPass::getEdgeCount(
+    StringRef FuncName, const UniqueBBID &SrcBBID,
+    const UniqueBBID &SinkBBID) const {
+  return BBSPR.getEdgeCount(FuncName, SrcBBID, SinkBBID);
+}
+
 BasicBlockSectionsProfileReader &
 BasicBlockSectionsProfileReaderWrapperPass::getBBSPR() {
   return BBSPR;