[CGData][llvm-cgdata] Support for stable function map (#112664)

This introduces a new cgdata format for stable function maps. The raw
data is embedded in the __llvm_merge section during compile time. This
data can be read and merged using the llvm-cgdata tool, into an indexed
cgdata file. Consequently, the tool is now capable of handling either
outlined hash trees, stable function maps, or both, as they are
orthogonal.

Depends on #112662.
This is a patch for
https://discourse.llvm.org/t/rfc-global-function-merging/82608.

lld/test/MachO/cgdata-generate.s

@@ -3,12 +3,12 @@
 
 # RUN: rm -rf %t; split-file %s %t
 
-# Synthesize raw cgdata without the header (24 byte) from the indexed cgdata.
+# Synthesize raw cgdata without the header (32 byte) from the indexed cgdata.
 # RUN: llvm-cgdata --convert --format binary %t/raw-1.cgtext -o %t/raw-1.cgdata
-# RUN: od -t x1 -j 24 -An %t/raw-1.cgdata | tr -d '\n\r\t' | sed 's/[ ][ ]*/ /g; s/^[ ]*//; s/[ ]*$//; s/[ ]/,0x/g; s/^/0x/' > %t/raw-1-bytes.txt
+# RUN: od -t x1 -j 32 -An %t/raw-1.cgdata | tr -d '\n\r\t' | sed 's/[ ][ ]*/ /g; s/^[ ]*//; s/[ ]*$//; s/[ ]/,0x/g; s/^/0x/' > %t/raw-1-bytes.txt
 # RUN: sed "s/<RAW_BYTES>/$(cat %t/raw-1-bytes.txt)/g" %t/merge-template.s > %t/merge-1.s
 # RUN: llvm-cgdata --convert --format binary %t/raw-2.cgtext -o %t/raw-2.cgdata
-# RUN: od -t x1 -j 24 -An %t/raw-2.cgdata | tr -d '\n\r\t' | sed 's/[ ][ ]*/ /g; s/^[ ]*//; s/[ ]*$//; s/[ ]/,0x/g; s/^/0x/' > %t/raw-2-bytes.txt
+# RUN: od -t x1 -j 32 -An %t/raw-2.cgdata | tr -d '\n\r\t' | sed 's/[ ][ ]*/ /g; s/^[ ]*//; s/[ ]*$//; s/[ ]/,0x/g; s/^/0x/' > %t/raw-2-bytes.txt
 # RUN: sed "s/<RAW_BYTES>/$(cat %t/raw-2-bytes.txt)/g" %t/merge-template.s > %t/merge-2.s
 
 # RUN: llvm-mc -filetype obj -triple arm64-apple-darwin %t/merge-1.s -o %t/merge-1.o

llvm/docs/CommandGuide/llvm-cgdata.rst

@@ -11,15 +11,13 @@ SYNOPSIS
 DESCRIPTION
 -----------
 
-The :program:llvm-cgdata utility parses raw codegen data embedded
-in compiled binary files and merges them into a single .cgdata file.
-It can also inspect and manipulate .cgdata files.
-Currently, the tool supports saving and restoring outlined hash trees,
-enabling global function outlining across modules, allowing for more
-efficient function outlining in subsequent compilations.
-The design is extensible, allowing for the incorporation of additional
-codegen summaries and optimization techniques, such as global function
-merging, in the future.
+The :program:llvm-cgdata utility parses raw codegen data embedded in compiled
+binary files and merges them into a single .cgdata file. It can also inspect
+and manipulate .cgdata files. Currently, the tool supports saving and restoring
+outlined hash trees and stable function maps, allowing for more efficient
+function outlining and function merging across modules in subsequent
+compilations. The design is extensible, allowing for the incorporation of
+additional codegen summaries and optimization techniques.
 
 COMMANDS
 --------

llvm/include/llvm/CGData/CodeGenData.h

@@ -19,6 +19,7 @@
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/CGData/OutlinedHashTree.h"
 #include "llvm/CGData/OutlinedHashTreeRecord.h"
+#include "llvm/CGData/StableFunctionMapRecord.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Caching.h"
@@ -41,7 +42,9 @@ enum class CGDataKind {
   Unknown = 0x0,
   // A function outlining info.
   FunctionOutlinedHashTree = 0x1,
-  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/FunctionOutlinedHashTree)
+  // A function merging info.
+  StableFunctionMergingMap = 0x2,
+  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/StableFunctionMergingMap)
 };
 
 const std::error_category &cgdata_category();
@@ -108,6 +111,8 @@ enum CGDataMode {
 class CodeGenData {
   /// Global outlined hash tree that has oulined hash sequences across modules.
   std::unique_ptr<OutlinedHashTree> PublishedHashTree;
+  /// Global stable function map that has stable function info across modules.
+  std::unique_ptr<StableFunctionMap> PublishedStableFunctionMap;
 
   /// This flag is set when -fcodegen-data-generate is passed.
   /// Or, it can be mutated with -fcodegen-data-thinlto-two-rounds.
@@ -131,6 +136,9 @@ class CodeGenData {
   bool hasOutlinedHashTree() {
     return PublishedHashTree && !PublishedHashTree->empty();
   }
+  bool hasStableFunctionMap() {
+    return PublishedStableFunctionMap && !PublishedStableFunctionMap->empty();
+  }
 
   /// Returns the outlined hash tree. This can be globally used in a read-only
   /// manner.
@@ -147,6 +155,12 @@ class CodeGenData {
     // Ensure we disable emitCGData as we do not want to read and write both.
     EmitCGData = false;
   }
+  void
+  publishStableFunctionMap(std::unique_ptr<StableFunctionMap> FunctionMap) {
+    PublishedStableFunctionMap = std::move(FunctionMap);
+    // Ensure we disable emitCGData as we do not want to read and write both.
+    EmitCGData = false;
+  }
 };
 
 namespace cgdata {
@@ -166,6 +180,11 @@ publishOutlinedHashTree(std::unique_ptr<OutlinedHashTree> HashTree) {
   CodeGenData::getInstance().publishOutlinedHashTree(std::move(HashTree));
 }
 
+inline void
+publishStableFunctionMap(std::unique_ptr<StableFunctionMap> FunctionMap) {
+  CodeGenData::getInstance().publishStableFunctionMap(std::move(FunctionMap));
+}
+
 struct StreamCacheData {
   /// Backing buffer for serialized data stream.
   SmallVector<SmallString<0>> Outputs;
@@ -249,6 +268,8 @@ enum CGDataVersion {
   // Version 1 is the first version. This version supports the outlined
   // hash tree.
   Version1 = 1,
+  // Version 2 supports the stable function merging map.
+  Version2 = 2,
   CurrentVersion = CG_DATA_INDEX_VERSION
 };
 const uint64_t Version = CGDataVersion::CurrentVersion;
@@ -258,6 +279,7 @@ struct Header {
   uint32_t Version;
   uint32_t DataKind;
   uint64_t OutlinedHashTreeOffset;
+  uint64_t StableFunctionMapOffset;
 
   // New fields should only be added at the end to ensure that the size
   // computation is correct. The methods below need to be updated to ensure that

llvm/include/llvm/CGData/CodeGenData.inc

@@ -20,27 +20,33 @@
 #define CG_DATA_DEFINED
 CG_DATA_SECT_ENTRY(CG_outline, CG_DATA_QUOTE(CG_DATA_OUTLINE_COMMON),
                    CG_DATA_OUTLINE_COFF, "__DATA,")
+CG_DATA_SECT_ENTRY(CG_merge, CG_DATA_QUOTE(CG_DATA_MERGE_COMMON),
+                   CG_DATA_MERGE_COFF, "__DATA,")
 
 #undef CG_DATA_SECT_ENTRY
 #endif
 
 /* section name strings common to all targets other
    than WIN32 */
 #define CG_DATA_OUTLINE_COMMON __llvm_outline
+#define CG_DATA_MERGE_COMMON __llvm_merge
 /* Since cg data sections are not allocated, we don't need to
  * access them at runtime.
  */
 #define CG_DATA_OUTLINE_COFF ".loutline"
+#define CG_DATA_MERGE_COFF ".lmerge"
 
 #ifdef _WIN32
 /* Runtime section names and name strings.  */
-#define CG_DATA_SECT_NAME CG_DATA_OUTLINE_COFF
+#define CG_DATA_OUTLINE_SECT_NAME CG_DATA_OUTLINE_COFF
+#define CG_DATA_MERGE_SECT_NAME CG_DATA_MERGE_COFF
 
 #else
 /* Runtime section names and name strings.  */
-#define CG_DATA_SECT_NAME CG_DATA_QUOTE(CG_DATA_OUTLINE_COMMON)
+#define CG_DATA_OUTLINE_SECT_NAME CG_DATA_QUOTE(CG_DATA_OUTLINE_COMMON)
+#define CG_DATA_MERGE_SECT_NAME CG_DATA_QUOTE(CG_DATA_MERGE_COMMON)
 
 #endif
 
 /* Indexed codegen data format version (start from 1). */
-#define CG_DATA_INDEX_VERSION 1
+#define CG_DATA_INDEX_VERSION 2

llvm/include/llvm/CGData/CodeGenDataReader.h

@@ -15,6 +15,7 @@
 
 #include "llvm/CGData/CodeGenData.h"
 #include "llvm/CGData/OutlinedHashTreeRecord.h"
+#include "llvm/CGData/StableFunctionMapRecord.h"
 #include "llvm/Support/LineIterator.h"
 #include "llvm/Support/VirtualFileSystem.h"
 
@@ -36,10 +37,15 @@ class CodeGenDataReader {
   virtual CGDataKind getDataKind() const = 0;
   /// Return true if the data has an outlined hash tree.
   virtual bool hasOutlinedHashTree() const = 0;
+  /// Return true if the data has a stable function map.
+  virtual bool hasStableFunctionMap() const = 0;
   /// Return the outlined hash tree that is released from the reader.
   std::unique_ptr<OutlinedHashTree> releaseOutlinedHashTree() {
     return std::move(HashTreeRecord.HashTree);
   }
+  std::unique_ptr<StableFunctionMap> releaseStableFunctionMap() {
+    return std::move(FunctionMapRecord.FunctionMap);
+  }
 
   /// Factory method to create an appropriately typed reader for the given
   /// codegen data file path and file system.
@@ -56,15 +62,21 @@ class CodeGenDataReader {
   /// is used by `llvm-cgdata --merge` or ThinLTO's two-codegen rounds.
   /// Optionally, \p CombinedHash can be used to compuate the combined hash of
   /// the merged data.
-  static Error mergeFromObjectFile(const object::ObjectFile *Obj,
-                                   OutlinedHashTreeRecord &GlobalOutlineRecord,
-                                   stable_hash *CombinedHash = nullptr);
+  static Error
+  mergeFromObjectFile(const object::ObjectFile *Obj,
+                      OutlinedHashTreeRecord &GlobalOutlineRecord,
+                      StableFunctionMapRecord &GlobalFunctionMapRecord,
+                      stable_hash *CombinedHash = nullptr);
 
 protected:
   /// The outlined hash tree that has been read. When it's released by
   /// releaseOutlinedHashTree(), it's no longer valid.
   OutlinedHashTreeRecord HashTreeRecord;
 
+  /// The stable function map that has been read. When it's released by
+  // releaseStableFunctionMap(), it's no longer valid.
+  StableFunctionMapRecord FunctionMapRecord;
+
   /// Set the current error and return same.
   Error error(cgdata_error Err, const std::string &ErrMsg = "") {
     LastError = Err;
@@ -115,6 +127,11 @@ class IndexedCodeGenDataReader : public CodeGenDataReader {
     return Header.DataKind &
            static_cast<uint32_t>(CGDataKind::FunctionOutlinedHashTree);
   }
+  /// Return true if the header indicates the data has a stable function map.
+  bool hasStableFunctionMap() const override {
+    return Header.DataKind &
+           static_cast<uint32_t>(CGDataKind::StableFunctionMergingMap);
+  }
 };
 
 /// This format is a simple text format that's suitable for test data.
@@ -150,6 +167,12 @@ class TextCodeGenDataReader : public CodeGenDataReader {
     return static_cast<uint32_t>(DataKind) &
            static_cast<uint32_t>(CGDataKind::FunctionOutlinedHashTree);
   }
+  /// Return true if the header indicates the data has a stable function map.
+  /// This does not mean that the data is still available.
+  bool hasStableFunctionMap() const override {
+    return static_cast<uint32_t>(DataKind) &
+           static_cast<uint32_t>(CGDataKind::StableFunctionMergingMap);
+  }
 };
 
 } // end namespace llvm

llvm/include/llvm/CGData/CodeGenDataWriter.h

@@ -15,6 +15,7 @@
 
 #include "llvm/CGData/CodeGenData.h"
 #include "llvm/CGData/OutlinedHashTreeRecord.h"
+#include "llvm/CGData/StableFunctionMapRecord.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/Error.h"
 
@@ -57,16 +58,22 @@ class CodeGenDataWriter {
   /// The outlined hash tree to be written.
   OutlinedHashTreeRecord HashTreeRecord;
 
+  /// The stable function map to be written.
+  StableFunctionMapRecord FunctionMapRecord;
+
   /// A bit mask describing the kind of the codegen data.
   CGDataKind DataKind = CGDataKind::Unknown;
 
 public:
   CodeGenDataWriter() = default;
   ~CodeGenDataWriter() = default;
 
-  /// Add the outlined hash tree record. The input Record is released.
+  /// Add the outlined hash tree record. The input hash tree is released.
   void addRecord(OutlinedHashTreeRecord &Record);
 
+  /// Add the stable function map record. The input function map is released.
+  void addRecord(StableFunctionMapRecord &Record);
+
   /// Write the codegen data to \c OS
   Error write(raw_fd_ostream &OS);
 
@@ -81,11 +88,19 @@ class CodeGenDataWriter {
     return static_cast<uint32_t>(DataKind) &
            static_cast<uint32_t>(CGDataKind::FunctionOutlinedHashTree);
   }
+  /// Return true if the header indicates the data has a stable function map.
+  bool hasStableFunctionMap() const {
+    return static_cast<uint32_t>(DataKind) &
+           static_cast<uint32_t>(CGDataKind::StableFunctionMergingMap);
+  }
 
 private:
   /// The offset of the outlined hash tree in the file.
   uint64_t OutlinedHashTreeOffset;
 
+  /// The offset of the stable function map in the file.
+  uint64_t StableFunctionMapOffset;
+
   /// Write the codegen data header to \c COS
   Error writeHeader(CGDataOStream &COS);
 

llvm/lib/CGData/CodeGenData.cpp

@@ -14,6 +14,7 @@
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/CGData/CodeGenDataReader.h"
 #include "llvm/CGData/OutlinedHashTreeRecord.h"
+#include "llvm/CGData/StableFunctionMapRecord.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Caching.h"
 #include "llvm/Support/CommandLine.h"
@@ -163,6 +164,8 @@ CodeGenData &CodeGenData::getInstance() {
       auto Reader = ReaderOrErr->get();
       if (Reader->hasOutlinedHashTree())
         Instance->publishOutlinedHashTree(Reader->releaseOutlinedHashTree());
+      if (Reader->hasStableFunctionMap())
+        Instance->publishStableFunctionMap(Reader->releaseStableFunctionMap());
     }
   });
   return *Instance;
@@ -185,18 +188,14 @@ Expected<Header> Header::readFromBuffer(const unsigned char *Curr) {
     return make_error<CGDataError>(cgdata_error::unsupported_version);
   H.DataKind = endian::readNext<uint32_t, endianness::little, unaligned>(Curr);
 
-  switch (H.Version) {
-    // When a new field is added to the header add a case statement here to
-    // compute the size as offset of the new field + size of the new field. This
-    // relies on the field being added to the end of the list.
-    static_assert(IndexedCGData::CGDataVersion::CurrentVersion == Version1,
-                  "Please update the size computation below if a new field has "
-                  "been added to the header, if not add a case statement to "
-                  "fall through to the latest version.");
-  case 1ull:
-    H.OutlinedHashTreeOffset =
+  static_assert(IndexedCGData::CGDataVersion::CurrentVersion == Version2,
+                "Please update the offset computation below if a new field has "
+                "been added to the header.");
+  H.OutlinedHashTreeOffset =
+      endian::readNext<uint64_t, endianness::little, unaligned>(Curr);
+  if (H.Version >= 2)
+    H.StableFunctionMapOffset =
         endian::readNext<uint64_t, endianness::little, unaligned>(Curr);
-  }
 
   return H;
 }
@@ -257,6 +256,7 @@ std::unique_ptr<Module> loadModuleForTwoRounds(BitcodeModule &OrigModule,
 
 Expected<stable_hash> mergeCodeGenData(ArrayRef<StringRef> ObjFiles) {
   OutlinedHashTreeRecord GlobalOutlineRecord;
+  StableFunctionMapRecord GlobalStableFunctionMapRecord;
   stable_hash CombinedHash = 0;
   for (auto File : ObjFiles) {
     if (File.empty())
@@ -270,12 +270,18 @@ Expected<stable_hash> mergeCodeGenData(ArrayRef<StringRef> ObjFiles) {
 
     std::unique_ptr<object::ObjectFile> &Obj = BinOrErr.get();
     if (auto E = CodeGenDataReader::mergeFromObjectFile(
-            Obj.get(), GlobalOutlineRecord, &CombinedHash))
+            Obj.get(), GlobalOutlineRecord, GlobalStableFunctionMapRecord,
+            &CombinedHash))
       return E;
   }
 
+  GlobalStableFunctionMapRecord.finalize();
+
   if (!GlobalOutlineRecord.empty())
     cgdata::publishOutlinedHashTree(std::move(GlobalOutlineRecord.HashTree));
+  if (!GlobalStableFunctionMapRecord.empty())
+    cgdata::publishStableFunctionMap(
+        std::move(GlobalStableFunctionMapRecord.FunctionMap));
 
   return CombinedHash;
 }