[LLD][COFF] Generate X64 thunks for ARM64EC entry points and patchable functions. (#105499)

cjacek · web-flow · commit a2d8743cc86f · 2024-08-22T22:03:05.000+02:00
This implements Fast-Forward Sequences documented in ARM64EC ABI https://learn.microsoft.com/en-us/windows/arm/arm64ec-abi. There are two conditions when linker should generate such thunks: - For each exported ARM64EC functions. It applies only to ARM64EC functions (we may also have pure x64 functions, for which no thunk is needed). MSVC linker creates `EXP+<mangled export name>` symbol in those cases that points to the thunk and uses that symbol for the export. It's observable from the module: it's possible to reference such symbols as I did in the test. Note that it uses export name, not name of the symbol that's exported (as in `foo` in `/EXPORT:foo=bar`). This implies that if the same function is exported multiple times, it will have multiple thunks. I followed this MSVC behavior. - For hybrid_patchable functions. The linker tries to generate a thunk for each undefined `EXP+*` symbol (and such symbols are created by the compiler as a target of weak alias from the demangled name). MSVC linker tries to find corresponding `*$hp_target` symbol and if fails to do so, it outputs a cryptic error like `LINK : fatal error LNK1000: Internal error during IMAGE::BuildImage`. I just skip generating the thunk in such case (which causes undefined reference error). MSVC linker additionally checks that the symbol complex type is a function (see also #102898). We generally don't do such checks in LLD, so I made it less strict. It should be fine: if it's some data symbol, it will not have `$hp_target` symbol, so we will skip it anyway.
diff --git a/lld/COFF/Chunks.cpp b/lld/COFF/Chunks.cpp
@@ -1073,4 +1073,9 @@ void AbsolutePointerChunk::writeTo(uint8_t *buf) const {
   }
 }
 
+void ECExportThunkChunk::writeTo(uint8_t *buf) const {
+  memcpy(buf, ECExportThunkCode, sizeof(ECExportThunkCode));
+  write32le(buf + 10, target->getRVA() - rva - 14);
+}
+
 } // namespace lld::coff
diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
@@ -749,6 +749,26 @@ class ECCodeMapChunk : public NonSectionChunk {
   std::vector<ECCodeMapEntry> &map;
 };
 
+static const uint8_t ECExportThunkCode[] = {
+    0x48, 0x8b, 0xc4,          // movq    %rsp, %rax
+    0x48, 0x89, 0x58, 0x20,    // movq    %rbx, 0x20(%rax)
+    0x55,                      // pushq   %rbp
+    0x5d,                      // popq    %rbp
+    0xe9, 0,    0,    0,    0, // jmp *0x0
+    0xcc,                      // int3
+    0xcc                       // int3
+};
+
+class ECExportThunkChunk : public NonSectionCodeChunk {
+public:
+  explicit ECExportThunkChunk(Defined *targetSym) : target(targetSym) {}
+  size_t getSize() const override { return sizeof(ECExportThunkCode); };
+  void writeTo(uint8_t *buf) const override;
+  MachineTypes getMachine() const override { return AMD64; }
+
+  Defined *target;
+};
+
 // MinGW specific, for the "automatic import of variables from DLLs" feature.
 // This provides the table of runtime pseudo relocations, for variable
 // references that turned out to need to be imported from a DLL even though
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
@@ -1317,6 +1317,72 @@ void LinkerDriver::convertResources() {
   f->includeResourceChunks();
 }
 
+void LinkerDriver::maybeCreateECExportThunk(StringRef name, Symbol *&sym) {
+  Defined *def;
+  if (!sym)
+    return;
+  if (auto undef = dyn_cast<Undefined>(sym))
+    def = undef->getWeakAlias();
+  else
+    def = dyn_cast<Defined>(sym);
+  if (!def)
+    return;
+
+  if (def->getChunk()->getArm64ECRangeType() != chpe_range_type::Arm64EC)
+    return;
+  StringRef expName;
+  if (auto mangledName = getArm64ECMangledFunctionName(name))
+    expName = saver().save("EXP+" + *mangledName);
+  else
+    expName = saver().save("EXP+" + name);
+  sym = addUndefined(expName);
+  if (auto undef = dyn_cast<Undefined>(sym)) {
+    if (!undef->getWeakAlias()) {
+      auto thunk = make<ECExportThunkChunk>(def);
+      replaceSymbol<DefinedSynthetic>(undef, undef->getName(), thunk);
+    }
+  }
+}
+
+void LinkerDriver::createECExportThunks() {
+  // Check if EXP+ symbols have corresponding $hp_target symbols and use them
+  // to create export thunks when available.
+  for (Symbol *s : ctx.symtab.expSymbols) {
+    if (!s->isUsedInRegularObj)
+      continue;
+    assert(s->getName().starts_with("EXP+"));
+    std::string targetName =
+        (s->getName().substr(strlen("EXP+")) + "$hp_target").str();
+    Symbol *sym = ctx.symtab.find(targetName);
+    if (!sym)
+      continue;
+    Defined *targetSym;
+    if (auto undef = dyn_cast<Undefined>(sym))
+      targetSym = undef->getWeakAlias();
+    else
+      targetSym = dyn_cast<Defined>(sym);
+    if (!targetSym)
+      continue;
+
+    auto *undef = dyn_cast<Undefined>(s);
+    if (undef && !undef->getWeakAlias()) {
+      auto thunk = make<ECExportThunkChunk>(targetSym);
+      replaceSymbol<DefinedSynthetic>(undef, undef->getName(), thunk);
+    }
+    if (!targetSym->isGCRoot) {
+      targetSym->isGCRoot = true;
+      ctx.config.gcroot.push_back(targetSym);
+    }
+  }
+
+  if (ctx.config.entry)
+    maybeCreateECExportThunk(ctx.config.entry->getName(), ctx.config.entry);
+  for (Export &e : ctx.config.exports) {
+    if (!e.data)
+      maybeCreateECExportThunk(e.extName.empty() ? e.name : e.extName, e.sym);
+  }
+}
+
 // In MinGW, if no symbols are chosen to be exported, then all symbols are
 // automatically exported by default. This behavior can be forced by the
 // -export-all-symbols option, so that it happens even when exports are
@@ -2520,6 +2586,9 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (!wrapped.empty())
     wrapSymbols(ctx, wrapped);
 
+  if (isArm64EC(config->machine))
+    createECExportThunks();
+
   // Resolve remaining undefined symbols and warn about imported locals.
   ctx.symtab.resolveRemainingUndefines();
   if (errorCount())
diff --git a/lld/COFF/Driver.h b/lld/COFF/Driver.h
@@ -270,6 +270,10 @@ class LinkerDriver {
   // Convert Windows resource files (.res files) to a .obj file.
   MemoryBufferRef convertResToCOFF(ArrayRef<MemoryBufferRef> mbs,
                                    ArrayRef<ObjFile *> objs);
+
+  // Create export thunks for exported and patchable Arm64EC function symbols.
+  void createECExportThunks();
+  void maybeCreateECExportThunk(StringRef name, Symbol *&sym);
 };
 
 // Create enum with OPT_xxx values for each option in Options.td
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
@@ -551,6 +551,9 @@ std::pair<Symbol *, bool> SymbolTable::insert(StringRef name) {
     sym->pendingArchiveLoad = false;
     sym->canInline = true;
     inserted = true;
+
+    if (isArm64EC(ctx.config.machine) && name.starts_with("EXP+"))
+      expSymbols.push_back(sym);
   }
   return {sym, inserted};
 }
diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
@@ -116,6 +116,9 @@ class SymbolTable {
   // A list of chunks which to be added to .rdata.
   std::vector<Chunk *> localImportChunks;
 
+  // A list of EC EXP+ symbols.
+  std::vector<Symbol *> expSymbols;
+
   // Iterates symbols in non-determinstic hash table order.
   template <typename T> void forEachSymbol(T callback) {
     for (auto &pair : symMap)
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
@@ -304,6 +304,7 @@ class Writer {
   uint64_t sizeOfHeaders;
 
   OutputSection *textSec;
+  OutputSection *hexpthkSec;
   OutputSection *rdataSec;
   OutputSection *buildidSec;
   OutputSection *dataSec;
@@ -984,6 +985,8 @@ void Writer::createSections() {
 
   // Try to match the section order used by link.exe.
   textSec = createSection(".text", code | r | x);
+  if (isArm64EC(ctx.config.machine))
+    hexpthkSec = createSection(".hexpthk", code | r | x);
   createSection(".bss", bss | r | w);
   rdataSec = createSection(".rdata", data | r);
   buildidSec = createSection(".buildid", data | r);
@@ -2046,6 +2049,14 @@ void Writer::maybeAddRVATable(SymbolRVASet tableSymbols, StringRef tableSym,
 
 // Create CHPE metadata chunks.
 void Writer::createECChunks() {
+  for (Symbol *s : ctx.symtab.expSymbols) {
+    auto sym = dyn_cast<Defined>(s);
+    if (!sym || !sym->getChunk())
+      continue;
+    if (auto thunk = dyn_cast<ECExportThunkChunk>(sym->getChunk()))
+      hexpthkSec->addChunk(thunk);
+  }
+
   auto codeMapChunk = make<ECCodeMapChunk>(codeMap);
   rdataSec->addChunk(codeMapChunk);
   Symbol *codeMapSym = ctx.symtab.findUnderscore("__hybrid_code_map");
diff --git a/lld/test/COFF/arm64ec-export-thunks.test b/lld/test/COFF/arm64ec-export-thunks.test
@@ -0,0 +1,188 @@
+REQUIRES: aarch64, x86
+RUN: split-file %s %t.dir && cd %t.dir
+
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows arm64ec-data.s -o arm64ec-data.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows arm64ec-func.s -o arm64ec-func.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows antidep-func.s -o antidep-func.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows arm64ec-data-sym.s -o arm64ec-data-sym.obj
+RUN: llvm-mc -filetype=obj -triple=x86_64-windows x86_64-func.s -o x86_64-func.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows %S/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
+
+RUN: lld-link -out:exports.dll -machine:arm64ec arm64ec-func.obj x86_64-func.obj loadconfig-arm64ec.obj \
+RUN:          arm64ec-data.obj -dll -noentry -export:arm64ec_func -export:func=arm64ec_func \
+RUN:          -export:x86_64_func -export:data_sym,DATA
+
+RUN: llvm-objdump -d exports.dll | FileCheck -check-prefix=EXP-DISASM %s
+EXP-DISASM:      Disassembly of section .text:
+EXP-DISASM-EMPTY:
+EXP-DISASM-NEXT: 0000000180001000 <.text>:
+EXP-DISASM-NEXT: 180001000: 90000008     adrp    x8, 0x180001000 <.text>
+EXP-DISASM-NEXT: 180001004: 52800040     mov     w0, #0x2
+EXP-DISASM-NEXT: 180001008: d65f03c0     ret
+EXP-DISASM-NEXT:                 ...
+EXP-DISASM-EMPTY:
+EXP-DISASM-NEXT: 0000000180002000 <x86_64_func>:
+EXP-DISASM-NEXT: 180002000: e8 fb ef ff ff               callq   0x180001000 <.text>
+EXP-DISASM-NEXT: 180002005: b8 03 00 00 00               movl    $0x3, %eax
+EXP-DISASM-NEXT: 18000200a: c3                           retq
+EXP-DISASM-EMPTY:
+EXP-DISASM-NEXT: Disassembly of section .hexpthk:
+EXP-DISASM-EMPTY:
+EXP-DISASM-NEXT: 0000000180003000 <func>:
+EXP-DISASM-NEXT: 180003000: 48 8b c4                     movq    %rsp, %rax
+EXP-DISASM-NEXT: 180003003: 48 89 58 20                  movq    %rbx, 0x20(%rax)
+EXP-DISASM-NEXT: 180003007: 55                           pushq   %rbp
+EXP-DISASM-NEXT: 180003008: 5d                           popq    %rbp
+EXP-DISASM-NEXT: 180003009: e9 f2 df ff ff               jmp     0x180001000 <.text>
+EXP-DISASM-NEXT: 18000300e: cc                           int3
+EXP-DISASM-NEXT: 18000300f: cc                           int3
+EXP-DISASM-EMPTY:
+EXP-DISASM-NEXT: 0000000180003010 <arm64ec_func>:
+EXP-DISASM-NEXT: 180003010: 48 8b c4                     movq    %rsp, %rax
+EXP-DISASM-NEXT: 180003013: 48 89 58 20                  movq    %rbx, 0x20(%rax)
+EXP-DISASM-NEXT: 180003017: 55                           pushq   %rbp
+EXP-DISASM-NEXT: 180003018: 5d                           popq    %rbp
+EXP-DISASM-NEXT: 180003019: e9 e2 df ff ff               jmp     0x180001000 <.text>
+EXP-DISASM-NEXT: 18000301e: cc                           int3
+EXP-DISASM-NEXT: 18000301f: cc                           int3
+
+RUN: llvm-objdump -p exports.dll | FileCheck -check-prefix=EXP-EXPORT %s
+EXP-EXPORT:      Ordinal      RVA  Name
+EXP-EXPORT-NEXT:       1   0x3010  arm64ec_func
+EXP-EXPORT-NEXT:       2   0x6000  data_sym
+EXP-EXPORT-NEXT:       3   0x3000  func
+EXP-EXPORT-NEXT:       4   0x2000  x86_64_func
+
+RUN: llvm-readobj --coff-load-config exports.dll | FileCheck -check-prefix=EXP-CHPE %s
+EXP-CHPE:       CodeMap [
+EXP-CHPE-NEXT:    0x1000 - 0x100C  ARM64EC
+EXP-CHPE-NEXT:    0x2000 - 0x3020  X64
+EXP-CHPE-NEXT:  ]
+
+RUN: llvm-objdump -s --section=.test exports.dll | FileCheck --check-prefix=EXP-DATA %s
+EXP-DATA: 180006000 00300000 10300000
+
+RUN: lld-link -out:exports2.dll -machine:arm64ec antidep-func.obj x86_64-func.obj loadconfig-arm64ec.obj \
+RUN:          arm64ec-data.obj -dll -noentry -export:arm64ec_func -export:func=arm64ec_func \
+RUN:          -export:x86_64_func -export:data_sym,DATA
+
+RUN: llvm-objdump -d exports2.dll | FileCheck -check-prefix=EXP-DISASM %s
+RUN: llvm-objdump -p exports2.dll | FileCheck -check-prefix=EXP-EXPORT %s
+RUN: llvm-objdump -s --section=.test exports2.dll | FileCheck --check-prefix=EXP-DATA %s
+RUN: llvm-readobj --coff-load-config exports2.dll | FileCheck -check-prefix=EXP-CHPE %s
+
+RUN: lld-link -out:entry.dll -machine:arm64ec arm64ec-func.obj loadconfig-arm64ec.obj -dll -entry:arm64ec_func
+
+RUN: llvm-objdump -d entry.dll | FileCheck -check-prefix=ENTRY-DISASM %s
+ENTRY-DISASM:      Disassembly of section .text:
+ENTRY-DISASM-EMPTY:
+ENTRY-DISASM-NEXT: 0000000180001000 <.text>:
+ENTRY-DISASM-NEXT: 180001000: 90000008     adrp    x8, 0x180001000 <.text>
+ENTRY-DISASM-NEXT: 180001004: 52800040     mov     w0, #0x2                // =2
+ENTRY-DISASM-NEXT: 180001008: d65f03c0     ret
+ENTRY-DISASM-EMPTY:
+ENTRY-DISASM-NEXT: Disassembly of section .hexpthk:
+ENTRY-DISASM-EMPTY:
+ENTRY-DISASM-NEXT: 0000000180002000 <.hexpthk>:
+ENTRY-DISASM-NEXT: 180002000: 48 8b c4                     movq    %rsp, %rax
+ENTRY-DISASM-NEXT: 180002003: 48 89 58 20                  movq    %rbx, 0x20(%rax)
+ENTRY-DISASM-NEXT: 180002007: 55                           pushq   %rbp
+ENTRY-DISASM-NEXT: 180002008: 5d                           popq    %rbp
+ENTRY-DISASM-NEXT: 180002009: e9 f2 ef ff ff               jmp     0x180001000 <.text>
+ENTRY-DISASM-NEXT: 18000200e: cc                           int3
+ENTRY-DISASM-NEXT: 18000200f: cc                           int3
+
+RUN: llvm-readobj --headers entry.dll | FileCheck -check-prefix=ENTRY %s
+ENTRY: AddressOfEntryPoint: 0x2000
+
+RUN: llvm-readobj --coff-load-config entry.dll | FileCheck -check-prefix=ENTRY-CHPE %s
+ENTRY-CHPE:       CodeMap [
+ENTRY-CHPE-NEXT:    0x1000 - 0x100C  ARM64EC
+ENTRY-CHPE-NEXT:    0x2000 - 0x2010  X64
+ENTRY-CHPE-NEXT:  ]
+
+
+Test exporting data symbol as a function:
+
+RUN: lld-link -out:data-func.dll -machine:arm64ec arm64ec-data-sym.obj loadconfig-arm64ec.obj -dll -noentry -export:data_sym
+
+RUN: llvm-readobj --hex-dump=.test data-func.dll | FileCheck --check-prefix=DATAFUNC-TEST %s
+DATAFUNC-TEST:      Hex dump of section '.test':
+DATAFUNC-TEST-NEXT: 0x180003000 00000000                            ....
+
+RUN: llvm-readobj --coff-exports --hex-dump=.test data-func.dll | FileCheck --check-prefix=DATAFUNC-EXP %s
+DATAFUNC-EXP:      Export {
+DATAFUNC-EXP-NEXT:   Ordinal: 1
+DATAFUNC-EXP-NEXT:   Name: data_sym
+DATAFUNC-EXP-NEXT:   RVA: 0x3000
+DATAFUNC-EXP-NEXT: }
+
+
+Test mingw-style auto-export:
+
+RUN: lld-link -out:export-all.dll -machine:arm64ec arm64ec-func.obj loadconfig-arm64ec.obj -dll -noentry -lldmingw
+RUN: llvm-objdump -d export-all.dll | FileCheck --check-prefix=EXPORT-ALL %s
+
+EXPORT-ALL:     Disassembly of section .text:
+EXPORT-ALL-EMPTY:
+EXPORT-ALL-NEXT: 0000000180001000 <.text>:
+EXPORT-ALL-NEXT: 180001000: 90000008     adrp    x8, 0x180001000 <.text>
+EXPORT-ALL-NEXT: 180001004: 52800040     mov     w0, #0x2                // =2
+EXPORT-ALL-NEXT: 180001008: d65f03c0     ret
+EXPORT-ALL-EMPTY:
+EXPORT-ALL-NEXT: Disassembly of section .hexpthk:
+EXPORT-ALL-EMPTY:
+EXPORT-ALL-NEXT: 0000000180002000 <arm64ec_func>:
+EXPORT-ALL-NEXT: 180002000: 48 8b c4                     movq    %rsp, %rax
+EXPORT-ALL-NEXT: 180002003: 48 89 58 20                  movq    %rbx, 0x20(%rax)
+EXPORT-ALL-NEXT: 180002007: 55                           pushq   %rbp
+EXPORT-ALL-NEXT: 180002008: 5d                           popq    %rbp
+EXPORT-ALL-NEXT: 180002009: e9 f2 ef ff ff               jmp     0x180001000 <.text>
+EXPORT-ALL-NEXT: 18000200e: cc                           int3
+EXPORT-ALL-NEXT: 18000200f: cc                           int3
+
+
+#--- arm64ec-func.s
+    .text
+    .globl arm64ec_func
+    .p2align 2, 0x0
+arm64ec_func:
+    adrp x8,arm64ec_func
+    mov w0, #2
+    ret
+
+#--- antidep-func.s
+    .text
+    .globl "#arm64ec_func"
+    .p2align 2, 0x0
+"#arm64ec_func":
+    adrp x8,arm64ec_func
+    mov w0, #2
+    ret
+
+    .weak_anti_dep arm64ec_func
+arm64ec_func = "#arm64ec_func"
+
+#--- arm64ec-data.s
+    .section .test, "r"
+    .globl data_sym
+    .p2align 2, 0x0
+data_sym:
+    .rva "EXP+#func"
+    .rva "EXP+#arm64ec_func"
+
+#--- x86_64-func.s
+    .text
+    .globl x86_64_func
+    .p2align 2, 0x0
+x86_64_func:
+    call arm64ec_func
+    movl $3, %eax
+    retq
+
+#--- arm64ec-data-sym.s
+    .section .test, "r"
+    .globl data_sym
+    .p2align 2, 0x0
+data_sym:
+    .word 0
diff --git a/lld/test/COFF/arm64ec-exports.s b/lld/test/COFF/arm64ec-exports.s
diff --git a/lld/test/COFF/arm64ec-patchable-thunks.test b/lld/test/COFF/arm64ec-patchable-thunks.test

Original file line number	Diff line number	Diff line change
`@@ -1073,4 +1073,9 @@ void AbsolutePointerChunk::writeTo(uint8_t *buf) const {`
`1073`	`1073`	`}`
`1074`	`1074`	`}`
`1075`	`1075`
	`1076`	`+void ECExportThunkChunk::writeTo(uint8_t *buf) const {`
	`1077`	`+ memcpy(buf, ECExportThunkCode, sizeof(ECExportThunkCode));`
	`1078`	`+ write32le(buf + 10, target->getRVA() - rva - 14);`
	`1079`	`+}`
	`1080`	`+`
`1076`	`1081`	`} // namespace lld::coff`
Original file line number	Diff line number	Diff line change
`@@ -551,6 +551,9 @@ std::pair<Symbol *, bool> SymbolTable::insert(StringRef name) {`
`551`	`551`	`sym->pendingArchiveLoad = false;`
`552`	`552`	`sym->canInline = true;`
`553`	`553`	`inserted = true;`
	`554`	`+`
	`555`	`+ if (isArm64EC(ctx.config.machine) && name.starts_with("EXP+"))`
	`556`	`+ expSymbols.push_back(sym);`
`554`	`557`	`}`
`555`	`558`	`return {sym, inserted};`
`556`	`559`	`}`