[lld] [LLD][COFF] Generate redirection metadata for custom ARM64EC export thunks. (PR #105901)

[Jacek Caban via llvm-commits]

https://github.com/cjacek created https://github.com/llvm/llvm-project/pull/105901

This allows using custom export thunks instead of default generated ones. This is useful for performance in cases where transfering between JIT and ARM64EC code is more expensive than just emulating the whole function (but it's still useful to have ARM64EC version so that ARM64EC callers don't call into the emulator). It's also useful for compatibility, where applications have specific expectations about function contents (like syscall functions).

>From 2c7be0acd7f79156189bb6ce6b78ab5097e40f07 Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek at codeweavers.com>
Date: Mon, 19 Aug 2024 11:21:21 +0200
Subject: [PATCH] [LLD][COFF] Generate redirection metadata for custom ARM64EC
 export thunks.

This allows using custom export thunks instead of default generated ones.
This is useful for performance in cases where transfering between JIT and
ARM64EC code is more expensive than just emulating the whole function (but
it's still useful to have ARM64EC version so that ARM64EC callers don't
call into the emulator). It's also useful for compatibility, where
applications have specific expectations about function contents (like
syscall functions).
---
 lld/COFF/Writer.cpp                       | 14 ++++
 lld/test/COFF/arm64ec-cust-export-thunk.s | 82 +++++++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 lld/test/COFF/arm64ec-cust-export-thunk.s

diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 0360e186ecf0cf..5fc05d7556f4f8 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -2062,6 +2062,20 @@ void Writer::createECChunks() {
     if (auto thunk = dyn_cast<ECExportThunkChunk>(sym->getChunk())) {
       hexpthkSec->addChunk(thunk);
       exportThunks.push_back({thunk, thunk->target});
+    } else if (auto def = dyn_cast<DefinedRegular>(sym)) {
+      // Allow section chunk to be treated as an export thunk if it looks like
+      // one.
+      SectionChunk *chunk = def->getChunk();
+      if (!chunk->live || chunk->getMachine() != AMD64)
+        continue;
+      assert(sym->getName().starts_with("EXP+"));
+      StringRef targetName = sym->getName().substr(strlen("EXP+"));
+      Symbol *targetSym = ctx.symtab.find((targetName + "$hp_target").str());
+      if (!targetSym)
+        targetSym = ctx.symtab.find(targetName);
+      Defined *t = dyn_cast_or_null<Defined>(targetSym);
+      if (t && isArm64EC(t->getChunk()->getMachine()))
+        exportThunks.push_back({chunk, t});
     }
   }
 
diff --git a/lld/test/COFF/arm64ec-cust-export-thunk.s b/lld/test/COFF/arm64ec-cust-export-thunk.s
new file mode 100644
index 00000000000000..e56f34c8752173
--- /dev/null
+++ b/lld/test/COFF/arm64ec-cust-export-thunk.s
@@ -0,0 +1,82 @@
+# REQUIRES: aarch64, x86
+# RUN: split-file %s %t.dir && cd %t.dir
+
+# Test that metadata is generated when a custom export thunk is supplied.
+
+# RUN: llvm-mc -filetype=obj -triple=arm64ec-windows func.s -o func.obj
+# RUN: llvm-mc -filetype=obj -triple=arm64ec-windows hp-func.s -o hp-func.obj
+# RUN: llvm-mc -filetype=obj -triple=x86_64-windows thunk.s -o thunk.obj
+# RUN: llvm-mc -filetype=obj -triple=arm64ec-windows %S/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
+
+# RUN: lld-link -out:out.dll -machine:arm64ec func.obj thunk.obj loadconfig-arm64ec.obj -dll -noentry "-export:#func,EXPORTAS,func"
+
+# RUN: llvm-objdump -d out.dll | FileCheck --check-prefixes=DISASM,DISASM-EXP %s
+# DISASM:      Disassembly of section .text:
+# DISASM-EMPTY:
+# DISASM-NEXT: 0000000180001000 <.text>:
+# DISASM-NEXT: 180001000: 52800040     mov     w0, #0x2                // =2
+# DISASM-NEXT: 180001004: d65f03c0     ret
+# DISASM-NEXT:                 ...
+# DISASM-EXP-EMPTY:
+# DISASM-EXP-NEXT: 0000000180002000 <func>:
+# DISASM-NEXT: 180002000: b8 03 00 00 00               movl    $0x3, %eax
+# DISASM-NEXT: 180002005: c3                           retq
+
+# RUN: llvm-objdump -p out.dll | FileCheck --check-prefix=EXPORT %s
+# EXPORT:      Ordinal      RVA  Name
+# EXPORT-NEXT:       1   0x2000  func
+
+# RUN: llvm-readobj --coff-load-config out.dll | FileCheck --check-prefix=CHPE %s
+# CHPE:       CodeMap [
+# CHPE-NEXT:    0x1000 - 0x1008  ARM64EC
+# CHPE-NEXT:    0x2000 - 0x2006  X64
+# CHPE-NEXT:  ]
+# CHPE-NEXT:  CodeRangesToEntryPoints [
+# CHPE-NEXT:    0x2000 - 0x2006 -> 0x2000
+# CHPE-NEXT:  ]
+# CHPE-NEXT:  RedirectionMetadata [
+# CHPE-NEXT:    0x2000 -> 0x1000
+# CHPE-NEXT:  ]
+
+# RUN: lld-link -out:out2.dll -machine:arm64ec hp-func.obj thunk.obj loadconfig-arm64ec.obj -dll -noentry
+# RUN: llvm-objdump -d out2.dll | FileCheck --check-prefix=DISASM %s
+# RUN: llvm-readobj --coff-load-config out2.dll | FileCheck --check-prefix=CHPE %s
+
+#--- func.s
+    .globl "#func"
+    .p2align 2, 0x0
+"#func":
+    mov w0, #2
+    ret
+
+#--- hp-func.s
+    .section .text,"xr",discard,"#func$hp_target"
+    .globl "#func$hp_target"
+    .p2align 2, 0x0
+"#func$hp_target":
+    mov w0, #2
+    ret
+
+    .def "EXP+#func"
+    .scl 2
+    .type 32
+    .endef
+    .weak func
+.set func, "EXP+#func"
+    .weak "#func"
+.set "#func", "#func$hp_target"
+
+    .data
+    .rva func
+
+#--- thunk.s
+    .def "EXP+#func"
+    .scl 2
+    .type 32
+    .endef
+    .section .wowthk$aa,"xr",discard,"EXP+#func"
+    .globl "EXP+#func"
+    .p2align 2, 0x0
+"EXP+#func":
+    movl $3, %eax
+    retq