[CGData][lld-macho] Merge CG Data by LLD

[kyulee-com]

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.

[kyulee-com]

cc. @nocchijiang

[alx32]

Can we merge the if(isec->getSegName() == segment_names::data) check here and below on 1341 ?
i.e. with else if for isec->getName() == section_names::functionmap

[alx32]

I think functionMap would be the name here if we were to match existing pattern.

[MaskRay]

drive-by: you can add cd %t to remove %t/ below

[ellishg]

I actually like keeping the %t/ in the commands because it allows me to copy and run the commands directly for debugging without needing to first cd %t.

[kyulee-com]

Yes, I'll leave it as is, since each command is self-contained and complete, making it convenient to reproduce.

[alx32]

LGTM !

[llvmbot]

@llvm/pr-subscribers-lld

@llvm/pr-subscribers-lld-macho

Author: Kyungwoo Lee (kyulee-com)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/112674.diff

3 Files Affected:

• (modified) lld/MachO/Driver.cpp (+21-3)
• (modified) lld/MachO/InputSection.h (+1)
• (added) lld/test/MachO/cgdata-generate-merge.s (+85)

diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
index ab4abb1fa97efc..e9482f85bdc078 100644
--- a/lld/MachO/Driver.cpp
+++ b/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;
diff --git a/lld/MachO/InputSection.h b/lld/MachO/InputSection.h
index 7ef0e31066f372..8e1f7ea0af01e7 100644
--- a/lld/MachO/InputSection.h
+++ b/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";
diff --git a/lld/test/MachO/cgdata-generate-merge.s b/lld/test/MachO/cgdata-generate-merge.s
new file mode 100644
index 00000000000000..3f7fb6777bc3cf
--- /dev/null
+++ b/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