[lld] Sort code section chunks by range types on Arm64EC targets. (#69099)

5 files changed, 349 insertions, 2 deletions

diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
index d14a258..4e500ca 100644
--- a/lld/COFF/Chunks.h
+++ b/lld/COFF/Chunks.h
@@ -24,10 +24,11 @@
 namespace lld::coff {
 
 using llvm::COFF::ImportDirectoryTableEntry;
-using llvm::object::COFFSymbolRef;
-using llvm::object::SectionRef;
+using llvm::object::chpe_range_type;
 using llvm::object::coff_relocation;
 using llvm::object::coff_section;
+using llvm::object::COFFSymbolRef;
+using llvm::object::SectionRef;
 
 class Baserel;
 class Defined;
@@ -114,6 +115,9 @@ public:
   // synthesized by the linker.
   bool isHotPatchable() const;
 
+  MachineTypes getMachine() const;
+  chpe_range_type getArm64ECRangeType() const;
+
 protected:
   Chunk(Kind k = OtherKind) : chunkKind(k), hasData(true), p2Align(0) {}
 
@@ -164,6 +168,8 @@ public:
   // Collect all locations that contain absolute addresses for base relocations.
   virtual void getBaserels(std::vector<Baserel> *res) {}
 
+  virtual MachineTypes getMachine() const { return IMAGE_FILE_MACHINE_UNKNOWN; }
+
   // Returns a human-readable name of this chunk. Chunks are unnamed chunks of
   // bytes, so this is used only for logging or debugging.
   virtual StringRef getDebugName() const { return ""; }
@@ -420,6 +426,24 @@ inline StringRef Chunk::getDebugName() const {
     return static_cast<const NonSectionChunk *>(this)->getDebugName();
 }
 
+inline MachineTypes Chunk::getMachine() const {
+  if (isa<SectionChunk>(this))
+    return static_cast<const SectionChunk *>(this)->getMachine();
+  else
+    return static_cast<const NonSectionChunk *>(this)->getMachine();
+}
+
+inline chpe_range_type Chunk::getArm64ECRangeType() const {
+  switch (getMachine()) {
+  case AMD64:
+    return chpe_range_type::Amd64;
+  case ARM64EC:
+    return chpe_range_type::Arm64EC;
+  default:
+    return chpe_range_type::Arm64;
+  }
+}
+
 // This class is used to implement an lld-specific feature (not implemented in
 // MSVC) that minimizes the output size by finding string literals sharing tail
 // parts and merging them.
@@ -506,6 +530,7 @@ public:
   explicit ImportThunkChunkX64(COFFLinkerContext &ctx, Defined *s);
   size_t getSize() const override { return sizeof(importThunkX86); }
   void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return AMD64; }
 };
 
 class ImportThunkChunkX86 : public ImportThunkChunk {
@@ -515,6 +540,7 @@ public:
   size_t getSize() const override { return sizeof(importThunkX86); }
   void getBaserels(std::vector<Baserel> *res) override;
   void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return I386; }
 };
 
 class ImportThunkChunkARM : public ImportThunkChunk {
@@ -526,6 +552,7 @@ public:
   size_t getSize() const override { return sizeof(importThunkARM); }
   void getBaserels(std::vector<Baserel> *res) override;
   void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return ARMNT; }
 };
 
 class ImportThunkChunkARM64 : public ImportThunkChunk {
@@ -536,6 +563,7 @@ public:
   }
   size_t getSize() const override { return sizeof(importThunkARM64); }
   void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return ARM64; }
 };
 
 class RangeExtensionThunkARM : public NonSectionChunk {
@@ -546,6 +574,7 @@ public:
   }
   size_t getSize() const override;
   void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return ARMNT; }
 
   Defined *target;
 
@@ -561,6 +590,7 @@ public:
   }
   size_t getSize() const override;
   void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return ARM64; }
 
   Defined *target;
 
diff --git a/lld/COFF/DLL.cpp b/lld/COFF/DLL.cpp
index 59779701..0b337a2 100644
--- a/lld/COFF/DLL.cpp
+++ b/lld/COFF/DLL.cpp
@@ -318,6 +318,7 @@ public:
   ThunkChunkX64(Defined *i, Chunk *tm) : imp(i), tailMerge(tm) {}
 
   size_t getSize() const override { return sizeof(thunkX64); }
+  MachineTypes getMachine() const override { return AMD64; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, thunkX64, sizeof(thunkX64));
@@ -334,6 +335,7 @@ public:
   TailMergeChunkX64(Chunk *d, Defined *h) : desc(d), helper(h) {}
 
   size_t getSize() const override { return sizeof(tailMergeX64); }
+  MachineTypes getMachine() const override { return AMD64; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, tailMergeX64, sizeof(tailMergeX64));
@@ -386,6 +388,7 @@ public:
       : imp(i), tailMerge(tm), ctx(ctx) {}
 
   size_t getSize() const override { return sizeof(thunkX86); }
+  MachineTypes getMachine() const override { return I386; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, thunkX86, sizeof(thunkX86));
@@ -410,6 +413,7 @@ public:
       : desc(d), helper(h), ctx(ctx) {}
 
   size_t getSize() const override { return sizeof(tailMergeX86); }
+  MachineTypes getMachine() const override { return I386; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, tailMergeX86, sizeof(tailMergeX86));
@@ -436,6 +440,7 @@ public:
   }
 
   size_t getSize() const override { return sizeof(thunkARM); }
+  MachineTypes getMachine() const override { return ARMNT; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, thunkARM, sizeof(thunkARM));
@@ -462,6 +467,7 @@ public:
   }
 
   size_t getSize() const override { return sizeof(tailMergeARM); }
+  MachineTypes getMachine() const override { return ARMNT; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, tailMergeARM, sizeof(tailMergeARM));
@@ -487,6 +493,7 @@ public:
   }
 
   size_t getSize() const override { return sizeof(thunkARM64); }
+  MachineTypes getMachine() const override { return ARM64; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, thunkARM64, sizeof(thunkARM64));
@@ -506,6 +513,7 @@ public:
   }
 
   size_t getSize() const override { return sizeof(tailMergeARM64); }
+  MachineTypes getMachine() const override { return ARM64; }
 
   void writeTo(uint8_t *buf) const override {
     memcpy(buf, tailMergeARM64, sizeof(tailMergeARM64));
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index d4f6ee6..43d8e7c 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -212,6 +212,7 @@ private:
   void locateImportTables();
   void createExportTable();
   void mergeSections();
+  void sortECChunks();
   void removeUnusedSections();
   void assignAddresses();
   bool isInRange(uint16_t relType, uint64_t s, uint64_t p, int margin);
@@ -676,6 +677,7 @@ void Writer::run() {
     createMiscChunks();
     createExportTable();
     mergeSections();
+    sortECChunks();
     removeUnusedSections();
     finalizeAddresses();
     removeEmptySections();
@@ -1377,6 +1379,21 @@ void Writer::mergeSections() {
   }
 }
 
+// EC targets may have chunks of various architectures mixed together at this
+// point. Group code chunks of the same architecture together by sorting chunks
+// by their EC range type.
+void Writer::sortECChunks() {
+  if (!isArm64EC(ctx.config.machine))
+    return;
+
+  for (OutputSection *sec : ctx.outputSections) {
+    if (sec->isCodeSection())
+      llvm::stable_sort(sec->chunks, [=](const Chunk *a, const Chunk *b) {
+        return a->getArm64ECRangeType() < b->getArm64ECRangeType();
+      });
+  }
+}
+
 // Visits all sections to assign incremental, non-overlapping RVAs and
 // file offsets.
 void Writer::assignAddresses() {
diff --git a/lld/test/COFF/Inputs/loadconfig-arm64ec.s b/lld/test/COFF/Inputs/loadconfig-arm64ec.s
new file mode 100644
index 0000000..1efd024
--- /dev/null
+++ b/lld/test/COFF/Inputs/loadconfig-arm64ec.s
@@ -0,0 +1,97 @@
+        .section .rdata,"dr"
+        .globl _load_config_used
+        .p2align 3, 0
+_load_config_used:
+        .word 0x140
+        .fill 0x54, 1, 0
+        .xword __security_cookie
+        .fill 0x10, 1, 0
+        .xword __guard_check_icall_fptr
+        .xword __guard_dispatch_icall_fptr
+        .xword __guard_fids_table
+        .xword __guard_fids_count
+        .xword __guard_flags
+        .xword 0
+        .xword __guard_iat_table
+        .xword __guard_iat_count
+        .xword __guard_longjmp_table
+        .xword __guard_longjmp_count
+        .xword 0
+        .xword __chpe_metadata
+        .fill 0x78, 1, 0
+
+__guard_check_icall_fptr:
+        .xword 0
+__guard_dispatch_icall_fptr:
+        .xword 0
+__os_arm64x_dispatch_call_no_redirect:
+        .xword 0
+__os_arm64x_dispatch_ret:
+        .xword 0
+__os_arm64x_check_call:
+        .xword 0
+__os_arm64x_check_icall:
+        .xword 0
+__os_arm64x_get_x64_information:
+        .xword 0
+__os_arm64x_set_x64_information:
+        .xword 0
+__os_arm64x_check_icall_cfg:
+        .xword 0
+__os_arm64x_dispatch_fptr:
+        .xword 0
+__os_arm64x_helper0:
+        .xword 0
+__os_arm64x_helper1:
+        .xword 0
+__os_arm64x_helper2:
+        .xword 0
+__os_arm64x_helper3:
+        .xword 0
+__os_arm64x_helper4:
+        .xword 0
+__os_arm64x_helper5:
+        .xword 0
+__os_arm64x_helper6:
+        .xword 0
+__os_arm64x_helper7:
+        .xword 0
+__os_arm64x_helper8:
+        .xword 0
+
+        .data
+        .globl __chpe_metadata
+        .p2align 3, 0
+__chpe_metadata:
+        .word 1
+        .rva code_map
+        .word code_map_count
+        .word 0 // __x64_code_ranges_to_entry_points
+        .word 0 //__arm64x_redirection_metadata
+        .rva __os_arm64x_dispatch_call_no_redirect
+        .rva __os_arm64x_dispatch_ret
+        .rva __os_arm64x_check_call
+        .rva __os_arm64x_check_icall
+        .rva __os_arm64x_check_icall_cfg
+        .word 0 // __arm64x_native_entrypoint
+        .word 0 // __hybrid_auxiliary_iat
+        .word 0 // __x64_code_ranges_to_entry_points_count
+        .word 0 // __arm64x_redirection_metadata_count
+        .rva __os_arm64x_get_x64_information
+        .rva __os_arm64x_set_x64_information
+        .word 0 // __arm64x_extra_rfe_table
+        .word 0 // __arm64x_extra_rfe_table_size
+        .rva __os_arm64x_dispatch_fptr
+        .word 0 // __hybrid_auxiliary_iat_copy
+        .rva __os_arm64x_helper0
+        .rva __os_arm64x_helper1
+        .rva __os_arm64x_helper2
+        .rva __os_arm64x_helper3
+        .rva __os_arm64x_helper4
+        .rva __os_arm64x_helper5
+        .rva __os_arm64x_helper6
+        .rva __os_arm64x_helper7
+        .rva __os_arm64x_helper8
+
+__security_cookie:
+        .xword 0
diff --git a/lld/test/COFF/arm64ec-codemap.test b/lld/test/COFF/arm64ec-codemap.test
new file mode 100644
index 0000000..424456a
--- /dev/null
+++ b/lld/test/COFF/arm64ec-codemap.test
@@ -0,0 +1,195 @@
+REQUIRES: aarch64, x86
+RUN: split-file %s %t.dir && cd %t.dir
+
+RUN: llvm-mc -filetype=obj -triple=arm64-windows arm64-func-sym.s -o arm64-func-sym.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows arm64ec-func-sym.s -o arm64ec-func-sym.obj
+RUN: llvm-mc -filetype=obj -triple=x86_64-windows x86_64-func-sym.s -o x86_64-func-sym.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows codemap.s -o codemap.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows codemap2.s -o codemap2.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows codemap3.s -o codemap3.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows %S/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
+
+Link ARM64EC DLL and verify that the code is arranged as expected.
+
+RUN: lld-link -out:test.dll -machine:arm64ec arm64ec-func-sym.obj x86_64-func-sym.obj \
+RUN:          codemap.obj loadconfig-arm64ec.obj -dll -noentry
+
+RUN: llvm-readobj --coff-load-config test.dll | FileCheck -check-prefix=CODEMAP %s
+CODEMAP:       CodeMap [
+CODEMAP-NEXT:    0x1000 - 0x1008  ARM64EC
+CODEMAP-NEXT:    0x2000 - 0x2006  X64
+CODEMAP-NEXT:    0x5000 - 0x5008  ARM64EC
+CODEMAP-NEXT:    0x6000 - 0x6006  X64
+CODEMAP-NEXT:  ]
+
+RUN: llvm-objdump -d test.dll | FileCheck -check-prefix=DISASM %s
+DISASM:      Disassembly of section .text:
+DISASM-EMPTY:
+DISASM-NEXT: 0000000180001000 <.text>:
+DISASM-NEXT: 180001000: 52800040     mov     w0, #0x2
+DISASM-NEXT: 180001004: d65f03c0     ret
+DISASM-NEXT:                 ...
+DISASM-NEXT: 180002000: b8 03 00 00 00               movl    $0x3, %eax
+DISASM-NEXT: 180002005: c3                           retq
+DISASM-EMPTY:
+DISASM-NEXT: Disassembly of section test:
+DISASM-EMPTY:
+DISASM-NEXT: 0000000180005000 <test>:
+DISASM-NEXT: 180005000: 528000a0     mov     w0, #0x5
+DISASM-NEXT: 180005004: d65f03c0     ret
+DISASM-NEXT:                 ...
+DISASM-NEXT: 180006000: b8 06 00 00 00               movl    $0x6, %eax
+DISASM-NEXT: 180006005: c3                           retq
+
+Order of arguments doesn't matter in this case, chunks are sorted by target type anyway.
+
+RUN: lld-link -out:test2.dll -machine:arm64ec x86_64-func-sym.obj arm64ec-func-sym.obj \
+RUN:          codemap.obj loadconfig-arm64ec.obj -dll -noentry
+RUN: llvm-readobj --coff-load-config test2.dll | FileCheck -check-prefix=CODEMAP %s
+RUN: llvm-objdump -d test2.dll | FileCheck -check-prefix=DISASM %s
+
+RUN: lld-link -out:testx.dll -machine:arm64x arm64-func-sym.obj arm64ec-func-sym.obj \
+RUN:          x86_64-func-sym.obj codemap2.obj loadconfig-arm64ec.obj -dll -noentry
+
+Do the same with ARM64X target.
+
+RUN: llvm-readobj --coff-load-config testx.dll | FileCheck -check-prefix=CODEMAPX %s
+CODEMAPX:       CodeMap [
+CODEMAPX-NEXT:    0x1000 - 0x1008  ARM64
+CODEMAPX-NEXT:    0x2000 - 0x2008  ARM64EC
+CODEMAPX-NEXT:    0x3000 - 0x3006  X64
+CODEMAPX-NEXT:    0x6000 - 0x6008  ARM64EC
+CODEMAPX-NEXT:    0x7000 - 0x7006  X64
+CODEMAPX-NEXT:  ]
+
+RUN: llvm-objdump -d testx.dll | FileCheck -check-prefix=DISASMX %s
+DISASMX:      Disassembly of section .text:
+DISASMX-EMPTY:
+DISASMX-NEXT: 0000000180001000 <.text>:
+DISASMX-NEXT: 180001000: 528000e0     mov     w0, #0x7
+DISASMX-NEXT: 180001004: d65f03c0     ret
+DISASMX-NEXT:                 ...
+DISASMX-NEXT: 180002000: 52800040     mov     w0, #0x2
+DISASMX-NEXT: 180002004: d65f03c0     ret
+DISASMX-NEXT:                 ...
+DISASMX-NEXT: 180003000: b8 03 00 00 00               movl    $0x3, %eax
+DISASMX-NEXT: 180003005: c3                           retq
+DISASMX-EMPTY:
+DISASMX-NEXT: Disassembly of section test:
+DISASMX-EMPTY:
+DISASMX-NEXT: 0000000180006000 <test>:
+DISASMX-NEXT: 180006000: 528000a0     mov     w0, #0x5
+DISASMX-NEXT: 180006004: d65f03c0     ret
+DISASMX-NEXT:                 ...
+DISASMX-NEXT: 180007000: b8 06 00 00 00               movl    $0x6, %eax
+DISASMX-NEXT: 180007005: c3                           retq
+
+Test merged sections.
+
+RUN: lld-link -out:testm.dll -machine:arm64ec arm64ec-func-sym.obj x86_64-func-sym.obj \
+RUN:          codemap3.obj loadconfig-arm64ec.obj -dll -noentry -merge:test=.text
+
+RUN: llvm-readobj --coff-load-config testm.dll | FileCheck -check-prefix=CODEMAPM %s
+CODEMAPM:      CodeMap [
+CODEMAPM-NEXT:   0x1000 - 0x1010  ARM64EC
+CODEMAPM-NEXT:   0x2000 - 0x3004  X64
+CODEMAPM-NEXT: ]
+
+RUN: llvm-objdump -d testm.dll | FileCheck -check-prefix=DISASMM %s
+DISASMM:      Disassembly of section .text:
+DISASMM-EMPTY:
+DISASMM-NEXT: 0000000180001000 <.text>:
+DISASMM-NEXT: 180001000: 52800040     mov     w0, #0x2
+DISASMM-NEXT: 180001004: d65f03c0     ret
+DISASMM-NEXT: 180001008: 528000a0     mov     w0, #0x5
+DISASMM-NEXT: 18000100c: d65f03c0     ret
+DISASMM-NEXT:                 ...
+DISASMM-NEXT: 180002000: b8 03 00 00 00               movl    $0x3, %eax
+DISASMM-NEXT: 180002005: c3                           retq
+DISASMM-NEXT:                 ...
+DISASMM-NEXT: 180002ffe: 00 00                        addb    %al, (%rax)
+DISASMM-NEXT: 180003000: b8 06 00 00 00               movl    $0x6, %eax
+
+#--- arm64-func-sym.s
+    .text
+    .globl arm64_func_sym
+    .p2align 2, 0x0
+arm64_func_sym:
+    mov w0, #7
+    ret
+
+#--- arm64ec-func-sym.s
+    .text
+    .globl arm64ec_func_sym
+    .p2align 12, 0x0
+arm64ec_func_sym:
+    mov w0, #2
+    ret
+
+    .section test, "xr"
+    .globl arm64ec_func_sym2
+    .p2align 2, 0x0
+arm64ec_func_sym2:
+    mov w0, #5
+    ret
+
+#--- x86_64-func-sym.s
+    .text
+    .globl x86_64_func_sym
+    .p2align 12, 0x0
+x86_64_func_sym:
+    movl $3, %eax
+    retq
+
+    .section test, "xr"
+    .globl x86_64_func_sym2
+    .p2align 12, 0x0
+x86_64_func_sym2:
+    movl $6, %eax
+    retq
+
+#--- codemap.s
+    .section .rdata,"dr"
+    .globl code_map
+code_map:
+    .rva arm64ec_func_sym + 1
+    .word 8
+    .rva x86_64_func_sym + 2
+    .word 6
+    .rva arm64ec_func_sym2 + 1
+    .word 8
+    .rva x86_64_func_sym2 + 2
+    .word 6
+
+    .globl code_map_count
+code_map_count = 4
+
+#--- codemap2.s
+    .section .rdata,"dr"
+    .globl code_map
+code_map:
+    .rva arm64_func_sym
+    .word 8
+    .rva arm64ec_func_sym + 1
+    .word 8
+    .rva x86_64_func_sym + 2
+    .word 6
+    .rva arm64ec_func_sym2 + 1
+    .word 8
+    .rva x86_64_func_sym2 + 2
+    .word 6
+
+    .globl code_map_count
+code_map_count = 5
+
+#--- codemap3.s
+    .section .rdata,"dr"
+    .globl code_map
+code_map:
+    .rva arm64ec_func_sym + 1
+    .word 16
+    .rva x86_64_func_sym + 2
+    .word 0x1004
+
+    .globl code_map_count
+code_map_count = 2