[lld][COFF] Support .pdata section on ARM64EC targets. (#72521)

ARM64EC needs to handle both ARM and x86_64 exception tables. This is achieved by separating their chunks and sorting them separately. EXCEPTION_TABLE directory references x86_64 variant, while ARM variant is exposed using CHPE metadata, which references __arm64x_extra_rfe_table and __arm64x_extra_rfe_table_size symbols.
alexbatashev · Dec 5, 2023 · 72c6ca6 · 72c6ca6
1 parent 7788ef4
commit 72c6ca6
Showing 4 changed files with 194 additions and 8 deletions.
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
@@ -2361,6 +2361,8 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   ctx.symtab.addAbsolute(mangle("__guard_eh_cont_table"), 0);
 
   if (isArm64EC(config->machine)) {
+    ctx.symtab.addAbsolute("__arm64x_extra_rfe_table", 0);
+    ctx.symtab.addAbsolute("__arm64x_extra_rfe_table_size", 0);
     ctx.symtab.addAbsolute("__hybrid_code_map", 0);
     ctx.symtab.addAbsolute("__hybrid_code_map_count", 0);
   }

diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
@@ -247,6 +247,7 @@ class Writer {
   void maybeAddRVATable(SymbolRVASet tableSymbols, StringRef tableSym,
                         StringRef countSym, bool hasFlag=false);
   void setSectionPermissions();
+  void setECSymbols();
   void writeSections();
   void writeBuildId();
   void writePEChecksum();
@@ -326,6 +327,9 @@ class Writer {
   // files, so we need to keep track of them separately.
   ChunkRange pdata;
 
+  // x86_64 .pdata sections on ARM64EC/ARM64X targets.
+  ChunkRange hybridPdata;
+
   COFFLinkerContext &ctx;
 };
 } // anonymous namespace
@@ -741,6 +745,7 @@ void Writer::run() {
     removeEmptySections();
     assignOutputSectionIndices();
     setSectionPermissions();
+    setECSymbols();
     createSymbolAndStringTable();
 
     if (fileSize > UINT32_MAX)
@@ -1411,8 +1416,28 @@ void Writer::createSymbolAndStringTable() {
 void Writer::mergeSections() {
   llvm::TimeTraceScope timeScope("Merge sections");
   if (!pdataSec->chunks.empty()) {
-    pdata.first = pdataSec->chunks.front();
-    pdata.last = pdataSec->chunks.back();
+    if (isArm64EC(ctx.config.machine)) {
+      // On ARM64EC .pdata may contain both ARM64 and X64 data. Split them by
+      // sorting and store their regions separately.
+      llvm::stable_sort(pdataSec->chunks, [=](const Chunk *a, const Chunk *b) {
+        return (a->getMachine() == AMD64) < (b->getMachine() == AMD64);
+      });
+
+      for (auto chunk : pdataSec->chunks) {
+        if (chunk->getMachine() == AMD64) {
+          hybridPdata.first = chunk;
+          hybridPdata.last = pdataSec->chunks.back();
+          break;
+        }
+
+        if (!pdata.first)
+          pdata.first = chunk;
+        pdata.last = chunk;
+      }
+    } else {
+      pdata.first = pdataSec->chunks.front();
+      pdata.last = pdataSec->chunks.back();
+    }
   }
 
   for (auto &p : ctx.config.merge) {
@@ -1668,10 +1693,15 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
     dir[RESOURCE_TABLE].RelativeVirtualAddress = rsrcSec->getRVA();
     dir[RESOURCE_TABLE].Size = rsrcSec->getVirtualSize();
   }
-  if (pdata.first) {
-    dir[EXCEPTION_TABLE].RelativeVirtualAddress = pdata.first->getRVA();
-    dir[EXCEPTION_TABLE].Size =
-        pdata.last->getRVA() + pdata.last->getSize() - pdata.first->getRVA();
+  // ARM64EC (but not ARM64X) contains x86_64 exception table in data directory.
+  ChunkRange &exceptionTable =
+      ctx.config.machine == ARM64EC ? hybridPdata : pdata;
+  if (exceptionTable.first) {
+    dir[EXCEPTION_TABLE].RelativeVirtualAddress =
+        exceptionTable.first->getRVA();
+    dir[EXCEPTION_TABLE].Size = exceptionTable.last->getRVA() +
+                                exceptionTable.last->getSize() -
+                                exceptionTable.first->getRVA();
   }
   if (relocSec->getVirtualSize()) {
     dir[BASE_RELOCATION_TABLE].RelativeVirtualAddress = relocSec->getRVA();
@@ -2084,6 +2114,24 @@ void Writer::setSectionPermissions() {
   }
 }
 
+// Set symbols used by ARM64EC metadata.
+void Writer::setECSymbols() {
+  if (!isArm64EC(ctx.config.machine))
+    return;
+
+  Symbol *rfeTableSym = ctx.symtab.findUnderscore("__arm64x_extra_rfe_table");
+  replaceSymbol<DefinedSynthetic>(rfeTableSym, "__arm64x_extra_rfe_table",
+                                  pdata.first);
+
+  if (pdata.first) {
+    Symbol *rfeSizeSym =
+        ctx.symtab.findUnderscore("__arm64x_extra_rfe_table_size");
+    cast<DefinedAbsolute>(rfeSizeSym)
+        ->setVA(pdata.last->getRVA() + pdata.last->getSize() -
+                pdata.first->getRVA());
+  }
+}
+
 // Write section contents to a mmap'ed file.
 void Writer::writeSections() {
   llvm::TimeTraceScope timeScope("Write sections");
@@ -2206,6 +2254,10 @@ void Writer::sortExceptionTables() {
   case AMD64:
     sortExceptionTable<EntryX64>(pdata);
     break;
+  case ARM64EC:
+  case ARM64X:
+    sortExceptionTable<EntryX64>(hybridPdata);
+    [[fallthrough]];
   case ARMNT:
   case ARM64:
     sortExceptionTable<EntryArm>(pdata);

diff --git a/lld/test/COFF/Inputs/loadconfig-arm64ec.s b/lld/test/COFF/Inputs/loadconfig-arm64ec.s
@@ -79,8 +79,8 @@ __chpe_metadata:
         .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 __arm64x_extra_rfe_table
+        .word __arm64x_extra_rfe_table_size
         .rva __os_arm64x_dispatch_fptr
         .word 0 // __hybrid_auxiliary_iat_copy
         .rva __os_arm64x_helper0

diff --git a/lld/test/COFF/pdata-arm64ec.test b/lld/test/COFF/pdata-arm64ec.test
@@ -0,0 +1,132 @@
+REQUIRES: aarch64, x86
+RUN: split-file %s %t.dir && cd %t.dir
+
+Test handlign of hybrid .pdata section on ARM64EC target.
+
+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 %p/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
+
+Only arm64ec code:
+
+RUN: lld-link -out:test1.dll -machine:arm64ec arm64ec-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+
+RUN: llvm-readobj --coff-load-config test1.dll | FileCheck -check-prefix=LOADCFG %s
+LOADCFG:      ExtraRFETable: 0x4000
+LOADCFG-NEXT: ExtraRFETableSize: 0x8
+
+RUN: llvm-readobj --headers test1.dll | FileCheck -check-prefix=NODIR %s
+NODIR: ExceptionTableSize: 0x0
+
+RUN: llvm-objdump -s --section=.pdata test1.dll | FileCheck -check-prefix=DATA %s
+DATA: 180004000 00100000 11000001
+
+Only x86_64 code:
+
+RUN: lld-link -out:test2.dll -machine:arm64ec x86_64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+
+RUN: llvm-readobj --coff-load-config test2.dll | FileCheck -check-prefix=NOLOADCFG %s
+NOLOADCFG: ExtraRFETableSize: 0x0
+
+RUN: llvm-readobj --headers test2.dll | FileCheck -check-prefix=DIR %s
+DIR:      ExceptionTableRVA: 0x4000
+DIR-NEXT: ExceptionTableSize: 0xC
+
+RUN: llvm-objdump -s --section=.pdata test2.dll | FileCheck -check-prefix=DATA2 %s
+DATA2: 180004000 00100000 0e100000
+
+Mixed arm64ec and x86_64 code:
+
+RUN: lld-link -out:test3.dll -machine:arm64ec arm64ec-func-sym.obj x86_64-func-sym.obj \
+RUN:          loadconfig-arm64ec.obj -dll -noentry
+
+RUN: llvm-readobj --coff-load-config test3.dll | FileCheck -check-prefix=LOADCFG2 %s
+LOADCFG2:      ExtraRFETable: 0x5000
+LOADCFG2-NEXT: ExtraRFETableSize: 0x8
+
+RUN: llvm-readobj --headers test3.dll | FileCheck -check-prefix=DIR2 %s
+DIR2:      ExceptionTableRVA: 0x5008
+DIR2-NEXT: ExceptionTableSize: 0xC
+
+RUN: llvm-objdump -s --section=.pdata test3.dll | FileCheck -check-prefix=DATA3 %s
+DATA3: 180005000 00100000 11000001 00200000 0e200000
+
+Mixed arm64x code:
+
+RUN: lld-link -out:test4.dll -machine:arm64x arm64-func-sym.obj arm64ec-func-sym.obj \
+RUN:          x86_64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+
+RUN: llvm-readobj --headers test4.dll | FileCheck -check-prefix=DIR3 %s
+DIR3:      ExceptionTableRVA: 0x6000
+DIR3-NEXT: ExceptionTableSize: 0x10
+
+RUN: llvm-objdump -s --section=.pdata test4.dll | FileCheck -check-prefix=DATA4 %s
+DATA4: 180006000 00100000 11000001 00200000 11000001  ......... ......
+DATA4: 180006010 00300000 0e300000
+
+Order of inputs doesn't matter, the data is sorted by type and RVA:
+
+RUN: lld-link -out:test5.dll -machine:arm64ec x86_64-func-sym.obj arm64ec-func-sym.obj \
+RUN:          loadconfig-arm64ec.obj -dll -noentry
+RUN: llvm-readobj --coff-load-config test5.dll | FileCheck -check-prefix=LOADCFG2 %s
+RUN: llvm-readobj --headers test5.dll | FileCheck -check-prefix=DIR2 %s
+RUN: llvm-objdump -s --section=.pdata test5.dll | FileCheck -check-prefix=DATA3 %s
+
+RUN: lld-link -out:test6.dll -machine:arm64x arm64ec-func-sym.obj x86_64-func-sym.obj \
+RUN:          arm64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+RUN: llvm-readobj --headers test6.dll | FileCheck -check-prefix=DIR3 %s
+RUN: llvm-objdump -s --section=.pdata test6.dll | FileCheck -check-prefix=DATA4 %s
+
+RUN: lld-link -out:test7.dll -machine:arm64x x86_64-func-sym.obj arm64ec-func-sym.obj \
+RUN:          arm64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+RUN: llvm-readobj --headers test7.dll | FileCheck -check-prefix=DIR3 %s
+RUN: llvm-objdump -s --section=.pdata test7.dll | FileCheck -check-prefix=DATA4 %s
+
+#--- arm64-func-sym.s
+    .text
+    .globl arm64_func_sym
+    .p2align 2, 0x0
+arm64_func_sym:
+    .seh_proc arm64_func_sym
+    sub sp, sp, #32
+    .seh_stackalloc 32
+    .seh_endprologue
+    mov w0, #2
+    .seh_startepilogue
+    add sp, sp, #32
+    .seh_stackalloc 32
+    .seh_endepilogue
+    ret
+    .seh_endproc
+
+#--- arm64ec-func-sym.s
+    .text
+    .globl arm64ec_func_sym
+    .p2align 2, 0x0
+arm64ec_func_sym:
+    .seh_proc arm64ec_func_sym
+    sub sp, sp, #32
+    .seh_stackalloc 32
+    .seh_endprologue
+    mov w0, #3
+    .seh_startepilogue
+    add sp, sp, #32
+    .seh_stackalloc 32
+    .seh_endepilogue
+    ret
+    .seh_endproc
+
+#--- x86_64-func-sym.s
+    .text
+    .globl x86_64_func_sym
+    .p2align 2, 0x0
+x86_64_func_sym:
+    .seh_proc x86_64_func_sym
+    subq    $40, %rsp
+    .seh_stackalloc 40
+    .seh_endprologue
+    movl $4, %eax
+    addq $40, %rsp
+    retq
+    .seh_endproc