[WebAssembly] Cooperative threading for WASIP3

clang/lib/Basic/Targets/WebAssembly.cpp

@@ -424,7 +424,7 @@ void WebAssemblyTargetInfo::adjust(DiagnosticsEngine &Diags, LangOptions &Opts,
   // Turn off POSIXThreads and ThreadModel so that we don't predefine _REENTRANT
   // or __STDCPP_THREADS__ if we will eventually end up stripping atomics
   // because they are unsupported.
-  if (!HasAtomics || !HasBulkMemory) {
+  if ((!HasCooperativeThreading && !HasAtomics) || !HasBulkMemory) {
     Opts.POSIXThreads = false;
     Opts.setThreadModel(LangOptions::ThreadModelKind::Single);
     Opts.ThreadsafeStatics = false;

clang/lib/Basic/Targets/WebAssembly.h

@@ -63,6 +63,7 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyTargetInfo : public TargetInfo {
   bool HasBulkMemory = false;
   bool HasBulkMemoryOpt = false;
   bool HasCallIndirectOverlong = false;
+  bool HasCooperativeThreading = false;
   bool HasCompactImports = false;
   bool HasExceptionHandling = false;
   bool HasExtendedConst = false;
@@ -111,8 +112,10 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyTargetInfo : public TargetInfo {
       PtrDiffType = SignedLong;
       IntPtrType = SignedLong;
     }
-    if (T.getOS() == llvm::Triple::WASIp3)
+    if (T.getOS() == llvm::Triple::WASIp3) {
       HasLibcallThreadContext = true;
+      HasCooperativeThreading = true;
+    }
   }
 
   StringRef getABI() const override;

clang/lib/Driver/ToolChains/WebAssembly.cpp

@@ -85,14 +85,23 @@ static bool WantsPthread(const llvm::Triple &Triple, const ArgList &Args) {
   if (Triple.isOSWASI() && Triple.getEnvironmentName() == "threads")
     WantsPthread = true;
 
+  // WASIp3 also implies pthreads support
+  if (Triple.getOS() == llvm::Triple::WASIp3)
+    WantsPthread = true;
+
   return WantsPthread;
 }
 
-static bool WantsLibcallThreadContext(const llvm::Triple &Triple,
-                                      const ArgList &Args) {
+static bool WantsCooperativeMultithreading(const llvm::Triple &Triple,
+                                           const ArgList &Args) {
   return Triple.getOS() == llvm::Triple::WASIp3;
 }
 
+static bool WantsSharedMemory(const llvm::Triple &Triple, const ArgList &Args) {
+  return WantsPthread(Triple, Args) &&
+         !WantsCooperativeMultithreading(Triple, Args);
+}
+
 void wasm::Linker::ConstructJob(Compilation &C, const JobAction &JA,
                                 const InputInfo &Output,
                                 const InputInfoList &Inputs,
@@ -174,10 +183,10 @@ 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 (WantsCooperativeMultithreading(ToolChain.getTriple(), Args))
+    CmdArgs.push_back("--cooperative-threading");
 
-  if (WantsPthread(ToolChain.getTriple(), Args))
+  if (WantsSharedMemory(ToolChain.getTriple(), Args))
     CmdArgs.push_back("--shared-memory");
 
   if (!Args.hasArg(options::OPT_nostdlib, options::OPT_nodefaultlibs)) {
@@ -332,9 +341,12 @@ void WebAssembly::addClangTargetOptions(const ArgList &DriverArgs,
                           options::OPT_fno_use_init_array, true))
     CC1Args.push_back("-fno-use-init-array");
 
-  // '-pthread' implies atomics, bulk-memory, mutable-globals, and sign-ext
+  // '-pthread' implies bulk-memory, mutable-globals, and sign-ext.
+  // It also implies atomics, so long as we're not targeting a cooperative
+  // threading environment.
   if (WantsPthread(getTriple(), DriverArgs)) {
-    if (DriverArgs.hasFlag(options::OPT_mno_atomics, options::OPT_matomics,
+    if (!WantsCooperativeMultithreading(getTriple(), DriverArgs) &&
+        DriverArgs.hasFlag(options::OPT_mno_atomics, options::OPT_matomics,
                            false))
       getDriver().Diag(diag::err_drv_argument_not_allowed_with)
           << "-pthread"
@@ -354,8 +366,10 @@ void WebAssembly::addClangTargetOptions(const ArgList &DriverArgs,
       getDriver().Diag(diag::err_drv_argument_not_allowed_with)
           << "-pthread"
           << "-mno-sign-ext";
-    CC1Args.push_back("-target-feature");
-    CC1Args.push_back("+atomics");
+    if (!WantsCooperativeMultithreading(getTriple(), DriverArgs)) {
+      CC1Args.push_back("-target-feature");
+      CC1Args.push_back("+atomics");
+    }
     CC1Args.push_back("-target-feature");
     CC1Args.push_back("+bulk-memory");
     CC1Args.push_back("-target-feature");

clang/test/Driver/wasm-toolchain.c

@@ -303,3 +303,9 @@
 // RUN:   | FileCheck -check-prefix=LINK_WALI_BASIC %s
 // LINK_WALI_BASIC: "-cc1" {{.*}} "-o" "[[temp:[^"]*]]"
 // LINK_WALI_BASIC: wasm-ld{{.*}}" "-L/foo/lib/wasm32-linux-muslwali" "crt1.o" "[[temp]]" "-lc" "{{.*[/\\]}}libclang_rt.builtins.a" "-o" "a.out"
+
+// Test that `wasm32-wasip3` passes `--cooperative-threading` to the linker.
+
+// RUN: %clang -### --target=wasm32-wasip3 -fuse-ld=lld %s --sysroot /foo 2>&1 \
+// RUN:   | FileCheck -check-prefix=LINK_WASIP3_COOP %s
+// LINK_WASIP3_COOP: wasm-ld{{.*}}" {{.*}} "--cooperative-threading"

lld/test/wasm/cooperative-threading.s

@@ -0,0 +1,85 @@
+# Test that --cooperative-threading uses the libcall ABI naming for
+# thread-context globals (__init_stack_pointer, __init_tls_base, etc.) and
+# works without --shared-memory and atomics.
+
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.o %s
+# RUN: wasm-ld --cooperative-threading -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
+
+# Test that --cooperative-threading and --shared-memory are mutually exclusive.
+# RUN: not wasm-ld --cooperative-threading --shared-memory %t.o -o %t2.wasm 2>&1 | FileCheck %s --check-prefix=INCOMPAT
+# INCOMPAT: --cooperative-threading is incompatible with --shared-memory
+
+.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"
+
+# Memory must NOT be marked as shared.
+# CHECK:      - Type:            MEMORY
+# CHECK-NEXT:   Memories:
+# CHECK-NEXT:     - Minimum:         0x2
+# CHECK-NOT:       Shared
+
+# Globals should use the libcall ABI naming, not the global ABI.
+# 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    {{[0-9]+}}
+# DIS-NEXT:       i32.const       0
+# DIS-NEXT:       i32.add
+# DIS-NEXT:       i32.load        0
+# DIS-NEXT:       call    {{[0-9]+}}
+# DIS-NEXT:       i32.const       4
+# DIS-NEXT:       i32.add
+# DIS-NEXT:       i32.load        0
+# DIS-NEXT:       i32.add
+# DIS-NEXT:       end

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

@@ -1,5 +1,5 @@
 # RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.o %s
-# RUN: wasm-ld --libcall-thread-context --no-gc-sections -o %t.libcall.wasm %t.o
+# RUN: wasm-ld --cooperative-threading --no-gc-sections -o %t.libcall.wasm %t.o
 # RUN: obj2yaml %t.libcall.wasm | FileCheck %s --check-prefix=LIBCALL
 # RUN: wasm-ld --no-gc-sections -o %t.global.wasm %t.o
 # RUN: obj2yaml %t.global.wasm | FileCheck %s --check-prefix=GLOBAL

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

@@ -3,10 +3,9 @@
 # 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
+# RUN: not wasm-ld --cooperative-threading %t.o -o %t.wasm 2>&1 | FileCheck %s
 
+# CHECK: object file uses globals for thread context, but --cooperative-threading was specified
 .globl _start
 _start:
   .functype _start () -> ()

lld/test/wasm/tls-libcall.s

@@ -1,5 +1,5 @@
 # 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: wasm-ld --cooperative-threading -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
 

lld/wasm/Config.h

@@ -65,6 +65,7 @@ struct Config {
   bool growableTable;
   bool gcSections;
   llvm::StringSet<> keepSections;
+  bool cooperativeThreading;
   bool libcallThreadContext;
   std::optional<std::pair<llvm::StringRef, llvm::StringRef>> memoryImport;
   std::optional<llvm::StringRef> memoryExport;
@@ -134,6 +135,8 @@ struct Config {
   std::optional<std::vector<std::string>> features;
   std::optional<std::vector<std::string>> extraFeatures;
   llvm::SmallVector<uint8_t, 0> buildIdVector;
+
+  bool isMultithreaded() const { return sharedMemory || cooperativeThreading; }
 };
 
 // The Ctx object hold all other (non-configuration) global state.

lld/wasm/Driver.cpp

@@ -561,7 +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.cooperativeThreading = args.hasArg(OPT_cooperative_threading);
   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));
@@ -755,6 +755,11 @@ static void setConfigs() {
   if (!ctx.arg.memoryExport.has_value() && !ctx.arg.memoryImport.has_value()) {
     ctx.arg.memoryExport = memoryName;
   }
+  if (ctx.arg.cooperativeThreading) {
+    if (ctx.arg.sharedMemory)
+      error("--cooperative-threading is incompatible with --shared-memory");
+    ctx.arg.libcallThreadContext = true;
+  }
 }
 
 // Some command line options or some combinations of them are not allowed.
@@ -964,7 +969,7 @@ static void createSyntheticSymbols() {
         createGlobalVariable(stack_pointer_name, !ctx.arg.libcallThreadContext);
   }
 
-  if (ctx.arg.sharedMemory) {
+  if (ctx.arg.isMultithreaded()) {
     // TLS symbols are all hidden/dso-local
     auto tls_base_name =
         ctx.arg.libcallThreadContext ? "__init_tls_base" : "__tls_base";
@@ -986,9 +991,11 @@ static void createSyntheticSymbols() {
       static WasmSignature setTLSBaseSignature{{}, {ValType::I32}};
       ctx.sym.setTLSBase =
           createUndefinedFunction("__wasm_set_tls_base", &setTLSBaseSignature);
+      ctx.sym.setTLSBase->markLive();
       static WasmSignature getTLSBaseSignature{{ValType::I32}, {}};
       ctx.sym.getTLSBase =
           createUndefinedFunction("__wasm_get_tls_base", &getTLSBaseSignature);
+      ctx.sym.getTLSBase->markLive();
     }
   }
 }
@@ -1019,16 +1026,12 @@ static void createOptionalSymbols() {
   if (ctx.sym.firstPageEnd)
     ctx.sym.firstPageEnd->setVA(ctx.arg.pageSize);
 
-  // For non-shared memory programs we still need to define __tls_base since we
-  // allow object files built with TLS to be linked into single threaded
-  // programs, and such object files can contain references to this symbol.
-  //
-  // However, in this case __tls_base is immutable and points directly to the
-  // start of the `.tdata` static segment.
-  //
-  // __tls_size and __tls_align are not needed in this case since they are only
-  // needed for __wasm_init_tls (which we do not create in this case).
-  if (!ctx.arg.sharedMemory)
+  // TLS object files may be linked into single-threaded programs, so
+  // __tls_base must always be defined. In this case it is immutable and points
+  // directly to the start of the `.tdata` segment. __tls_size and __tls_align
+  // are omitted since they are only used by __wasm_init_tls, which is not
+  // created in this case.
+  if (!ctx.sym.tlsBase)
     ctx.sym.tlsBase = createOptionalGlobal("__tls_base", false);
 }
 

lld/wasm/Options.td

@@ -238,8 +238,8 @@ 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 cooperative_threading: FF<"cooperative-threading">,
+  HelpText<"Enable cooperative multithreading.">;
 
 def max_memory: JJ<"max-memory=">,
   HelpText<"Maximum size of the linear memory">;

lld/wasm/Relocations.cpp

@@ -125,7 +125,7 @@ void scanRelocations(InputChunk *chunk) {
       // In single-threaded builds TLS is lowered away and TLS data can be
       // merged with normal data and allowing TLS relocation in non-TLS
       // segments.
-      if (ctx.arg.sharedMemory) {
+      if (ctx.arg.isMultithreaded()) {
         if (!sym->isTLS()) {
           error(toString(file) + ": relocation " +
                 relocTypeToString(reloc.Type) +

lld/wasm/SyntheticSections.cpp

@@ -57,7 +57,7 @@ void writeGetTLSBase(const Ctx &ctx, raw_ostream &os) {
     writeU8(os, WASM_OPCODE_CALL, "call");
     writeUleb128(os, ctx.sym.getTLSBase->getFunctionIndex(), "function index");
   } else {
-    writeU8(os, WASM_OPCODE_GLOBAL_GET, "GLOBAL_SET");
+    writeU8(os, WASM_OPCODE_GLOBAL_GET, "GLOBAL_GET");
     writeUleb128(os, ctx.sym.tlsBase->getGlobalIndex(), "__tls_base");
   }
 }
@@ -532,7 +532,7 @@ void GlobalSection::writeBody() {
         mutable_ = true;
       // With multi-threading any TLS globals must be mutable since they get
       // set during `__wasm_apply_global_tls_relocs`
-      if (ctx.arg.sharedMemory && sym->isTLS())
+      if (ctx.arg.isMultithreaded() && sym->isTLS())
         mutable_ = true;
     }
     WasmGlobalType type{itype, mutable_};
@@ -569,10 +569,11 @@ void GlobalSection::writeBody() {
     } else {
       WasmInitExpr initExpr;
       if (auto *d = dyn_cast<DefinedData>(sym))
-        // In the sharedMemory case TLS globals are set during
-        // `__wasm_apply_global_tls_relocs`, but in the non-shared case
+        // In the multithreaded case, TLS globals are set during
+        // `__wasm_apply_global_tls_relocs`, but in the single-threaded case
         // we know the absolute value at link time.
-        initExpr = intConst(d->getVA(/*absolute=*/!ctx.arg.sharedMemory), is64);
+        initExpr =
+            intConst(d->getVA(/*absolute=*/!ctx.arg.isMultithreaded()), is64);
       else if (auto *f = dyn_cast<FunctionSymbol>(sym))
         initExpr = intConst(f->isStub ? 0 : f->getTableIndex(), is64);
       else {
@@ -680,7 +681,7 @@ bool DataCountSection::isNeeded() const {
   // instructions are not yet supported in input files.  However, in the case
   // of shared memory, lld itself will generate these instructions as part of
   // `__wasm_init_memory`. See Writer::createInitMemoryFunction.
-  return numSegments && ctx.arg.sharedMemory;
+  return numSegments && ctx.arg.isMultithreaded();
 }
 
 void LinkingSection::writeBody() {