From 907d90105a5d7ff3f74fb6e147ced03781b8e8ac Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem@gmail.com>
Date: Tue, 14 Jun 2022 23:54:19 +0200
Subject: [PATCH] runtime: use the main (startup) stack for the main goroutine

Instead of always starting a new goroutine for the main goroutine, run
the main goroutine on the system stack.
The system stack is not occupied with scheduling, instead each goroutine
that wants to pause itself calls into the scheduler which will switch to
the next task (goroutine) to run, or sleeps.

There are various advantages of this over the previous system:

  * When the program doesn't start a goroutine, the code size and RAM
    consumption is close to what you'd get with `-scheduler=none`.
  * When the program does start a goroutine, there is still a reduction
    in RAM consumption because only one extra stack is needed.
  * Because tasks directly switch to the next task to run, only a single
    task switch is needed instead of two (goroutine -> scheduler ->
    goroutine). This should improve task switching performance.

I kept the current behavior for WebAssembly/Asyncify. I looked into how
the same benefits can be realized for WebAssembly but couldn't easily
find how to do that. Maybe this can be done separately, or maybe we'll
just wait for the stack switching proposal to finish.

The code for Cortex-M is currently more complicated than I'd like, and
therefore can sometimes result in a slight increase in code size. I'd
like to fix this eventually but am still looking into good ways to do
this. I still think this change is generally beneficial because many
programs see big reductions in code size when compiling for Cortex-M.
---
 builder/build.go                              | 16 +++-
 builder/elfpatch.go                           |  2 +-
 src/internal/task/task_asyncify.go            | 10 +-
 src/internal/task/task_none.go                |  6 +-
 src/internal/task/task_stack.go               | 35 ++++---
 src/internal/task/task_stack_386.S            |  7 +-
 src/internal/task/task_stack_386.go           | 14 +--
 src/internal/task/task_stack_amd64.S          |  6 +-
 src/internal/task/task_stack_amd64.go         | 14 +--
 src/internal/task/task_stack_amd64_windows.S  |  4 +-
 src/internal/task/task_stack_amd64_windows.go | 14 +--
 src/internal/task/task_stack_arm.go           | 14 +--
 src/internal/task/task_stack_arm64.go         | 14 +--
 src/internal/task/task_stack_avr.S            |  6 --
 src/internal/task/task_stack_avr.go           | 15 +--
 src/internal/task/task_stack_baremetal.go     | 12 +++
 src/internal/task/task_stack_cortexm.S        | 93 +++++++++++++++----
 src/internal/task/task_stack_cortexm.go       | 20 ++--
 src/internal/task/task_stack_esp32.go         | 14 +--
 src/internal/task/task_stack_esp8266.go       | 14 +--
 src/internal/task/task_stack_tinygoriscv.go   | 14 +--
 src/runtime/gc_stack_raw.go                   |  4 +-
 src/runtime/scheduler.go                      | 10 +-
 ...scheduler_any.go => scheduler_asyncify.go} |  4 +-
 src/runtime/scheduler_tasks.go                | 24 +++++
 targets/arm.ld                                |  1 +
 targets/avr.ld                                |  1 +
 targets/esp32.ld                              |  1 +
 targets/esp32c3.ld                            |  1 +
 targets/esp8266.ld                            |  5 +
 targets/maixbit.ld                            |  1 +
 targets/mimxrt1062-teensy40.ld                |  1 +
 targets/riscv.ld                              |  1 +
 33 files changed, 221 insertions(+), 177 deletions(-)
 create mode 100644 src/internal/task/task_stack_baremetal.go
 rename src/runtime/{scheduler_any.go => scheduler_asyncify.go} (90%)

diff --git a/builder/build.go b/builder/build.go
index ab5b2b44e..0ece7ff5e 100644
--- a/builder/build.go
+++ b/builder/build.go
@@ -1228,10 +1228,15 @@ func determineStackSizes(mod llvm.Module, executable string) ([]string, map[stri
 	// Goroutines need to be started and finished and take up some stack space
 	// that way. This can be measured by measuing the stack size of
 	// tinygo_startTask.
-	if numFuncs := len(functions["tinygo_startTask"]); numFuncs != 1 {
-		return nil, nil, fmt.Errorf("expected exactly one definition of tinygo_startTask, got %d", numFuncs)
+	var baseStackSize uint64
+	var baseStackSizeType stacksize.SizeType
+	var baseStackSizeFailedAt *stacksize.CallNode
+	if len(gowrappers) != 0 {
+		if numFuncs := len(functions["tinygo_startTask"]); numFuncs != 1 {
+			return nil, nil, fmt.Errorf("expected exactly one definition of tinygo_startTask, got %d", numFuncs)
+		}
+		baseStackSize, baseStackSizeType, baseStackSizeFailedAt = functions["tinygo_startTask"][0].StackSize()
 	}
-	baseStackSize, baseStackSizeType, baseStackSizeFailedAt := functions["tinygo_startTask"][0].StackSize()
 
 	sizes := make(map[string]functionStackSize)
 
@@ -1303,6 +1308,11 @@ func modifyStackSizes(executable string, stackSizeLoads []string, stackSizes map
 	if err != nil {
 		return err
 	}
+	if data == nil {
+		// The .tinygo_stacksizes section doesn't exist, so assume this
+		// modification isn't needed.
+		return nil
+	}
 
 	if len(stackSizeLoads)*4 != len(data) {
 		// Note: while AVR should use 2 byte stack sizes, even 64-bit platforms
diff --git a/builder/elfpatch.go b/builder/elfpatch.go
index 6a407db6f..fc13c97ea 100644
--- a/builder/elfpatch.go
+++ b/builder/elfpatch.go
@@ -15,7 +15,7 @@ func getElfSectionData(executable string, sectionName string) ([]byte, elf.FileH
 
 	section := elfFile.Section(sectionName)
 	if section == nil {
-		return nil, elf.FileHeader{}, fmt.Errorf("could not find %s section", sectionName)
+		return nil, elf.FileHeader{}, nil
 	}
 
 	data, err := section.Data()
diff --git a/src/internal/task/task_asyncify.go b/src/internal/task/task_asyncify.go
index 939008bcd..1cee9fc92 100644
--- a/src/internal/task/task_asyncify.go
+++ b/src/internal/task/task_asyncify.go
@@ -99,9 +99,9 @@ func Pause() {
 //export tinygo_unwind
 func (*stackState) unwind()
 
-// Resume the task until it pauses or completes.
+// Switch to this task until it pauses or completes.
 // This may only be called from the scheduler.
-func (t *Task) Resume() {
+func (t *Task) Switch() {
 	// The current task must be saved and restored because this can nest on WASM with JS.
 	prevTask := currentTask
 	t.gcData.swap()
@@ -121,9 +121,3 @@ func (t *Task) Resume() {
 
 //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
-}
diff --git a/src/internal/task/task_none.go b/src/internal/task/task_none.go
index 19dbee624..ccb0c3dfe 100644
--- a/src/internal/task/task_none.go
+++ b/src/internal/task/task_none.go
@@ -28,12 +28,12 @@ func start(fn uintptr, args unsafe.Pointer, stackSize uintptr) {
 
 type state struct{}
 
-func (t *Task) Resume() {
+func (t *Task) Switch() {
 	runtimePanic("scheduler is disabled")
 }
 
-// OnSystemStack returns whether the caller is running on the system stack.
-func OnSystemStack() bool {
+// MainTask returns whether the caller is running in the main goroutine.
+func MainTask() bool {
 	// This scheduler does not do any stack switching.
 	return true
 }
diff --git a/src/internal/task/task_stack.go b/src/internal/task/task_stack.go
index 59af65031..0b48c6838 100644
--- a/src/internal/task/task_stack.go
+++ b/src/internal/task/task_stack.go
@@ -27,8 +27,12 @@ type state struct {
 	canaryPtr *uintptr
 }
 
-// currentTask is the current running task, or nil if currently in the scheduler.
-var currentTask *Task
+// The task struct for the main goroutine.
+var mainTask Task
+
+// currentTask is the current running task. The default value is the main
+// goroutine.
+var currentTask = &mainTask
 
 // Current returns the current active task.
 func Current() *Task {
@@ -36,29 +40,36 @@ func Current() *Task {
 }
 
 // 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 or in an interrupt.
+// This function may only be called when running on a goroutine stack, not when in an interrupt.
 func Pause() {
 	// Check whether the canary (the lowest address of the stack) is still
 	// valid. If it is not, a stack overflow has occured.
-	if *currentTask.state.canaryPtr != stackCanary {
+	if currentTask.state.canaryPtr != nil && *currentTask.state.canaryPtr != stackCanary {
 		runtimePanic("goroutine stack overflow")
 	}
-	currentTask.state.pause()
+	scheduler()
 }
 
+//go:linkname scheduler runtime.scheduler
+func scheduler()
+
 //export tinygo_pause
 func pause() {
 	Pause()
 }
 
-// Resume the task until it pauses or completes.
+// Switch to the given task until it pauses or completes.
 // This may only be called from the scheduler.
-func (t *Task) Resume() {
+func (t *Task) Switch() {
+	current := currentTask
+	if current == t {
+		// Nothing to switch to: we're already in this task.
+		return
+	}
 	currentTask = t
 	t.gcData.swap()
-	t.state.resume()
+	t.state.switchTo(current)
 	t.gcData.swap()
-	currentTask = nil
 }
 
 // initialize the state and prepare to call the specified function with the specified argument bundle.
@@ -103,8 +114,8 @@ func start(fn uintptr, args unsafe.Pointer, stackSize uintptr) {
 	runqueuePushBack(t)
 }
 
-// OnSystemStack returns whether the caller is running on the system stack.
-func OnSystemStack() bool {
+// MainTask returns whether the caller is running in the main goroutine.
+func MainTask() bool {
 	// If there is not an active goroutine, then this must be running on the system stack.
-	return Current() == nil
+	return currentTask == &mainTask
 }
diff --git a/src/internal/task/task_stack_386.S b/src/internal/task/task_stack_386.S
index c82213e98..eb4fb95cf 100644
--- a/src/internal/task/task_stack_386.S
+++ b/src/internal/task/task_stack_386.S
@@ -20,11 +20,8 @@ tinygo_startTask:
     // Branch to the "goroutine start" function.
     calll *%ebx
 
-    // Rebalance the stack (to undo the above push).
-    addl $4, %esp
-
-    // After return, exit this goroutine. This is a tail call.
-    jmp tinygo_pause
+    // After return, exit this goroutine.
+    calll tinygo_pause
     .cfi_endproc
 
 .global tinygo_swapTask
diff --git a/src/internal/task/task_stack_386.go b/src/internal/task/task_stack_386.go
index 10744757a..8b698a662 100644
--- a/src/internal/task/task_stack_386.go
+++ b/src/internal/task/task_stack_386.go
@@ -5,8 +5,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_386.S that relies on the exact
 // layout of this struct.
@@ -46,18 +44,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.esi = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_amd64.S b/src/internal/task/task_stack_amd64.S
index f9182d49f..36287fc24 100644
--- a/src/internal/task/task_stack_amd64.S
+++ b/src/internal/task/task_stack_amd64.S
@@ -28,11 +28,11 @@ tinygo_startTask:
     // Branch to the "goroutine start" function.
     callq *%r12
 
-    // After return, exit this goroutine. This is a tail call.
+    // After return, exit this goroutine.
     #ifdef __MACH__
-    jmp _tinygo_pause
+    callq _tinygo_pause
     #else
-    jmp tinygo_pause
+    callq tinygo_pause
     #endif
     .cfi_endproc
 
diff --git a/src/internal/task/task_stack_amd64.go b/src/internal/task/task_stack_amd64.go
index b39189316..8e9e8d4c2 100644
--- a/src/internal/task/task_stack_amd64.go
+++ b/src/internal/task/task_stack_amd64.go
@@ -5,8 +5,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_amd64.S that relies on the exact
 // layout of this struct.
@@ -45,18 +43,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.r13 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_amd64_windows.S b/src/internal/task/task_stack_amd64_windows.S
index bdbcfca66..1545f275b 100644
--- a/src/internal/task/task_stack_amd64_windows.S
+++ b/src/internal/task/task_stack_amd64_windows.S
@@ -17,8 +17,8 @@ tinygo_startTask:
     // Branch to the "goroutine start" function.
     callq *%r12
 
-    // After return, exit this goroutine. This is a tail call.
-    jmp tinygo_pause
+    // After return, exit this goroutine.
+    callq tinygo_pause
 
 .global tinygo_swapTask
 .section .text.tinygo_swapTask,"ax"
diff --git a/src/internal/task/task_stack_amd64_windows.go b/src/internal/task/task_stack_amd64_windows.go
index 41ddbe6c9..890f59544 100644
--- a/src/internal/task/task_stack_amd64_windows.go
+++ b/src/internal/task/task_stack_amd64_windows.go
@@ -8,8 +8,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_amd64_windows.S that relies on
 // the exact layout of this struct.
@@ -50,18 +48,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.r13 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_arm.go b/src/internal/task/task_stack_arm.go
index 18982e875..a6da5f12f 100644
--- a/src/internal/task/task_stack_arm.go
+++ b/src/internal/task/task_stack_arm.go
@@ -5,8 +5,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_arm.S that relies on the exact
 // layout of this struct.
@@ -45,18 +43,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.r5 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_arm64.go b/src/internal/task/task_stack_arm64.go
index e4fd09b20..0d1d5e277 100644
--- a/src/internal/task/task_stack_arm64.go
+++ b/src/internal/task/task_stack_arm64.go
@@ -5,8 +5,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_arm64.S that relies on the exact
 // layout of this struct.
@@ -48,18 +46,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.x20 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_avr.S b/src/internal/task/task_stack_avr.S
index 2cba528f2..d4c5629ad 100644
--- a/src/internal/task/task_stack_avr.S
+++ b/src/internal/task/task_stack_avr.S
@@ -1,9 +1,3 @@
-.section .bss.tinygo_systemStack
-.global tinygo_systemStack
-.type   tinygo_systemStack, %object
-tinygo_systemStack:
-    .short 0
-
 .section .text.tinygo_startTask
 .global  tinygo_startTask
 .type    tinygo_startTask, %function
diff --git a/src/internal/task/task_stack_avr.go b/src/internal/task/task_stack_avr.go
index 96b7034d9..aa082ecc3 100644
--- a/src/internal/task/task_stack_avr.go
+++ b/src/internal/task/task_stack_avr.go
@@ -5,9 +5,6 @@ package task
 
 import "unsafe"
 
-//go:extern tinygo_systemStack
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_avr.S that relies on the exact
 // layout of this struct.
@@ -51,18 +48,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.r4r5 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_baremetal.go b/src/internal/task/task_stack_baremetal.go
new file mode 100644
index 000000000..cd868f654
--- /dev/null
+++ b/src/internal/task/task_stack_baremetal.go
@@ -0,0 +1,12 @@
+//go:build scheduler.tasks && baremetal
+// +build scheduler.tasks,baremetal
+
+package task
+
+//go:extern _stack_bottom
+var _stack_bottom uintptr
+
+func init() {
+	mainTask.state.canaryPtr = &_stack_bottom
+	*mainTask.state.canaryPtr = stackCanary
+}
diff --git a/src/internal/task/task_stack_cortexm.S b/src/internal/task/task_stack_cortexm.S
index afb23eb27..8923d372b 100644
--- a/src/internal/task/task_stack_cortexm.S
+++ b/src/internal/task/task_stack_cortexm.S
@@ -30,10 +30,10 @@ tinygo_startTask:
     .cfi_endproc
 .size tinygo_startTask, .-tinygo_startTask
 
-.section .text.tinygo_switchToScheduler
-.global  tinygo_switchToScheduler
-.type    tinygo_switchToScheduler, %function
-tinygo_switchToScheduler:
+.section .text.tinygo_switchToMain
+.global  tinygo_switchToMain
+.type    tinygo_switchToMain, %function
+tinygo_switchToMain:
     .cfi_startproc
     // r0 = sp *uintptr
 
@@ -43,9 +43,9 @@ tinygo_switchToScheduler:
     subs r1, #36
     str r1, [r0]
 
-    b tinygo_swapTask
+    b tinygo_swapTaskInternal
     .cfi_endproc
-.size tinygo_switchToScheduler, .-tinygo_switchToScheduler
+.size tinygo_switchToMain, .-tinygo_switchToMain
 
 .section .text.tinygo_switchToTask
 .global  tinygo_switchToTask
@@ -55,17 +55,17 @@ tinygo_switchToTask:
     // r0 = sp uintptr
 
     // Currently on the scheduler stack (SP=MSP). We'll have to update the PSP,
-    // and then we can invoke swapTask.
+    // and then we can invoke swapTaskInternal.
     msr PSP, r0
 
-    b.n tinygo_swapTask
+    b.n tinygo_swapTaskInternal
     .cfi_endproc
 .size tinygo_switchToTask, .-tinygo_switchToTask
 
-.section .text.tinygo_swapTask
-.global  tinygo_swapTask
-.type    tinygo_swapTask, %function
-tinygo_swapTask:
+.section .text.tinygo_swapTaskInternal
+.global  tinygo_swapTaskInternal
+.type    tinygo_swapTaskInternal, %function
+tinygo_swapTaskInternal:
     .cfi_startproc
     // This function stores the current register state to the stack, switches to
     // the other stack (MSP/PSP), and loads the register state from the other
@@ -77,15 +77,15 @@ tinygo_swapTask:
     // used directly. Only very few operations work on them, such as mov. That's
     // why the higher register values are first stored in the temporary register
     // r3 when loading/storing them.
-    // It is possible to reduce the swapTask by two instructions (~2 cycles) on
-    // Cortex-M0 by reordering the layout of the pushed registers from {r4-r11,
-    // lr} to {r8-r11, r4-r8, lr}. However, that also requires a change on the
-    // Go side (depending on thumb1/thumb2!) and so is not really worth the
-    // complexity.
+    // It is possible to reduce the swapTaskInternal by two instructions (~2
+    // cycles) on Cortex-M0 by reordering the layout of the pushed registers
+    // from {r4-r11, lr} to {r8-r11, r4-r8, lr}. However, that also requires a
+    // change on the Go side (depending on thumb1/thumb2!) and so is not really
+    // worth the complexity.
 
     // Store state to old task. It saves the lr instead of the pc, because that
     // will be the pc after returning back to the old task (in a different
-    // invocation of swapTask).
+    // invocation of swapTaskInternal).
     #if defined(__thumb2__)
     push {r4-r11, lr}
     .cfi_def_cfa_offset 9*4
@@ -108,6 +108,63 @@ tinygo_swapTask:
     msr  CONTROL, r0 // store CONTROL register
     isb              // required to flush the pipeline
 
+    // Load state from new task and branch to the previous position in the
+    // program.
+    #if defined(__thumb2__)
+    pop {r4-r11, pc}
+    #else
+    pop {r4-r7}
+    .cfi_def_cfa_offset 5*9
+    pop {r0-r3}
+    .cfi_def_cfa_offset 1*9
+    mov r8, r0
+    mov r9, r1
+    mov r10, r2
+    mov r11, r3
+    pop {pc}
+    #endif
+    .cfi_endproc
+.size tinygo_swapTaskInternal, .-tinygo_swapTaskInternal
+
+
+.section .text.tinygo_swapTask
+.global  tinygo_swapTask
+.type    tinygo_swapTask, %function
+tinygo_swapTask:
+    .cfi_startproc
+    // This function is like tinygo_swapTaskInternal but only switches the stack
+    // pointer without touching the SPSEL bit (to switch between MSP and PSP).
+    // Therefore, it is used to switch between two tasks that are not the main
+    // goroutine.
+
+    // Store state to old task. It saves the lr instead of the pc, because that
+    // will be the pc after returning back to the old task (in a different
+    // invocation of swapTask).
+    #if defined(__thumb2__)
+    push {r4-r11, lr}
+    .cfi_def_cfa_offset 9*4
+    #else
+    mov r0, r8
+    mov r1, r9
+    mov r2, r10
+    mov r3, r11
+    push {r0-r3, lr}
+    .cfi_def_cfa_offset 5*4
+    push {r4-r7}
+    .cfi_def_cfa_offset 9*4
+    #endif
+
+    // Save the current stack pointer in oldStack.
+    #if defined(__thumb2__)
+    str sp, [r1]
+    #else
+    mov r2, sp
+    str r2, [r1]
+    #endif
+
+    // Switch to the new stack pointer.
+    mov sp, r0
+
     // Load state from new task and branch to the previous position in the
     // program.
     #if defined(__thumb2__)
diff --git a/src/internal/task/task_stack_cortexm.go b/src/internal/task/task_stack_cortexm.go
index f61688808..d39367d7e 100644
--- a/src/internal/task/task_stack_cortexm.go
+++ b/src/internal/task/task_stack_cortexm.go
@@ -52,19 +52,23 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.r5 = uintptr(args)
 }
 
-func (s *state) resume() {
-	switchToTask(s.sp)
+func (s *state) switchTo(current *Task) {
+	// This code is probably more complex than it needs to be.
+	// TODO: simplify this, ideally to a single function call.
+	if s == &mainTask.state {
+		switchToMain(&current.state.sp)
+	} else if current == &mainTask {
+		switchToTask(s.sp)
+	} else {
+		swapTask(s.sp, &current.state.sp)
+	}
 }
 
 //export tinygo_switchToTask
 func switchToTask(uintptr)
 
-//export tinygo_switchToScheduler
-func switchToScheduler(*uintptr)
-
-func (s *state) pause() {
-	switchToScheduler(&s.sp)
-}
+//export tinygo_switchToMain
+func switchToMain(*uintptr)
 
 // SystemStack returns the system stack pointer. On Cortex-M, it is always
 // available.
diff --git a/src/internal/task/task_stack_esp32.go b/src/internal/task/task_stack_esp32.go
index a76688447..475287c47 100644
--- a/src/internal/task/task_stack_esp32.go
+++ b/src/internal/task/task_stack_esp32.go
@@ -16,8 +16,6 @@ import (
 	"unsafe"
 )
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_esp8266.S that relies on the
 // exact layout of this struct.
@@ -60,18 +58,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.a2 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_esp8266.go b/src/internal/task/task_stack_esp8266.go
index 4a01a6fff..d8724975b 100644
--- a/src/internal/task/task_stack_esp8266.go
+++ b/src/internal/task/task_stack_esp8266.go
@@ -17,8 +17,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see task_stack_esp8266.S that relies on the
 // exact layout of this struct.
@@ -54,18 +52,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.a13 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/internal/task/task_stack_tinygoriscv.go b/src/internal/task/task_stack_tinygoriscv.go
index 36ef742e2..6b94e4fba 100644
--- a/src/internal/task/task_stack_tinygoriscv.go
+++ b/src/internal/task/task_stack_tinygoriscv.go
@@ -5,8 +5,6 @@ package task
 
 import "unsafe"
 
-var systemStack uintptr
-
 // calleeSavedRegs is the list of registers that must be saved and restored when
 // switching between tasks. Also see scheduler_riscv.S that relies on the
 // exact layout of this struct.
@@ -50,18 +48,12 @@ func (s *state) archInit(r *calleeSavedRegs, fn uintptr, args unsafe.Pointer) {
 	r.s1 = uintptr(args)
 }
 
-func (s *state) resume() {
-	swapTask(s.sp, &systemStack)
-}
-
-func (s *state) pause() {
-	newStack := systemStack
-	systemStack = 0
-	swapTask(newStack, &s.sp)
+func (s *state) switchTo(current *Task) {
+	swapTask(s.sp, &current.state.sp)
 }
 
 // SystemStack returns the system stack pointer when called from a task stack.
 // When called from the system stack, it returns 0.
 func SystemStack() uintptr {
-	return systemStack
+	return mainTask.state.sp
 }
diff --git a/src/runtime/gc_stack_raw.go b/src/runtime/gc_stack_raw.go
index e8fb0dc13..f5deeed24 100644
--- a/src/runtime/gc_stack_raw.go
+++ b/src/runtime/gc_stack_raw.go
@@ -14,7 +14,7 @@ func markStack() {
 	// Scan the current stack, and all current registers.
 	scanCurrentStack()
 
-	if !task.OnSystemStack() {
+	if !task.MainTask() {
 		// Mark system stack.
 		markRoots(getSystemStackPointer(), stackTop)
 	}
@@ -28,7 +28,7 @@ func scanstack(sp uintptr) {
 	// Mark current stack.
 	// This function is called by scanCurrentStack, after pushing all registers onto the stack.
 	// Callee-saved registers have been pushed onto stack by tinygo_localscan, so this will scan them too.
-	if task.OnSystemStack() {
+	if task.MainTask() {
 		// This is the system stack.
 		// Scan all words on the stack.
 		markRoots(sp, stackTop)
diff --git a/src/runtime/scheduler.go b/src/runtime/scheduler.go
index 196f6e3a7..f0b57a3fe 100644
--- a/src/runtime/scheduler.go
+++ b/src/runtime/scheduler.go
@@ -112,9 +112,7 @@ func addSleepTask(t *task.Task, duration timeUnit) {
 	*q = t
 }
 
-// Run the scheduler until all tasks have finished.
 func scheduler() {
-	// Main scheduler loop.
 	var now timeUnit
 	for !schedulerDone {
 		scheduleLog("")
@@ -164,7 +162,11 @@ func scheduler() {
 
 		// Run the given task.
 		scheduleLogTask("  run:", t)
-		t.Resume()
+		t.Switch()
+		if GOARCH != "wasm" {
+			// Don't return when using Asyncify.
+			return
+		}
 	}
 }
 
@@ -181,7 +183,7 @@ func minSched() {
 		}
 
 		scheduleLogTask("  run:", t)
-		t.Resume()
+		t.Switch()
 	}
 	scheduleLog("stop nested scheduler")
 }
diff --git a/src/runtime/scheduler_any.go b/src/runtime/scheduler_asyncify.go
similarity index 90%
rename from src/runtime/scheduler_any.go
rename to src/runtime/scheduler_asyncify.go
index d80979151..b3dd6a97f 100644
--- a/src/runtime/scheduler_any.go
+++ b/src/runtime/scheduler_asyncify.go
@@ -1,5 +1,5 @@
-//go:build !scheduler.none
-// +build !scheduler.none
+//go:build scheduler.asyncify
+// +build scheduler.asyncify
 
 package runtime
 
diff --git a/src/runtime/scheduler_tasks.go b/src/runtime/scheduler_tasks.go
index 67b475496..b1a675983 100644
--- a/src/runtime/scheduler_tasks.go
+++ b/src/runtime/scheduler_tasks.go
@@ -16,3 +16,27 @@ func getSystemStackPointer() uintptr {
 	}
 	return sp
 }
+
+// Pause the current task for a given time.
+//
+//go:linkname sleep time.Sleep
+func sleep(duration int64) {
+	if duration <= 0 {
+		return
+	}
+
+	addSleepTask(task.Current(), nanosecondsToTicks(duration))
+	task.Pause()
+}
+
+// run is called by the program entry point to execute the go program.
+// The main goroutine runs on the system stack, so initAll and callMain can be
+// called directly.
+func run() {
+	initHeap()
+	initAll()
+	callMain()
+	schedulerDone = true
+}
+
+const hasScheduler = true
diff --git a/targets/arm.ld b/targets/arm.ld
index e4155b909..5cab4835d 100644
--- a/targets/arm.ld
+++ b/targets/arm.ld
@@ -27,6 +27,7 @@ SECTIONS
     .stack (NOLOAD) :
     {
         . = ALIGN(4);
+        _stack_bottom = .;
         . += _stack_size;
         _stack_top = .;
     } >RAM
diff --git a/targets/avr.ld b/targets/avr.ld
index 720465ca4..2997da201 100644
--- a/targets/avr.ld
+++ b/targets/avr.ld
@@ -19,6 +19,7 @@ SECTIONS
 
     .stack (NOLOAD) :
     {
+        _stack_bottom = .;
         . += _stack_size;
         _stack_top = .;
     } >RAM
diff --git a/targets/esp32.ld b/targets/esp32.ld
index a8d161288..6fa2f822a 100644
--- a/targets/esp32.ld
+++ b/targets/esp32.ld
@@ -41,6 +41,7 @@ SECTIONS
     .stack (NOLOAD) :
     {
         . = ALIGN(16);
+        _stack_bottom = .;
         . += _stack_size;
         _stack_top = .;
     } >DRAM
diff --git a/targets/esp32c3.ld b/targets/esp32c3.ld
index b6be6e7ac..0a67b6367 100644
--- a/targets/esp32c3.ld
+++ b/targets/esp32c3.ld
@@ -70,6 +70,7 @@ SECTIONS
     .stack (NOLOAD) :
     {
         . = ALIGN(16);
+        _stack_bottom = .;
         . += _stack_size;
         _stack_top = .;
     } >DRAM
diff --git a/targets/esp8266.ld b/targets/esp8266.ld
index 9cbc42c14..8d35cc56e 100644
--- a/targets/esp8266.ld
+++ b/targets/esp8266.ld
@@ -67,6 +67,11 @@ _heap_end = ORIGIN(DRAM) + LENGTH(DRAM);
  */
 _stack_top = 0x40000000;
 
+/* This is the stack bottom according to this forum post:
+ * https://bbs.espressif.com/viewtopic.php?t=8879#p19038
+ */
+_stack_bottom = 0x3fffeb2c;
+
 /* Functions normally provided by a libc.
  * Source:
  * https://github.com/espressif/ESP8266_NONOS_SDK/blob/master/ld/eagle.rom.addr.v6.ld
diff --git a/targets/maixbit.ld b/targets/maixbit.ld
index 015a9beaf..7f099329c 100644
--- a/targets/maixbit.ld
+++ b/targets/maixbit.ld
@@ -57,6 +57,7 @@ SECTIONS
 }
 
 PROVIDE(_stack_top = ORIGIN(RAM) + LENGTH(RAM));
+_stack_bottom = _stack_top - _stack_size;
 
 /* For the memory allocator. */
 _heap_start = _ebss;
diff --git a/targets/mimxrt1062-teensy40.ld b/targets/mimxrt1062-teensy40.ld
index 70a717b72..3148ee0dc 100644
--- a/targets/mimxrt1062-teensy40.ld
+++ b/targets/mimxrt1062-teensy40.ld
@@ -59,6 +59,7 @@ SECTIONS
   .stack (NOLOAD) : {
 
     . = ALIGN(8);
+    _stack_bottom = .;
     . += _stack_size;
     _stack_top = .;
 
diff --git a/targets/riscv.ld b/targets/riscv.ld
index eecac6b47..cbc8dbe1c 100644
--- a/targets/riscv.ld
+++ b/targets/riscv.ld
@@ -19,6 +19,7 @@ SECTIONS
     .stack (NOLOAD) :
     {
         . = ALIGN(16);
+        _stack_bottom = .;
         . += _stack_size;
         _stack_top = .;
     } >RAM