[CGData][lld-macho] Merge CG Data by LLD (#112674)

LLD now processes raw CG data for stable functions, similar to how it
handles raw CG data for the outliner's hash tree. This data is encoded
in the custom section (`__llvm_merge`) within object files. LLD merges
this information into the indexed CG data file specified by the
`-codegen-data-generate-path={path}` option. For the linker that does
not support this feature, we could use `llvm-cgdata` tool --
https://github.com/llvm/llvm-project/blob/main/llvm/docs/CommandGuide/llvm-cgdata.rst.

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

lld/MachO/Driver.cpp

@@ -1326,21 +1326,39 @@ static void codegenDataGenerate() {
   TimeTraceScope timeScope("Generating codegen data");
 
   OutlinedHashTreeRecord globalOutlineRecord;
-  for (ConcatInputSection *isec : inputSections)
-    if (isec->getSegName() == segment_names::data &&
-        isec->getName() == section_names::outlinedHashTree) {
+  StableFunctionMapRecord globalMergeRecord;
+  for (ConcatInputSection *isec : inputSections) {
+    if (isec->getSegName() != segment_names::data)
+      continue;
+    if (isec->getName() == section_names::outlinedHashTree) {
       // Read outlined hash tree from each section.
       OutlinedHashTreeRecord localOutlineRecord;
+      // Use a pointer to allow modification by the function.
       auto *data = isec->data.data();
       localOutlineRecord.deserialize(data);
 
       // Merge it to the global hash tree.
       globalOutlineRecord.merge(localOutlineRecord);
     }
+    if (isec->getName() == section_names::functionMap) {
+      // Read stable functions from each section.
+      StableFunctionMapRecord localMergeRecord;
+      // Use a pointer to allow modification by the function.
+      auto *data = isec->data.data();
+      localMergeRecord.deserialize(data);
+
+      // Merge it to the global function map.
+      globalMergeRecord.merge(localMergeRecord);
+    }
+  }
+
+  globalMergeRecord.finalize();
 
   CodeGenDataWriter Writer;
   if (!globalOutlineRecord.empty())
     Writer.addRecord(globalOutlineRecord);
+  if (!globalMergeRecord.empty())
+    Writer.addRecord(globalMergeRecord);
 
   std::error_code EC;
   auto fileName = config->codegenDataGeneratePath;

lld/MachO/InputSection.h

@@ -339,6 +339,7 @@ constexpr const char const_[] = "__const";
 constexpr const char lazySymbolPtr[] = "__la_symbol_ptr";
 constexpr const char lazyBinding[] = "__lazy_binding";
 constexpr const char literals[] = "__literals";
+constexpr const char functionMap[] = "__llvm_merge";
 constexpr const char moduleInitFunc[] = "__mod_init_func";
 constexpr const char moduleTermFunc[] = "__mod_term_func";
 constexpr const char nonLazySymbolPtr[] = "__nl_symbol_ptr";

lld/test/MachO/cgdata-generate-merge.s

@@ -0,0 +1,85 @@
+# UNSUPPORTED: system-windows
+# REQUIRES: aarch64
+
+# RUN: rm -rf %t; split-file %s %t
+
+# 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 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 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
+# RUN: llvm-mc -filetype obj -triple arm64-apple-darwin %t/merge-2.s -o %t/merge-2.o
+# RUN: llvm-mc -filetype obj -triple arm64-apple-darwin %t/main.s -o %t/main.o
+
+# This checks if the codegen data from the linker is identical to the merged codegen data
+# from each object file, which is obtained using the llvm-cgdata tool.
+# RUN: %no-arg-lld -dylib -arch arm64 -platform_version ios 14.0 15.0 -o %t/out \
+# RUN: %t/merge-1.o %t/merge-2.o %t/main.o --codegen-data-generate-path=%t/out-cgdata
+# RUN: llvm-cgdata --merge %t/merge-1.o %t/merge-2.o %t/main.o -o %t/merge-cgdata
+# RUN: diff %t/out-cgdata %t/merge-cgdata
+
+# Merge order doesn't matter in the yaml format. `main.o` is dropped due to missing __llvm_merge.
+# RUN: llvm-cgdata --merge %t/merge-2.o %t/merge-1.o -o %t/merge-cgdata-shuffle
+# RUN: llvm-cgdata --convert %t/out-cgdata -o %t/out-cgdata.yaml
+# RUN: llvm-cgdata --convert %t/merge-cgdata-shuffle -o %t/merge-cgdata-shuffle.yaml
+# RUN: diff %t/out-cgdata.yaml %t/merge-cgdata-shuffle.yaml
+
+# We can also generate the merged codegen data from the executable that is not dead-stripped.
+# RUN: llvm-objdump -h %t/out| FileCheck %s
+# CHECK: __llvm_merge
+# RUN: llvm-cgdata --merge %t/out -o %t/merge-cgdata-exe
+# RUN: diff %t/merge-cgdata-exe %t/merge-cgdata
+
+# Dead-strip will remove __llvm_merge sections from the final executable.
+# But the codeden data is still correctly produced from the linker.
+# RUN: %no-arg-lld -dylib -arch arm64 -platform_version ios 14.0 15.0 -o %t/out-strip \
+# RUN: %t/merge-1.o %t/merge-2.o %t/main.o -dead_strip --codegen-data-generate-path=%t/out-cgdata-strip
+# RUN: llvm-cgdata --merge %t/merge-1.o %t/merge-2.o %t/main.o -o %t/merge-cgdata-strip
+# RUN: diff %t/out-cgdata-strip %t/merge-cgdata-strip
+# RUN: diff %t/out-cgdata-strip %t/merge-cgdata
+
+# Ensure no __llvm_merge section remains in the executable.
+# RUN: llvm-objdump -h %t/out-strip | FileCheck %s --check-prefix=STRIP
+# STRIP-NOT: __llvm_merge
+
+#--- raw-1.cgtext
+:stable_function_map
+---
+- Hash:            123
+  FunctionName:    f1
+  ModuleName:      'foo.bc'
+  InstCount:       7
+  IndexOperandHashes:
+    - InstIndex:       3
+      OpndIndex:       0
+      OpndHash:        456
+...
+
+#--- raw-2.cgtext
+:stable_function_map
+---
+- Hash:            123
+  FunctionName:    f2
+  ModuleName:      'goo.bc'
+  InstCount:       7
+  IndexOperandHashes:
+    - InstIndex:       3
+      OpndIndex:       0
+      OpndHash:        789
+...
+
+#--- merge-template.s
+.section __DATA,__llvm_merge
+_data:
+.byte <RAW_BYTES>
+
+#--- main.s
+.globl _main
+
+.text
+_main:
+  ret