asyncify wip

Author: niaow

compileopts/config.go

@@ -150,7 +150,7 @@ func (c *Config) OptLevels() (optLevel, sizeLevel int, inlinerThreshold uint) {
 // target.
 func (c *Config) FuncImplementation() string {
 	switch c.Scheduler() {
-	case "tasks":
+	case "tasks", "asyncify":
 		// A func value is implemented as a pair of pointers:
 		//     {context, function pointer}
 		// where the context may be a pointer to a heap-allocated struct

compileopts/options.go

@@ -8,7 +8,7 @@ import (
 
 var (
 	validGCOptions            = []string{"none", "leaking", "extalloc", "conservative"}
-	validSchedulerOptions     = []string{"none", "tasks", "coroutines"}
+	validSchedulerOptions     = []string{"none", "tasks", "coroutines", "asyncify"}
 	validSerialOptions        = []string{"none", "uart", "usb"}
 	validPrintSizeOptions     = []string{"none", "short", "full"}
 	validPanicStrategyOptions = []string{"print", "trap"}

compiler/goroutine.go

@@ -92,7 +92,7 @@ func (b *builder) createGo(instr *ssa.Go) {
 		switch b.Scheduler {
 		case "none", "coroutines":
 			// There are no additional parameters needed for the goroutine start operation.
-		case "tasks":
+		case "tasks", "asyncify":
 			// Add the function pointer as a parameter to start the goroutine.
 			params = append(params, funcPtr)
 		default:
@@ -107,7 +107,7 @@ func (b *builder) createGo(instr *ssa.Go) {
 	paramBundle := b.emitPointerPack(params)
 	var callee, stackSize llvm.Value
 	switch b.Scheduler {
-	case "none", "tasks":
+	case "none", "tasks", "asyncify":
 		callee = b.createGoroutineStartWrapper(funcPtr, prefix, hasContext, instr.Pos())
 		if b.AutomaticStackSize {
 			// The stack size is not known until after linking. Call a dummy
@@ -119,7 +119,7 @@ func (b *builder) createGo(instr *ssa.Go) {
 		} else {
 			// The stack size is fixed at compile time. By emitting it here as a
 			// constant, it can be optimized.
-			if b.Scheduler == "tasks" && b.DefaultStackSize == 0 {
+			if (b.Scheduler == "tasks" || b.Scheduler == "asyncify") && b.DefaultStackSize == 0 {
 				b.addError(instr.Pos(), "default stack size for goroutines is not set")
 			}
 			stackSize = llvm.ConstInt(b.uintptrType, b.DefaultStackSize, false)

src/internal/task/task_asyncify.go

@@ -0,0 +1,101 @@
+// +build scheduler.asyncify
+
+package task
+
+import (
+	"unsafe"
+)
+
+// state is a structure which holds a reference to the state of the task.
+// When the task is suspended, the stack pointers are saved here.
+type state struct {
+	// entry is the entry function of the task.
+	// This is needed every time the function is invoked so that asyncify knows what to rewind.
+	entry uintptr
+
+	// args are a pointer to a struct holding the arguments of the function.
+	args unsafe.Pointer
+
+	// stackState is the state of the stack while unwound.
+	stackState
+}
+
+// stackState is the saved state of a stack while unwound.
+// The stack is arranged with asyncify at the bottom, C stack at the top, and a gap of available stack space between the two.
+type stackState struct {
+	// asyncify is the stack pointer of the asyncify stack.
+	// This starts from the bottom and grows upwards.
+	asyncifysp uintptr
+
+	// asyncify is stack pointer of the C stack.
+	// This starts from the top and grows downwards.
+	csp uintptr
+}
+
+// start creates and starts a new goroutine with the given function and arguments.
+// The new goroutine is immediately started.
+func start(fn uintptr, args unsafe.Pointer, stackSize uintptr) {
+	t := &Task{}
+	t.state.initialize(fn, args, stackSize)
+	prev := currentTask
+	currentTask = t
+	t.state.launch()
+	currentTask = prev
+}
+
+//export tinygo_launch
+func (*state) launch()
+
+// initialize the state and prepare to call the specified function with the specified argument bundle.
+func (s *state) initialize(fn uintptr, args unsafe.Pointer, stackSize uintptr) {
+	// Save the entry call.
+	s.entry = fn
+	s.args = args
+
+	// Create a stack.
+	stack := make([]uintptr, stackSize/unsafe.Sizeof(uintptr(0)))
+
+	// Calculate stack base addresses.
+	s.asyncifysp = uintptr(unsafe.Pointer(&stack[0]))
+	s.csp = uintptr(unsafe.Pointer(&stack[len(stack)-1])) + unsafe.Sizeof(uintptr(0)) - 1
+}
+
+//go:linkname runqueuePushBack runtime.runqueuePushBack
+func runqueuePushBack(*Task)
+
+// currentTask is the current running task, or nil if currently in the scheduler.
+var currentTask *Task
+
+// Current returns the current active task.
+func Current() *Task {
+	return currentTask
+}
+
+// Pause suspends the current task and returns to the scheduler.
+// This function may only be called when running on a goroutine stack, not when running on the system stack.
+func Pause() {
+	println("pausing", currentTask)
+	currentTask.state.unwind()
+}
+
+//export tinygo_unwind
+func (*stackState) unwind()
+
+// Resume the task until it pauses or completes.
+// This may only be called from the scheduler.
+func (t *Task) Resume() {
+	println("resume", t)
+	currentTask = t
+	t.state.rewind()
+	currentTask = nil
+	println("rewound", t)
+}
+
+//export tinygo_rewind
+func (*state) rewind()
+
+// OnSystemStack returns whether the caller is running on the system stack.
+func OnSystemStack() bool {
+	// If there is not an active goroutine, then this must be running on the system stack.
+	return Current() == nil
+}

src/internal/task/task_asyncify_wasm.S

@@ -0,0 +1,102 @@
+.globaltype __stack_pointer, i32
+
+.global  tinygo_unwind
+.type    tinygo_unwind,@function
+tinygo_unwind: // func (state *stackState) unwind()
+    .functype tinygo_unwind (i32) -> ()
+    .local i32
+    // Check if we are rewinding.
+    i32.const 0
+    i32.load8_u tinygo_rewinding
+    if // if tinygo_rewinding {
+    // Stop rewinding.
+    call stop_rewind
+    i32.const 0
+    i32.const 0
+    i32.store8 rewinding // tinygo_rewinding = false;
+    else
+    // Save the C stack pointer (destination structure pointer is in local 0).
+    local.get 0
+    global.get __stack_pointer
+    i32.store 4 // state.csp = getCurrentStackPointer()
+    // Ask asyncify to unwind.
+    // When resuming, asyncify will return this function with tinygo_rewinding set to true.
+    local.get 0
+    call start_unwind // asyncify.start_unwind(state)
+    end_if
+    return
+    end_function
+
+.global tinygo_launch
+.type tinygo_launch,@function
+tinygo_launch: // func (state *state) launch()
+    .functype tinygo_launch (i32) -> ()
+    .local i32
+    // Switch to the goroutine's C stack.
+    global.get __stack_pointer // prev := getCurrentStackPointer()
+    local.get 0
+    i32.load 12
+    global.set __stack_pointer // setStackPointer(state.csp)
+    // Get the argument pack and entry pointer.
+    local.get 0
+    i32.load 4 // args := state.args
+    local.get 0
+    i32.load 0 // fn := state.entry
+    // Launch the entry function.
+    call_indirect (i32) -> () // fn(args)
+    // Stop unwinding.
+    call stop_unwind
+    // Restore the C stack.
+    global.set __stack_pointer // setStackPointer(prev)
+    return
+    end_function
+
+.global  tinygo_rewind
+.type    tinygo_rewind,@function
+tinygo_rewind: // func (state *state) rewind()
+    .functype tinygo_rewind (i32) -> ()
+    .local i32
+    // Switch to the goroutine's C stack.
+    global.get __stack_pointer // prev := getCurrentStackPointer()
+    local.get 0
+    i32.load 12
+    global.set __stack_pointer // setStackPointer(state.csp)
+    // Get the argument pack and entry pointer.
+    local.get 0
+    i32.load 4 // args := state.args
+    local.get 0
+    i32.load 0 // fn := state.entry
+    // Prepare to rewind.
+    i32.const 0
+    i32.const 1
+    i32.store8 tinygo_rewinding // tinygo_rewinding = true;
+    local.get 0
+    i32.const 8
+    i32.add
+    call start_rewind // asyncify.start_rewind(&state.stackState)
+    // Launch the entry function.
+    // This will actually rewind the call stack.
+    call_indirect (i32) -> () // fn(args)
+    // Stop unwinding.
+    call stop_unwind
+    // Restore the C stack.
+    global.set __stack_pointer // setStackPointer(prev)
+    return
+    end_function
+
+.functype start_unwind (i32) -> ()
+.import_module start_unwind, asyncify
+.functype stop_unwind () -> ()
+.import_module stop_unwind, asyncify
+.functype start_rewind (i32) -> ()
+.import_module start_rewind, asyncify
+.functype stop_rewind () -> ()
+.import_module stop_rewind, asyncify
+
+        .hidden tinygo_rewinding                # @tinygo_rewinding
+        .type   tinygo_rewinding,@object
+        .section        .bss.tinygo_rewinding,"",@
+        .globl  tinygo_rewinding
+tinygo_rewinding:
+        .int8   0                               # 0x0
+        .size   tinygo_rewinding, 1

targets/wasi.json

@@ -5,6 +5,9 @@
 	"goarch":        "arm",
 	"linker":        "wasm-ld",
 	"libc":          "wasi-libc",
+	"scheduler":     "asyncify",
+	"gc":            "leaking",
+	"default-stack-size": 1000000,
 	"cflags": [
 		"--target=wasm32--wasi",
 		"--sysroot={root}/lib/wasi-libc/sysroot",
@@ -16,6 +19,9 @@
 		"--export-dynamic",
 		"--no-demangle"
 	],
+	"extra-files": [
+		"src/internal/task/task_asyncify_wasm.S"
+	],
 	"emulator":      ["wasmtime"],
 	"wasm-abi":      "generic"
 }

targets/wasm.json

@@ -5,6 +5,9 @@
 	"goarch":        "wasm",
 	"linker":        "wasm-ld",
 	"libc":          "wasi-libc",
+	"scheduler":     "asyncify",
+	"gc":            "leaking",
+	"default-stack-size": 2048,
 	"cflags": [
 		"--target=wasm32--wasi",
 		"--sysroot={root}/lib/wasi-libc/sysroot",
@@ -15,6 +18,9 @@
 		"--stack-first",
 		"--no-demangle"
 	],
+	"extra-files": [
+		"src/internal/task/task_asyncify_wasm.S"
+	],
 	"emulator":      ["node", "targets/wasm_exec.js"],
 	"wasm-abi":      "js"
 }

transform/optimizer.go

@@ -125,7 +125,7 @@ func Optimize(mod llvm.Module, config *compileopts.Config, optLevel, sizeLevel i
 		if err != nil {
 			return []error{err}
 		}
-	case "tasks":
+	case "tasks", "asyncify":
 		// No transformations necessary.
 	case "none":
 		// Check for any goroutine starts.
@@ -219,7 +219,7 @@ func getFunctionsUsedInTransforms(config *compileopts.Config) []string {
 	case "none":
 	case "coroutines":
 		fnused = append(append([]string{}, fnused...), coroFunctionsUsedInTransforms...)
-	case "tasks":
+	case "tasks", "asyncify":
 		fnused = append(append([]string{}, fnused...), taskFunctionsUsedInTransforms...)
 	default:
 		panic(fmt.Errorf("invalid scheduler %q", config.Scheduler()))