[WebAssembly] WASIP3 Library Call Thread Context Support (#175800)

The [WebAssembly Component
Model](https://component-model.bytecodealliance.org/) has added support
for [cooperative
multithreading](WebAssembly/component-model#557).
This has been implemented in the [Wasmtime
engine](bytecodealliance/wasmtime#11751) and is
part of the wider project of [WASI preview
3](https://wasi.dev/roadmap#upcoming-wasi-03-releases), which is
currently tracked
[here](https://github.com/orgs/bytecodealliance/projects/16).

These changes require updating the way that `__stack_pointer` and
`__tls_base` work purely for a new `wasm32-wasip3` target; other targets
will not be touched. Specifically, rather than using a Wasm global for
tracking the stack pointer and TLS base, the new
[`context.get/set`](https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md#-canon-contextget)
component model builtin functions will be used (the intention being that
runtimes will need to aggressively optimize these calls into single
load/stores). For justification on this choice rather than switching out
the global at context-switch boundaries, see [this
comment](WebAssembly/wasi-libc#691 (comment))
and [this
comment](WebAssembly/wasi-libc#691 (comment)).

This PR adds support for using library calls instead of globals for
holding the stack pointer and TLS base. When used, this thread context
ABI emits calls to `__wasm_{get,set}_{stack_pointer,tls_base}` when
needed. These functions can then be implemented in `libc`. This is
enabled only for the WASIp3 target.

There is a temporary macro define for `__wasm_libcall_thread_context__`
which can be removed once `wasi-libc` has fully migrated to the new ABI
for the WASIp3 target.

Author: TartanLlama

clang/lib/Basic/Targets/WebAssembly.cpp

@@ -123,6 +123,8 @@ void WebAssemblyTargetInfo::getTargetDefines(const LangOptions &Opts,
     Builder.defineMacro("__wasm_tail_call__");
   if (HasWideArithmetic)
     Builder.defineMacro("__wasm_wide_arithmetic__");
+  if (HasLibcallThreadContext)
+    Builder.defineMacro("__wasm_libcall_thread_context__");
   // Note that not all wasm features appear here.   For example,
   // HasCompatctImports
 

clang/lib/Basic/Targets/WebAssembly.h

@@ -68,6 +68,7 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyTargetInfo : public TargetInfo {
   bool HasExtendedConst = false;
   bool HasFP16 = false;
   bool HasGC = false;
+  bool HasLibcallThreadContext = false;
   bool HasMultiMemory = false;
   bool HasMultivalue = false;
   bool HasMutableGlobals = false;
@@ -110,6 +111,8 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyTargetInfo : public TargetInfo {
       PtrDiffType = SignedLong;
       IntPtrType = SignedLong;
     }
+    if (T.getOS() == llvm::Triple::WASIp3)
+      HasLibcallThreadContext = true;
   }
 
   StringRef getABI() const override;

clang/lib/Driver/ToolChains/WebAssembly.cpp

@@ -88,6 +88,11 @@ static bool WantsPthread(const llvm::Triple &Triple, const ArgList &Args) {
   return WantsPthread;
 }
 
+static bool WantsLibcallThreadContext(const llvm::Triple &Triple,
+                                      const ArgList &Args) {
+  return Triple.getOS() == llvm::Triple::WASIp3;
+}
+
 void wasm::Linker::ConstructJob(Compilation &C, const JobAction &JA,
                                 const InputInfo &Output,
                                 const InputInfoList &Inputs,
@@ -169,6 +174,9 @@ void wasm::Linker::ConstructJob(Compilation &C, const JobAction &JA,
 
   AddLinkerInputs(ToolChain, Inputs, Args, CmdArgs, JA);
 
+  if (WantsLibcallThreadContext(ToolChain.getTriple(), Args))
+    CmdArgs.push_back("--libcall-thread-context");
+
   if (WantsPthread(ToolChain.getTriple(), Args))
     CmdArgs.push_back("--shared-memory");
 

clang/test/Preprocessor/wasm-target-features.c

@@ -217,6 +217,7 @@
 // GENERIC-NOT: #define __wasm_simd128__ 1{{$}}
 // GENERIC-NOT: #define __wasm_tail_call__ 1{{$}}
 // GENERIC-NOT: #define __wasm_wide_arithmetic__ 1{{$}}
+// GENERIC-NOT: #define __wasm_libcall_thread_context__ 1{{$}}
 
 // RUN: %clang -E -dM %s -o - 2>&1 \
 // RUN:     -target wasm32-unknown-unknown -mcpu=bleeding-edge \
@@ -251,3 +252,12 @@
 // RUN:   | FileCheck %s -check-prefix=BLEEDING-EDGE-NO-SIMD128
 //
 // BLEEDING-EDGE-NO-SIMD128-NOT: #define __wasm_simd128__ 1{{$}}
+
+// RUN: %clang -E -dM %s -o - 2>&1 \
+// RUN:     -target wasm32-wasip3 \
+// RUN:   | FileCheck %s -check-prefix=LIBCALL-THREAD-CONTEXT
+// RUN: %clang -E -dM %s -o - 2>&1 \
+// RUN:     -target wasm64-wasip3 \
+// RUN:   | FileCheck %s -check-prefix=LIBCALL-THREAD-CONTEXT
+
+// LIBCALL-THREAD-CONTEXT: #define __wasm_libcall_thread_context__ 1{{$}}

lld/test/wasm/stack-pointer-abi.s

@@ -0,0 +1,13 @@
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.o %s
+# RUN: wasm-ld --libcall-thread-context -o %t.libcall.wasm %t.o
+# RUN: obj2yaml %t.libcall.wasm | FileCheck %s --check-prefix=LIBCALL
+# RUN: wasm-ld -o %t.global.wasm %t.o
+# RUN: obj2yaml %t.global.wasm | FileCheck %s --check-prefix=GLOBAL
+
+  .globl  _start
+_start:
+  .functype _start () -> ()
+  end_function
+
+# LIBCALL: Name: __init_stack_pointer
+# GLOBAL: Name: __stack_pointer

lld/test/wasm/thread-context-abi-mismatch.s

@@ -0,0 +1,22 @@
+# Test that linking object files with mismatched thread context ABIs fails with an error.
+# The presence of an import of __stack_pointer from the env module should be treated 
+# as an indication that the global thread context ABI is being used.
+
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.o %s
+# RUN: not wasm-ld --libcall-thread-context %t.o -o %t.wasm 2>&1 | FileCheck %s
+
+# CHECK: object file uses globals for thread context, but --libcall-thread-context was specified
+
+.globl _start
+_start:
+  .functype _start () -> ()
+  end_function
+
+.globaltype __stack_pointer, i32
+
+.globl use_stack_pointer
+use_stack_pointer:
+  .functype use_stack_pointer () -> ()
+  global.get __stack_pointer
+  drop
+  end_function

lld/test/wasm/tls-libcall.s

@@ -0,0 +1,71 @@
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.o %s
+# RUN: wasm-ld --libcall-thread-context --shared-memory -no-gc-sections -o %t.wasm %t.o
+# RUN: obj2yaml %t.wasm | FileCheck %s
+# RUN: llvm-objdump -d --no-print-imm-hex --no-show-raw-insn %t.wasm | FileCheck %s --check-prefix=DIS
+
+.globl         __wasm_get_tls_base
+__wasm_get_tls_base:
+    .functype   __wasm_get_tls_base () -> (i32)
+    i32.const 0
+    end_function
+
+.globl _start
+_start:
+  .functype _start () -> (i32)
+  call __wasm_get_tls_base
+  i32.const tls1@TLSREL
+  i32.add
+  i32.load 0
+  call __wasm_get_tls_base
+  i32.const tls2@TLSREL
+  i32.add
+  i32.load 0
+  i32.add
+  end_function
+
+.section  .tdata.tls1,"",@
+.globl  tls1
+tls1:
+  .int32  1
+  .size tls1, 4
+
+.section  .tdata.tls2,"",@
+.globl  tls2
+tls2:
+  .int32  2
+  .size tls2, 4
+
+.section  .custom_section.target_features,"",@
+  .int8 2
+  .int8 43
+  .int8 11
+  .ascii  "bulk-memory"
+  .int8 43
+  .int8 7
+  .ascii  "atomics"
+
+
+# CHECK:      GlobalNames:
+# CHECK-NEXT:      - Index:           0
+# CHECK-NEXT:        Name:            __init_stack_pointer
+# CHECK-NEXT:      - Index:           1
+# CHECK-NEXT:        Name:            __init_tls_base
+# CHECK-NEXT:      - Index:           2
+# CHECK-NEXT:        Name:            __tls_size
+# CHECK-NEXT:      - Index:           3
+# CHECK-NEXT:        Name:            __tls_align
+
+# DIS-LABEL: <__wasm_init_memory>:
+
+# DIS-LABEL: <_start>:
+# DIS-EMPTY:
+# DIS-NEXT:       call    4
+# DIS-NEXT:       i32.const       0
+# DIS-NEXT:       i32.add 
+# DIS-NEXT:       i32.load        0
+# DIS-NEXT:       call    4
+# DIS-NEXT:       i32.const       4
+# DIS-NEXT:       i32.add 
+# DIS-NEXT:       i32.load        0
+# DIS-NEXT:       i32.add 
+# DIS-NEXT:       end

lld/wasm/Config.h

@@ -35,6 +35,7 @@ class Symbol;
 class DefinedData;
 class GlobalSymbol;
 class DefinedFunction;
+class UndefinedFunction;
 class DefinedGlobal;
 class UndefinedGlobal;
 class TableSymbol;
@@ -64,6 +65,7 @@ struct Config {
   bool growableTable;
   bool gcSections;
   llvm::StringSet<> keepSections;
+  bool libcallThreadContext;
   std::optional<std::pair<llvm::StringRef, llvm::StringRef>> memoryImport;
   std::optional<llvm::StringRef> memoryExport;
   bool sharedMemory;
@@ -252,6 +254,14 @@ struct Ctx {
     // Used as an address space for function pointers, with each function that
     // is used as a function pointer being allocated a slot.
     TableSymbol *indirectFunctionTable;
+
+    // __wasm_set_tls_base
+    // Function used to set TLS base in libcall thread context modules.
+    UndefinedFunction *setTLSBase;
+
+    // __wasm_get_tls_base
+    // Function used to get TLS base in libcall thread context modules.
+    UndefinedFunction *getTLSBase;
   };
   WasmSym sym;
 

lld/wasm/Driver.cpp

@@ -561,6 +561,7 @@ static void readConfigs(opt::InputArgList &args) {
   ctx.arg.soName = args.getLastArgValue(OPT_soname);
   ctx.arg.importTable = args.hasArg(OPT_import_table);
   ctx.arg.importUndefined = args.hasArg(OPT_import_undefined);
+  ctx.arg.libcallThreadContext = args.hasArg(OPT_libcall_thread_context);
   ctx.arg.ltoo = args::getInteger(args, OPT_lto_O, 2);
   if (ctx.arg.ltoo > 3)
     error("invalid optimization level for LTO: " + Twine(ctx.arg.ltoo));
@@ -883,6 +884,16 @@ createUndefinedGlobal(StringRef name, llvm::wasm::WasmGlobalType *type) {
   return sym;
 }
 
+static UndefinedFunction *createUndefinedFunction(StringRef name,
+                                                  WasmSignature *signature) {
+  auto *sym = cast<UndefinedFunction>(symtab->addUndefinedFunction(
+      name, std::nullopt, std::nullopt, WASM_SYMBOL_UNDEFINED, nullptr,
+      signature, true));
+  ctx.arg.allowUndefinedSymbols.insert(sym->getName());
+  sym->isUsedInRegularObj = true;
+  return sym;
+}
+
 static InputGlobal *createGlobal(StringRef name, bool isMutable) {
   llvm::wasm::WasmGlobal wasmGlobal;
   bool is64 = ctx.arg.is64.value_or(false);
@@ -917,17 +928,26 @@ static void createSyntheticSymbols() {
                                                             true};
   static llvm::wasm::WasmGlobalType mutableGlobalTypeI64 = {WASM_TYPE_I64,
                                                             true};
+
   ctx.sym.callCtors = symtab->addSyntheticFunction(
       "__wasm_call_ctors", WASM_SYMBOL_VISIBILITY_HIDDEN,
       make<SyntheticFunction>(nullSignature, "__wasm_call_ctors"));
 
   bool is64 = ctx.arg.is64.value_or(false);
 
+  auto stack_pointer_name =
+      ctx.arg.libcallThreadContext ? "__init_stack_pointer" : "__stack_pointer";
   if (ctx.isPic) {
-    ctx.sym.stackPointer =
-        createUndefinedGlobal("__stack_pointer", ctx.arg.is64.value_or(false)
-                                                     ? &mutableGlobalTypeI64
-                                                     : &mutableGlobalTypeI32);
+    if (ctx.arg.libcallThreadContext) {
+      ctx.sym.stackPointer = createUndefinedGlobal(
+          stack_pointer_name,
+          ctx.arg.is64.value_or(false) ? &globalTypeI64 : &globalTypeI32);
+    } else {
+      ctx.sym.stackPointer = createUndefinedGlobal(stack_pointer_name,
+                                                   ctx.arg.is64.value_or(false)
+                                                       ? &mutableGlobalTypeI64
+                                                       : &mutableGlobalTypeI32);
+    }
     // For PIC code, we import two global variables (__memory_base and
     // __table_base) from the environment and use these as the offset at
     // which to load our static data and function table.
@@ -940,14 +960,18 @@ static void createSyntheticSymbols() {
     ctx.sym.tableBase->markLive();
   } else {
     // For non-PIC code
-    ctx.sym.stackPointer = createGlobalVariable("__stack_pointer", true);
+    ctx.sym.stackPointer =
+        createGlobalVariable(stack_pointer_name, !ctx.arg.libcallThreadContext);
     ctx.sym.stackPointer->markLive();
   }
 
   if (ctx.arg.sharedMemory) {
     // TLS symbols are all hidden/dso-local
+    auto tls_base_name =
+        ctx.arg.libcallThreadContext ? "__init_tls_base" : "__tls_base";
     ctx.sym.tlsBase =
-        createGlobalVariable("__tls_base", true, WASM_SYMBOL_VISIBILITY_HIDDEN);
+        createGlobalVariable(tls_base_name, !ctx.arg.libcallThreadContext,
+                             WASM_SYMBOL_VISIBILITY_HIDDEN);
     ctx.sym.tlsSize = createGlobalVariable("__tls_size", false,
                                            WASM_SYMBOL_VISIBILITY_HIDDEN);
     ctx.sym.tlsAlign = createGlobalVariable("__tls_align", false,
@@ -956,6 +980,17 @@ static void createSyntheticSymbols() {
         "__wasm_init_tls", WASM_SYMBOL_VISIBILITY_HIDDEN,
         make<SyntheticFunction>(is64 ? i64ArgSignature : i32ArgSignature,
                                 "__wasm_init_tls"));
+    if (ctx.arg.libcallThreadContext) {
+      ctx.sym.tlsBase->markLive();
+      ctx.sym.tlsSize->markLive();
+      ctx.sym.tlsAlign->markLive();
+      static WasmSignature setTLSBaseSignature{{}, {ValType::I32}};
+      ctx.sym.setTLSBase =
+          createUndefinedFunction("__wasm_set_tls_base", &setTLSBaseSignature);
+      static WasmSignature getTLSBaseSignature{{ValType::I32}, {}};
+      ctx.sym.getTLSBase =
+          createUndefinedFunction("__wasm_get_tls_base", &getTLSBaseSignature);
+    }
   }
 }
 

lld/wasm/Options.td

@@ -238,6 +238,9 @@ def page_size: JJ<"page-size=">,
 def initial_memory: JJ<"initial-memory=">,
   HelpText<"Initial size of the linear memory">;
 
+def libcall_thread_context: FF<"libcall-thread-context">,
+  HelpText<"Use library calls for thread context access instead of globals.">;
+  
 def max_memory: JJ<"max-memory=">,
   HelpText<"Maximum size of the linear memory">;