[LLD][COFF] Add support for ARM64EC auxiliary IAT (#108304)

In addition to the regular IAT, ARM64EC also includes an auxiliary IAT.
At runtime, the regular IAT is populated with the addresses of imported
functions, which may be x86_64 functions or the export thunks of ARM64EC
functions. The auxiliary IAT contains versions of functions that are
guaranteed to be directly callable by ARM64 code.

The linker fills the auxiliary IAT with the addresses of `__impchk_`
thunks. These thunks perform a call on the IAT address using
`__icall_helper_arm64ec` with the target address from the IAT. If the
imported function is an ARM64EC function, the OS may replace the address
in the auxiliary IAT with the address of the ARM64EC version of the
function (not its export thunk), avoiding the runtime call checker for
better performance.

Author: cjacek

lld/COFF/DLL.cpp

@@ -142,6 +142,30 @@ class NullChunk : public NonSectionChunk {
   size_t size;
 };
 
+// A chunk for ARM64EC auxiliary IAT.
+class AuxImportChunk : public NonSectionChunk {
+public:
+  explicit AuxImportChunk(ImportFile *file) : file(file) {
+    setAlignment(sizeof(uint64_t));
+  }
+  size_t getSize() const override { return sizeof(uint64_t); }
+
+  void writeTo(uint8_t *buf) const override {
+    uint64_t impchkVA = 0;
+    if (file->impchkThunk)
+      impchkVA = file->impchkThunk->getRVA() + file->ctx.config.imageBase;
+    write64le(buf, impchkVA);
+  }
+
+  void getBaserels(std::vector<Baserel> *res) override {
+    if (file->impchkThunk)
+      res->emplace_back(rva, file->ctx.config.machine);
+  }
+
+private:
+  ImportFile *file;
+};
+
 static std::vector<std::vector<DefinedImportData *>>
 binImports(COFFLinkerContext &ctx,
            const std::vector<DefinedImportData *> &imports) {
@@ -160,7 +184,15 @@ binImports(COFFLinkerContext &ctx,
     // Sort symbols by name for each group.
     std::vector<DefinedImportData *> &syms = kv.second;
     llvm::sort(syms, [](DefinedImportData *a, DefinedImportData *b) {
-      return a->getName() < b->getName();
+      auto getBaseName = [](DefinedImportData *sym) {
+        StringRef name = sym->getName();
+        name.consume_front("__imp_");
+        // Skip aux_ part of ARM64EC function symbol name.
+        if (sym->file->impchkThunk)
+          name.consume_front("aux_");
+        return name;
+      };
+      return getBaseName(a) < getBaseName(b);
     });
     v.push_back(std::move(syms));
   }
@@ -687,16 +719,24 @@ void IdataContents::create(COFFLinkerContext &ctx) {
       if (s->getExternalName().empty()) {
         lookups.push_back(make<OrdinalOnlyChunk>(ctx, ord));
         addresses.push_back(make<OrdinalOnlyChunk>(ctx, ord));
-        continue;
+      } else {
+        auto *c = make<HintNameChunk>(s->getExternalName(), ord);
+        lookups.push_back(make<LookupChunk>(ctx, c));
+        addresses.push_back(make<LookupChunk>(ctx, c));
+        hints.push_back(c);
+      }
+
+      if (s->file->impECSym) {
+        auto chunk = make<AuxImportChunk>(s->file);
+        auxIat.push_back(chunk);
+        s->file->impECSym->setLocation(chunk);
       }
-      auto *c = make<HintNameChunk>(s->getExternalName(), ord);
-      lookups.push_back(make<LookupChunk>(ctx, c));
-      addresses.push_back(make<LookupChunk>(ctx, c));
-      hints.push_back(c);
     }
     // Terminate with null values.
     lookups.push_back(make<NullChunk>(ctx.config.wordsize));
     addresses.push_back(make<NullChunk>(ctx.config.wordsize));
+    if (ctx.config.machine == ARM64EC)
+      auxIat.push_back(make<NullChunk>(ctx.config.wordsize));
 
     for (int i = 0, e = syms.size(); i < e; ++i)
       syms[i]->setLocation(addresses[base + i]);

lld/COFF/DLL.h

@@ -31,6 +31,7 @@ class IdataContents {
   std::vector<Chunk *> addresses;
   std::vector<Chunk *> hints;
   std::vector<Chunk *> dllNames;
+  std::vector<Chunk *> auxIat;
 };
 
 // Windows-specific.

lld/COFF/Driver.cpp

@@ -2447,6 +2447,7 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
     ctx.symtab.addAbsolute("__arm64x_extra_rfe_table_size", 0);
     ctx.symtab.addAbsolute("__arm64x_redirection_metadata", 0);
     ctx.symtab.addAbsolute("__arm64x_redirection_metadata_count", 0);
+    ctx.symtab.addAbsolute("__hybrid_auxiliary_iat", 0);
     ctx.symtab.addAbsolute("__hybrid_code_map", 0);
     ctx.symtab.addAbsolute("__hybrid_code_map_count", 0);
     ctx.symtab.addAbsolute("__x64_code_ranges_to_entry_points", 0);

lld/COFF/InputFiles.cpp

@@ -1071,19 +1071,39 @@ void ImportFile::parse() {
   this->hdr = hdr;
   externalName = extName;
 
-  impSym = ctx.symtab.addImportData(impName, this);
+  bool isCode = hdr->getType() == llvm::COFF::IMPORT_CODE;
+
+  if (ctx.config.machine != ARM64EC) {
+    impSym = ctx.symtab.addImportData(impName, this, location);
+  } else {
+    // In addition to the regular IAT, ARM64EC also contains an auxiliary IAT,
+    // which holds addresses that are guaranteed to be callable directly from
+    // ARM64 code. Function symbol naming is swapped: __imp_ symbols refer to
+    // the auxiliary IAT, while __imp_aux_ symbols refer to the regular IAT. For
+    // data imports, the naming is reversed.
+    StringRef auxImpName = saver().save("__imp_aux_" + name);
+    if (isCode) {
+      impSym = ctx.symtab.addImportData(auxImpName, this, location);
+      impECSym = ctx.symtab.addImportData(impName, this, auxLocation);
+    } else {
+      impSym = ctx.symtab.addImportData(impName, this, location);
+      impECSym = ctx.symtab.addImportData(auxImpName, this, auxLocation);
+    }
+    if (!impECSym)
+      return;
+  }
   // If this was a duplicate, we logged an error but may continue;
   // in this case, impSym is nullptr.
   if (!impSym)
     return;
 
   if (hdr->getType() == llvm::COFF::IMPORT_CONST)
-    static_cast<void>(ctx.symtab.addImportData(name, this));
+    static_cast<void>(ctx.symtab.addImportData(name, this, location));
 
   // If type is function, we need to create a thunk which jump to an
   // address pointed by the __imp_ symbol. (This allows you to call
   // DLL functions just like regular non-DLL functions.)
-  if (hdr->getType() == llvm::COFF::IMPORT_CODE) {
+  if (isCode) {
     if (ctx.config.machine != ARM64EC) {
       thunkSym = ctx.symtab.addImportThunk(name, impSym, makeImportThunk());
     } else {

lld/COFF/InputFiles.h

@@ -362,6 +362,10 @@ class ImportFile : public InputFile {
   const coff_import_header *hdr;
   Chunk *location = nullptr;
 
+  // Auxiliary IAT symbol and chunk on ARM64EC.
+  DefinedImportData *impECSym = nullptr;
+  Chunk *auxLocation = nullptr;
+
   // We want to eliminate dllimported symbols if no one actually refers to them.
   // These "Live" bits are used to keep track of which import library members
   // are actually in use.

lld/COFF/SymbolTable.cpp

@@ -584,7 +584,7 @@ void SymbolTable::initializeECThunks() {
 
     Symbol *sym = exitThunks.lookup(file->thunkSym);
     if (!sym)
-      sym = exitThunks.lookup(file->impSym);
+      sym = exitThunks.lookup(file->impECSym);
     file->impchkThunk->exitThunk = dyn_cast_or_null<Defined>(sym);
   }
 }
@@ -785,11 +785,12 @@ Symbol *SymbolTable::addCommon(InputFile *f, StringRef n, uint64_t size,
   return s;
 }
 
-DefinedImportData *SymbolTable::addImportData(StringRef n, ImportFile *f) {
+DefinedImportData *SymbolTable::addImportData(StringRef n, ImportFile *f,
+                                              Chunk *&location) {
   auto [s, wasInserted] = insert(n, nullptr);
   s->isUsedInRegularObj = true;
   if (wasInserted || isa<Undefined>(s) || s->isLazy()) {
-    replaceSymbol<DefinedImportData>(s, n, f);
+    replaceSymbol<DefinedImportData>(s, n, f, location);
     return cast<DefinedImportData>(s);
   }
 

lld/COFF/SymbolTable.h

@@ -103,7 +103,8 @@ class SymbolTable {
   Symbol *addCommon(InputFile *f, StringRef n, uint64_t size,
                     const llvm::object::coff_symbol_generic *s = nullptr,
                     CommonChunk *c = nullptr);
-  DefinedImportData *addImportData(StringRef n, ImportFile *f);
+  DefinedImportData *addImportData(StringRef n, ImportFile *f,
+                                   Chunk *&location);
   Symbol *addImportThunk(StringRef name, DefinedImportData *s,
                          ImportThunkChunk *chunk);
   void addLibcall(StringRef name);

lld/COFF/Symbols.h

@@ -354,23 +354,23 @@ class Undefined : public Symbol {
 // table in an output. The former has "__imp_" prefix.
 class DefinedImportData : public Defined {
 public:
-  DefinedImportData(StringRef n, ImportFile *f)
-      : Defined(DefinedImportDataKind, n), file(f) {
-  }
+  DefinedImportData(StringRef n, ImportFile *file, Chunk *&location)
+      : Defined(DefinedImportDataKind, n), file(file), location(location) {}
 
   static bool classof(const Symbol *s) {
     return s->kind() == DefinedImportDataKind;
   }
 
-  uint64_t getRVA() { return file->location->getRVA(); }
-  Chunk *getChunk() { return file->location; }
-  void setLocation(Chunk *addressTable) { file->location = addressTable; }
+  uint64_t getRVA() { return getChunk()->getRVA(); }
+  Chunk *getChunk() { return location; }
+  void setLocation(Chunk *addressTable) { location = addressTable; }
 
   StringRef getDLLName() { return file->dllName; }
   StringRef getExternalName() { return file->externalName; }
   uint16_t getOrdinal() { return file->hdr->OrdinalHint; }
 
   ImportFile *file;
+  Chunk *&location;
 
   // This is a pointer to the synthetic symbol associated with the load thunk
   // for this symbol that will be called if the DLL is delay-loaded. This is

lld/COFF/Writer.cpp

@@ -914,6 +914,8 @@ void Writer::addSyntheticIdata() {
   if (!idata.hints.empty())
     add(".idata$6", idata.hints);
   add(".idata$7", idata.dllNames);
+  if (!idata.auxIat.empty())
+    add(".idata$9", idata.auxIat);
 }
 
 void Writer::appendECImportTables() {
@@ -936,6 +938,15 @@ void Writer::appendECImportTables() {
     rdataSec->contribSections.insert(rdataSec->contribSections.begin(),
                                      importAddresses);
   }
+
+  // The auxiliary IAT is always placed at the end of the .rdata section
+  // and is aligned to 4KB.
+  if (PartialSection *auxIat = findPartialSection(".idata$9", rdata)) {
+    auxIat->chunks.front()->setAlignment(0x1000);
+    rdataSec->chunks.insert(rdataSec->chunks.end(), auxIat->chunks.begin(),
+                            auxIat->chunks.end());
+    rdataSec->addContributingPartialSection(auxIat);
+  }
 }
 
 // Locate the first Chunk and size of the import directory list and the
@@ -1095,7 +1106,8 @@ void Writer::createSections() {
 
     // ARM64EC has specific placement and alignment requirements for the IAT.
     // Delay adding its chunks until appendECImportTables.
-    if (isArm64EC(ctx.config.machine) && pSec->name == ".idata$5")
+    if (isArm64EC(ctx.config.machine) &&
+        (pSec->name == ".idata$5" || pSec->name == ".idata$9"))
       continue;
 
     OutputSection *sec = createSection(name, outChars);
@@ -2254,6 +2266,11 @@ void Writer::setECSymbols() {
   Symbol *entryPointCountSym =
       ctx.symtab.findUnderscore("__arm64x_redirection_metadata_count");
   cast<DefinedAbsolute>(entryPointCountSym)->setVA(exportThunks.size());
+
+  Symbol *iatSym = ctx.symtab.findUnderscore("__hybrid_auxiliary_iat");
+  replaceSymbol<DefinedSynthetic>(iatSym, "__hybrid_auxiliary_iat",
+                                  idata.auxIat.empty() ? nullptr
+                                                       : idata.auxIat.front());
 }
 
 // Write section contents to a mmap'ed file.

lld/test/COFF/Inputs/loadconfig-arm64ec.s

@@ -76,7 +76,7 @@ __chpe_metadata:
         .rva __os_arm64x_check_icall
         .rva __os_arm64x_check_icall_cfg
         .word 0 // __arm64x_native_entrypoint
-        .word 0 // __hybrid_auxiliary_iat
+        .rva __hybrid_auxiliary_iat
         .word __x64_code_ranges_to_entry_points_count
         .word __arm64x_redirection_metadata_count
         .rva __os_arm64x_get_x64_information

lld/test/COFF/arm64ec-import.test

@@ -63,14 +63,37 @@ DISASM-NEXT: 180002000: ff 25 02 10 00 00            jmpq    *0x1002(%rip)
 
 RUN: llvm-readobj --hex-dump=.test out.dll | FileCheck --check-prefix=TESTSEC %s
 RUN: llvm-readobj --hex-dump=.test out2.dll | FileCheck --check-prefix=TESTSEC %s
-TESTSEC:      0x180006000 08300000 00300000 10300000 20300000
-TESTSEC-NEXT: 0x180006010 08100000 1c100000          00200000
+TESTSEC:      0x180007000 08500000 00300000 10500000 20500000
+TESTSEC-NEXT: 0x180007010 08300000 00500000 10300000 20300000
+TESTSEC-NEXT: 0x180007020 08100000 1c100000 00200000
 
 RUN: llvm-readobj --headers out.dll | FileCheck -check-prefix=HEADERS %s
 HEADERS:  LoadConfigTableRVA: 0x4010
 HEADERS:  IATRVA: 0x3000
 HEADERS:  IATSize: 0x1000
 
+RUN: llvm-readobj --coff-load-config out.dll | FileCheck -check-prefix=LOADCONFIG %s
+LOADCONFIG: AuxiliaryIAT: 0x5000
+
+RUN: llvm-readobj --hex-dump=.rdata out.dll | FileCheck -check-prefix=RDATA %s
+RDATA:      0x180005000 00000000 00000000 08100080 01000000
+RDATA-NEXT: 0x180005010 1c100080 01000000 00000000 00000000
+RDATA-NEXT: 0x180005020 30100080 01000000 00000000 00000000
+
+RUN: llvm-readobj --coff-basereloc out.dll | FileCheck -check-prefix=BASERELOC %s
+BASERELOC:      BaseReloc [
+BASERELOC-NOT:      Address: 0x5000
+BASERELOC:          Address: 0x5008
+BASERELOC-NEXT:   }
+BASERELOC-NEXT:   Entry {
+BASERELOC-NEXT:     Type: DIR64
+BASERELOC-NEXT:     Address: 0x5010
+BASERELOC-NEXT:   }
+BASERELOC-NEXT:   Entry {
+BASERELOC-NEXT:     Type: DIR64
+BASERELOC-NEXT:     Address: 0x5020
+BASERELOC-NEXT:   }
+
 #--- test.s
     .section .test, "r"
     .globl arm64ec_data_sym
@@ -80,6 +103,10 @@ arm64ec_data_sym:
     .rva __imp_data
     .rva __imp_func2
     .rva __imp_t2func
+    .rva __imp_aux_func
+    .rva __imp_aux_data
+    .rva __imp_aux_func2
+    .rva __imp_aux_t2func
     .rva __impchk_func
     .rva __impchk_func2
     .rva func