authorgravatar for luuk@degram.devLuuk de Gram <luuk@degram.dev> 2023-06-27 18:28:26+02:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-06-27 18:28:26+02:00
log622c5f3200f66e9db74d696b5dbec308f61d506e
treef237dd84263ee0d11e18f3c5329bf769525dd12a
parentff37ccd298f0ab28a9d0e0ee1110dadc6db4df1e
parent87b8a0567b0f54415aeecd879d3f1a4e12014d22
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #16207 from Luukdegram/wasi-threads

WASI: Implement experimental threading support

7 files changed, 373 insertions(+), 8 deletions(-)

lib/std/Thread.zig+290
......@@ -28,6 +28,8 @@ else if (use_pthreads)
2828 PosixThreadImpl
2929else if (target.os.tag == .linux)
3030 LinuxThreadImpl
31else if (target.os.tag == .wasi)
32 WasiThreadImpl
3133else
3234 UnsupportedImpl;
3335
......@@ -266,6 +268,7 @@ pub const Id = switch (target.os.tag) {
266268 .freebsd,
267269 .openbsd,
268270 .haiku,
271 .wasi,
269272 => u32,
270273 .macos, .ios, .watchos, .tvos => u64,
271274 .windows => os.windows.DWORD,
......@@ -296,6 +299,8 @@ pub const SpawnConfig = struct {
296299
297300 /// Size in bytes of the Thread's stack
298301 stack_size: usize = 16 * 1024 * 1024,
302 /// The allocator to be used to allocate memory for the to-be-spawned thread
303 allocator: ?std.mem.Allocator = null,
299304};
300305
301306pub const SpawnError = error{
......@@ -733,6 +738,291 @@ const PosixThreadImpl = struct {
733738 }
734739};
735740
741const WasiThreadImpl = struct {
742 thread: *WasiThread,
743
744 pub const ThreadHandle = i32;
745 threadlocal var tls_thread_id: Id = 0;
746
747 const WasiThread = struct {
748 /// Thread ID
749 tid: Atomic(i32) = Atomic(i32).init(0),
750 /// Contains all memory which was allocated to bootstrap this thread, including:
751 /// - Guard page
752 /// - Stack
753 /// - TLS segment
754 /// - `Instance`
755 /// All memory is freed upon call to `join`
756 memory: []u8,
757 /// The allocator used to allocate the thread's memory,
758 /// which is also used during `join` to ensure clean-up.
759 allocator: std.mem.Allocator,
760 /// The current state of the thread.
761 state: State = State.init(.running),
762 };
763
764 /// A meta-data structure used to bootstrap a thread
765 const Instance = struct {
766 thread: WasiThread,
767 /// Contains the offset to the new __tls_base.
768 /// The offset starting from the memory's base.
769 tls_offset: usize,
770 /// Contains the offset to the stack for the newly spawned thread.
771 /// The offset is calculated starting from the memory's base.
772 stack_offset: usize,
773 /// Contains the raw pointer value to the wrapper which holds all arguments
774 /// for the callback.
775 raw_ptr: usize,
776 /// Function pointer to a wrapping function which will call the user's
777 /// function upon thread spawn. The above mentioned pointer will be passed
778 /// to this function pointer as its argument.
779 call_back: *const fn (usize) void,
780 /// When a thread is in `detached` state, we must free all of its memory
781 /// upon thread completion. However, as this is done while still within
782 /// the thread, we must first jump back to the main thread's stack or else
783 /// we end up freeing the stack that we're currently using.
784 original_stack_pointer: [*]u8,
785 };
786
787 const State = Atomic(enum(u8) { running, completed, detached });
788
789 fn getCurrentId() Id {
790 return tls_thread_id;
791 }
792
793 fn getHandle(self: Impl) ThreadHandle {
794 return self.thread.tid.load(.SeqCst);
795 }
796
797 fn detach(self: Impl) void {
798 switch (self.thread.state.swap(.detached, .SeqCst)) {
799 .running => {},
800 .completed => self.join(),
801 .detached => unreachable,
802 }
803 }
804
805 fn join(self: Impl) void {
806 defer {
807 // Create a copy of the allocator so we do not free the reference to the
808 // original allocator while freeing the memory.
809 var allocator = self.thread.allocator;
810 allocator.free(self.thread.memory);
811 }
812
813 var spin: u8 = 10;
814 while (true) {
815 const tid = self.thread.tid.load(.SeqCst);
816 if (tid == 0) {
817 break;
818 }
819
820 if (spin > 0) {
821 spin -= 1;
822 std.atomic.spinLoopHint();
823 continue;
824 }
825
826 const result = asm (
827 \\ local.get %[ptr]
828 \\ local.get %[expected]
829 \\ i64.const -1 # infinite
830 \\ memory.atomic.wait32 0
831 \\ local.set %[ret]
832 : [ret] "=r" (-> u32),
833 : [ptr] "r" (&self.thread.tid.value),
834 [expected] "r" (tid),
835 );
836 switch (result) {
837 0 => continue, // ok
838 1 => continue, // expected =! loaded
839 2 => unreachable, // timeout (infinite)
840 else => unreachable,
841 }
842 }
843 }
844
845 fn spawn(config: std.Thread.SpawnConfig, comptime f: anytype, args: anytype) !WasiThreadImpl {
846 if (config.allocator == null) return error.OutOfMemory; // an allocator is required to spawn a WASI-thread
847
848 // Wrapping struct required to hold the user-provided function arguments.
849 const Wrapper = struct {
850 args: @TypeOf(args),
851 fn entry(ptr: usize) void {
852 const w: *@This() = @ptrFromInt(ptr);
853 @call(.auto, f, w.args);
854 }
855 };
856
857 var stack_offset: usize = undefined;
858 var tls_offset: usize = undefined;
859 var wrapper_offset: usize = undefined;
860 var instance_offset: usize = undefined;
861
862 // Calculate the bytes we have to allocate to store all thread information, including:
863 // - The actual stack for the thread
864 // - The TLS segment
865 // - `Instance` - containing information about how to call the user's function.
866 const map_bytes = blk: {
867 // start with atleast a single page, which is used as a guard to prevent
868 // other threads clobbering our new thread.
869 // Unfortunately, WebAssembly has no notion of read-only segments, so this
870 // is only a best effort.
871 var bytes: usize = std.wasm.page_size;
872
873 bytes = std.mem.alignForward(usize, bytes, 16); // align stack to 16 bytes
874 stack_offset = bytes;
875 bytes += @max(std.wasm.page_size, config.stack_size);
876
877 bytes = std.mem.alignForward(usize, bytes, __tls_align());
878 tls_offset = bytes;
879 bytes += __tls_size();
880
881 bytes = std.mem.alignForward(usize, bytes, @alignOf(Wrapper));
882 wrapper_offset = bytes;
883 bytes += @sizeOf(Wrapper);
884
885 bytes = std.mem.alignForward(usize, bytes, @alignOf(Instance));
886 instance_offset = bytes;
887 bytes += @sizeOf(Instance);
888
889 bytes = std.mem.alignForward(usize, bytes, std.wasm.page_size);
890 break :blk bytes;
891 };
892
893 // Allocate the amount of memory required for all meta data.
894 const allocated_memory = try config.allocator.?.alloc(u8, map_bytes);
895
896 const wrapper: *Wrapper = @ptrCast(@alignCast(&allocated_memory[wrapper_offset]));
897 wrapper.* = .{ .args = args };
898
899 const instance: *Instance = @ptrCast(@alignCast(&allocated_memory[instance_offset]));
900 instance.* = .{
901 .thread = .{ .memory = allocated_memory, .allocator = config.allocator.? },
902 .tls_offset = tls_offset,
903 .stack_offset = stack_offset,
904 .raw_ptr = @intFromPtr(wrapper),
905 .call_back = &Wrapper.entry,
906 .original_stack_pointer = __get_stack_pointer(),
907 };
908
909 const tid = spawnWasiThread(instance);
910 // The specification says any value lower than 0 indicates an error.
911 // The values of such error are unspecified. WASI-Libc treats it as EAGAIN.
912 if (tid < 0) {
913 return error.SystemResources;
914 }
915 instance.thread.tid.store(tid, .SeqCst);
916
917 return .{ .thread = &instance.thread };
918 }
919
920 /// Bootstrap procedure, called by the host environment after thread creation.
921 export fn wasi_thread_start(tid: i32, arg: *Instance) void {
922 if (builtin.single_threaded) {
923 // ensure function is not analyzed in single-threaded mode
924 return;
925 }
926 __set_stack_pointer(arg.thread.memory.ptr + arg.stack_offset);
927 __wasm_init_tls(arg.thread.memory.ptr + arg.tls_offset);
928 @atomicStore(u32, &WasiThreadImpl.tls_thread_id, @intCast(tid), .SeqCst);
929
930 // Finished bootstrapping, call user's procedure.
931 arg.call_back(arg.raw_ptr);
932
933 switch (arg.thread.state.swap(.completed, .SeqCst)) {
934 .running => {
935 // reset the Thread ID
936 asm volatile (
937 \\ local.get %[ptr]
938 \\ i32.const 0
939 \\ i32.atomic.store 0
940 :
941 : [ptr] "r" (&arg.thread.tid.value),
942 );
943
944 // Wake the main thread listening to this thread
945 asm volatile (
946 \\ local.get %[ptr]
947 \\ i32.const 1 # waiters
948 \\ memory.atomic.notify 0
949 \\ drop # no need to know the waiters
950 :
951 : [ptr] "r" (&arg.thread.tid.value),
952 );
953 },
954 .completed => unreachable,
955 .detached => {
956 // restore the original stack pointer so we can free the memory
957 // without having to worry about freeing the stack
958 __set_stack_pointer(arg.original_stack_pointer);
959 // Ensure a copy so we don't free the allocator reference itself
960 var allocator = arg.thread.allocator;
961 allocator.free(arg.thread.memory);
962 },
963 }
964 }
965
966 /// Asks the host to create a new thread for us.
967 /// Newly created thread will call `wasi_tread_start` with the thread ID as well
968 /// as the input `arg` that was provided to `spawnWasiThread`
969 const spawnWasiThread = @"thread-spawn";
970 extern "wasi" fn @"thread-spawn"(arg: *Instance) i32;
971
972 /// Initializes the TLS data segment starting at `memory`.
973 /// This is a synthetic function, generated by the linker.
974 extern fn __wasm_init_tls(memory: [*]u8) void;
975
976 /// Returns a pointer to the base of the TLS data segment for the current thread
977 inline fn __tls_base() [*]u8 {
978 return asm (
979 \\ .globaltype __tls_base, i32
980 \\ global.get __tls_base
981 \\ local.set %[ret]
982 : [ret] "=r" (-> [*]u8),
983 );
984 }
985
986 /// Returns the size of the TLS segment
987 inline fn __tls_size() u32 {
988 return asm volatile (
989 \\ .globaltype __tls_size, i32, immutable
990 \\ global.get __tls_size
991 \\ local.set %[ret]
992 : [ret] "=r" (-> u32),
993 );
994 }
995
996 /// Returns the alignment of the TLS segment
997 inline fn __tls_align() u32 {
998 return asm (
999 \\ .globaltype __tls_align, i32, immutable
1000 \\ global.get __tls_align
1001 \\ local.set %[ret]
1002 : [ret] "=r" (-> u32),
1003 );
1004 }
1005
1006 /// Allows for setting the stack pointer in the WebAssembly module.
1007 inline fn __set_stack_pointer(addr: [*]u8) void {
1008 asm volatile (
1009 \\ local.get %[ptr]
1010 \\ global.set __stack_pointer
1011 :
1012 : [ptr] "r" (addr),
1013 );
1014 }
1015
1016 /// Returns the current value of the stack pointer
1017 inline fn __get_stack_pointer() [*]u8 {
1018 return asm (
1019 \\ global.get __stack_pointer
1020 \\ local.set %[stack_ptr]
1021 : [stack_ptr] "=r" (-> [*]u8),
1022 );
1023 }
1024};
1025
7361026const LinuxThreadImpl = struct {
7371027 const linux = os.linux;
7381028
lib/std/Thread/Futex.zig+45
......@@ -73,6 +73,8 @@ else if (builtin.os.tag == .openbsd)
7373 OpenbsdImpl
7474else if (builtin.os.tag == .dragonfly)
7575 DragonflyImpl
76else if (builtin.target.isWasm())
77 WasmImpl
7678else if (std.Thread.use_pthreads)
7779 PosixImpl
7880else
......@@ -446,6 +448,49 @@ const DragonflyImpl = struct {
446448 }
447449};
448450
451const WasmImpl = struct {
452 fn wait(ptr: *const Atomic(u32), expect: u32, timeout: ?u64) error{Timeout}!void {
453 if (!comptime std.Target.wasm.featureSetHas(builtin.target.cpu.features, .atomics)) {
454 @compileError("WASI target missing cpu feature 'atomics'");
455 }
456 const to: i64 = if (timeout) |to| @intCast(to) else -1;
457 const result = asm (
458 \\local.get %[ptr]
459 \\local.get %[expected]
460 \\local.get %[timeout]
461 \\memory.atomic.wait32 0
462 \\local.set %[ret]
463 : [ret] "=r" (-> u32),
464 : [ptr] "r" (&ptr.value),
465 [expected] "r" (@as(i32, @bitCast(expect))),
466 [timeout] "r" (to),
467 );
468 switch (result) {
469 0 => {}, // ok
470 1 => {}, // expected =! loaded
471 2 => return error.Timeout,
472 else => unreachable,
473 }
474 }
475
476 fn wake(ptr: *const Atomic(u32), max_waiters: u32) void {
477 if (!comptime std.Target.wasm.featureSetHas(builtin.target.cpu.features, .atomics)) {
478 @compileError("WASI target missing cpu feature 'atomics'");
479 }
480 assert(max_waiters != 0);
481 const woken_count = asm (
482 \\local.get %[ptr]
483 \\local.get %[waiters]
484 \\memory.atomic.notify 0
485 \\local.set %[ret]
486 : [ret] "=r" (-> u32),
487 : [ptr] "r" (&ptr.value),
488 [waiters] "r" (max_waiters),
489 );
490 _ = woken_count; // can be 0 when linker flag 'shared-memory' is not enabled
491 }
492};
493
449494/// Modified version of linux's futex and Go's sema to implement userspace wait queues with pthread:
450495/// https://code.woboq.org/linux/linux/kernel/futex.c.html
451496/// https://go.dev/src/runtime/sema.go
src/Compilation.zig+3
......@@ -557,6 +557,7 @@ pub const InitOptions = struct {
557557 linker_allow_shlib_undefined: ?bool = null,
558558 linker_bind_global_refs_locally: ?bool = null,
559559 linker_import_memory: ?bool = null,
560 linker_export_memory: ?bool = null,
560561 linker_import_symbols: bool = false,
561562 linker_import_table: bool = false,
562563 linker_export_table: bool = false,
......@@ -1463,6 +1464,7 @@ pub fn create(gpa: Allocator, options: InitOptions) !*Compilation {
14631464 .module_definition_file = options.linker_module_definition_file,
14641465 .sort_section = options.linker_sort_section,
14651466 .import_memory = options.linker_import_memory orelse false,
1467 .export_memory = options.linker_export_memory orelse !(options.linker_import_memory orelse false),
14661468 .import_symbols = options.linker_import_symbols,
14671469 .import_table = options.linker_import_table,
14681470 .export_table = options.linker_export_table,
......@@ -2324,6 +2326,7 @@ fn addNonIncrementalStuffToCacheManifest(comp: *Compilation, man: *Cache.Manifes
23242326
23252327 // WASM specific stuff
23262328 man.hash.add(comp.bin_file.options.import_memory);
2329 man.hash.add(comp.bin_file.options.export_memory);
23272330 man.hash.addOptional(comp.bin_file.options.initial_memory);
23282331 man.hash.addOptional(comp.bin_file.options.max_memory);
23292332 man.hash.add(comp.bin_file.options.shared_memory);
src/link.zig+1
......@@ -133,6 +133,7 @@ pub const Options = struct {
133133 compress_debug_sections: CompressDebugSections,
134134 bind_global_refs_locally: bool,
135135 import_memory: bool,
136 export_memory: bool,
136137 import_symbols: bool,
137138 import_table: bool,
138139 export_table: bool,
src/link/Wasm.zig+5
......@@ -4251,6 +4251,7 @@ fn linkWithLLD(wasm: *Wasm, comp: *Compilation, prog_node: *std.Progress.Node) !
42514251 man.hash.addOptional(wasm.base.options.stack_size_override);
42524252 man.hash.add(wasm.base.options.build_id);
42534253 man.hash.add(wasm.base.options.import_memory);
4254 man.hash.add(wasm.base.options.export_memory);
42544255 man.hash.add(wasm.base.options.import_table);
42554256 man.hash.add(wasm.base.options.export_table);
42564257 man.hash.addOptional(wasm.base.options.initial_memory);
......@@ -4338,6 +4339,10 @@ fn linkWithLLD(wasm: *Wasm, comp: *Compilation, prog_node: *std.Progress.Node) !
43384339 try argv.append("--import-memory");
43394340 }
43404341
4342 if (wasm.base.options.export_memory) {
4343 try argv.append("--export-memory");
4344 }
4345
43414346 if (wasm.base.options.import_table) {
43424347 assert(!wasm.base.options.export_table);
43434348 try argv.append("--import-table");
src/main.zig+27-7
......@@ -544,6 +544,7 @@ const usage_build_generic =
544544 \\ -dead_strip (Darwin) remove functions and data that are unreachable by the entry point or exported symbols
545545 \\ -dead_strip_dylibs (Darwin) remove dylibs that are unreachable by the entry point or exported symbols
546546 \\ --import-memory (WebAssembly) import memory from the environment
547 \\ --export-memory (WebAssembly) export memory to the host (Default unless --import-memory used)
547548 \\ --import-symbols (WebAssembly) import missing symbols from the host environment
548549 \\ --import-table (WebAssembly) import function table from the host environment
549550 \\ --export-table (WebAssembly) export function table to the host environment
......@@ -787,6 +788,7 @@ fn buildOutputType(
787788 var linker_allow_shlib_undefined: ?bool = null;
788789 var linker_bind_global_refs_locally: ?bool = null;
789790 var linker_import_memory: ?bool = null;
791 var linker_export_memory: ?bool = null;
790792 var linker_import_symbols: bool = false;
791793 var linker_import_table: bool = false;
792794 var linker_export_table: bool = false;
......@@ -1419,6 +1421,8 @@ fn buildOutputType(
14191421 }
14201422 } else if (mem.eql(u8, arg, "--import-memory")) {
14211423 linker_import_memory = true;
1424 } else if (mem.eql(u8, arg, "--export-memory")) {
1425 linker_export_memory = true;
14221426 } else if (mem.eql(u8, arg, "--import-symbols")) {
14231427 linker_import_symbols = true;
14241428 } else if (mem.eql(u8, arg, "--import-table")) {
......@@ -1982,6 +1986,8 @@ fn buildOutputType(
19821986 linker_bind_global_refs_locally = true;
19831987 } else if (mem.eql(u8, arg, "--import-memory")) {
19841988 linker_import_memory = true;
1989 } else if (mem.eql(u8, arg, "--export-memory")) {
1990 linker_export_memory = true;
19851991 } else if (mem.eql(u8, arg, "--import-symbols")) {
19861992 linker_import_symbols = true;
19871993 } else if (mem.eql(u8, arg, "--import-table")) {
......@@ -2422,15 +2428,28 @@ fn buildOutputType(
24222428 link_libcpp = true;
24232429 }
24242430
2425 if (target_info.target.cpu.arch.isWasm() and linker_shared_memory) {
2426 if (output_mode == .Obj) {
2427 fatal("shared memory is not allowed in object files", .{});
2431 if (target_info.target.cpu.arch.isWasm()) blk: {
2432 if (single_threaded == null) {
2433 single_threaded = true;
24282434 }
2435 if (linker_shared_memory) {
2436 if (output_mode == .Obj) {
2437 fatal("shared memory is not allowed in object files", .{});
2438 }
24292439
2430 if (!target_info.target.cpu.features.isEnabled(@intFromEnum(std.Target.wasm.Feature.atomics)) or
2431 !target_info.target.cpu.features.isEnabled(@intFromEnum(std.Target.wasm.Feature.bulk_memory)))
2432 {
2433 fatal("'atomics' and 'bulk-memory' features must be enabled to use shared memory", .{});
2440 if (!target_info.target.cpu.features.isEnabled(@intFromEnum(std.Target.wasm.Feature.atomics)) or
2441 !target_info.target.cpu.features.isEnabled(@intFromEnum(std.Target.wasm.Feature.bulk_memory)))
2442 {
2443 fatal("'atomics' and 'bulk-memory' features must be enabled to use shared memory", .{});
2444 }
2445 break :blk;
2446 }
2447
2448 // Single-threaded is the default for WebAssembly, so only when the user specified `-fno_single-threaded`
2449 // can they enable multithreaded WebAssembly builds.
2450 const is_single_threaded = single_threaded.?;
2451 if (!is_single_threaded) {
2452 fatal("'-fno-single-threaded' requires the linker feature shared-memory to be enabled using '--shared-memory'", .{});
24342453 }
24352454 }
24362455
......@@ -3113,6 +3132,7 @@ fn buildOutputType(
31133132 .linker_allow_shlib_undefined = linker_allow_shlib_undefined,
31143133 .linker_bind_global_refs_locally = linker_bind_global_refs_locally,
31153134 .linker_import_memory = linker_import_memory,
3135 .linker_export_memory = linker_export_memory,
31163136 .linker_import_symbols = linker_import_symbols,
31173137 .linker_import_table = linker_import_table,
31183138 .linker_export_table = linker_export_table,
src/target.zig+2-1
......@@ -208,7 +208,8 @@ pub fn supports_fpic(target: std.Target) bool {
208208}
209209
210210pub fn isSingleThreaded(target: std.Target) bool {
211 return target.isWasm();
211 _ = target;
212 return false;
212213}
213214
214215/// Valgrind supports more, but Zig does not support them yet.