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

[Jacek Caban via llvm-commits]

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

>From 539517ccc98bf154f5164526b426460d18a9d38c 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                       | 18 +++++
 lld/test/COFF/arm64ec-cust-export-thunk.s | 82 +++++++++++++++++++++++
 2 files changed, 100 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..35e0f98926ee8a 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -2062,6 +2062,24 @@ 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+"));
+      // If EXP+#foo is an export thunk of a hybrid patchable function,
+      // we should use the #foo$hp_target symbol as the redirection target.
+      // First, try to look up the $hp_target symbol. If it can't be found,
+      // assume it's a regular function and look for #foo instead.
+      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