compiler,runtime: implement stack-based scheduler

This scheduler is intended to live along the (stackless) coroutine based
scheduler which is needed for WebAssembly and unsupported platforms. The
stack based scheduler is somewhat simpler in implementation as it does
not require full program transform passes and supports things like
function pointers and interface methods out of the box with no changes.

Code size is reduced in most cases, even in the case where no scheduler
scheduler is used at all. I'm not exactly sure why but these changes
likely allowed some further optimizations somewhere. Even RAM is
slightly reduced, perhaps some global was elminated in the process as
well.

Author: aykevl

compiler/channel.go

@@ -69,7 +69,7 @@ func (c *Compiler) emitChanRecv(frame *Frame, unop *ssa.UnOp) llvm.Value {
 	c.emitLifetimeEnd(valueAllocaCast, valueAllocaSize)
 
 	if unop.CommaOk {
-		commaOk := c.createRuntimeCall("getTaskPromiseData", []llvm.Value{coroutine}, "chan.commaOk.wide")
+		commaOk := c.createRuntimeCall("getTaskStateData", []llvm.Value{coroutine}, "chan.commaOk.wide")
 		commaOk = c.builder.CreateTrunc(commaOk, c.ctx.Int1Type(), "chan.commaOk")
 		tuple := llvm.Undef(c.ctx.StructType([]llvm.Type{valueType, c.ctx.Int1Type()}, false))
 		tuple = c.builder.CreateInsertValue(tuple, received, 0, "")
@@ -95,7 +95,7 @@ func (c *Compiler) emitSelect(frame *Frame, expr *ssa.Select) llvm.Value {
 		if expr.Blocking {
 			// Blocks forever:
 			//     select {}
-			c.createRuntimeCall("deadlockStub", nil, "")
+			c.createRuntimeCall("deadlock", nil, "")
 			return llvm.Undef(llvmType)
 		} else {
 			// No-op:

compiler/compiler.go

@@ -30,6 +30,20 @@ func init() {
 // The TinyGo import path.
 const tinygoPath = "github.com/tinygo-org/tinygo"
 
+// functionsUsedInTransform is a list of function symbols that may be used
+// during TinyGo optimization passes so they have to be marked as external
+// linkage until all TinyGo passes have finished.
+var functionsUsedInTransforms = []string{
+	"runtime.alloc",
+	"runtime.free",
+	"runtime.sleepTask",
+	"runtime.setTaskStatePtr",
+	"runtime.getTaskStatePtr",
+	"runtime.activateTask",
+	"runtime.scheduler",
+	"runtime.startGoroutine",
+}
+
 // Configure the compiler.
 type Config struct {
 	Triple        string   // LLVM target triple, e.g. x86_64-unknown-linux-gnu (empty string means default)
@@ -38,6 +52,7 @@ type Config struct {
 	GOOS          string   //
 	GOARCH        string   //
 	GC            string   // garbage collection strategy
+	Scheduler     string   // scheduler implementation ("coroutines" or "tasks")
 	PanicStrategy string   // panic strategy ("print" or "trap")
 	CFlags        []string // cflags to pass to cgo
 	LDFlags       []string // ldflags to pass to cgo
@@ -173,6 +188,17 @@ func (c *Compiler) selectGC() string {
 	return "conservative"
 }
 
+// selectScheduler picks an appropriate scheduler for the target if none was
+// given.
+func (c *Compiler) selectScheduler() string {
+	if c.Scheduler != "" {
+		// A scheduler was specified in the target description.
+		return c.Scheduler
+	}
+	// Fall back to coroutines, which are supported everywhere.
+	return "coroutines"
+}
+
 // Compile the given package path or .go file path. Return an error when this
 // fails (in any stage).
 func (c *Compiler) Compile(mainPath string) []error {
@@ -189,6 +215,7 @@ func (c *Compiler) Compile(mainPath string) []error {
 	if err != nil {
 		return []error{err}
 	}
+	buildTags := append([]string{"tinygo", "gc." + c.selectGC(), "scheduler." + c.selectScheduler()}, c.BuildTags...)
 	lprogram := &loader.Program{
 		Build: &build.Context{
 			GOARCH:      c.GOARCH,
@@ -198,7 +225,7 @@ func (c *Compiler) Compile(mainPath string) []error {
 			CgoEnabled:  true,
 			UseAllFiles: false,
 			Compiler:    "gc", // must be one of the recognized compilers
-			BuildTags:   append([]string{"tinygo", "gc." + c.selectGC()}, c.BuildTags...),
+			BuildTags:   buildTags,
 		},
 		OverlayBuild: &build.Context{
 			GOARCH:      c.GOARCH,
@@ -208,7 +235,7 @@ func (c *Compiler) Compile(mainPath string) []error {
 			CgoEnabled:  true,
 			UseAllFiles: false,
 			Compiler:    "gc", // must be one of the recognized compilers
-			BuildTags:   append([]string{"tinygo", "gc." + c.selectGC()}, c.BuildTags...),
+			BuildTags:   buildTags,
 		},
 		OverlayPath: func(path string) string {
 			// Return the (overlay) import path when it should be overlaid, and
@@ -335,13 +362,15 @@ func (c *Compiler) Compile(mainPath string) []error {
 	// would be optimized away.
 	realMain := c.mod.NamedFunction(c.ir.MainPkg().Pkg.Path() + ".main")
 	realMain.SetLinkage(llvm.ExternalLinkage) // keep alive until goroutine lowering
-	c.mod.NamedFunction("runtime.alloc").SetLinkage(llvm.ExternalLinkage)
-	c.mod.NamedFunction("runtime.free").SetLinkage(llvm.ExternalLinkage)
-	c.mod.NamedFunction("runtime.sleepTask").SetLinkage(llvm.ExternalLinkage)
-	c.mod.NamedFunction("runtime.setTaskPromisePtr").SetLinkage(llvm.ExternalLinkage)
-	c.mod.NamedFunction("runtime.getTaskPromisePtr").SetLinkage(llvm.ExternalLinkage)
-	c.mod.NamedFunction("runtime.activateTask").SetLinkage(llvm.ExternalLinkage)
-	c.mod.NamedFunction("runtime.scheduler").SetLinkage(llvm.ExternalLinkage)
+
+	// Make sure these functions are kept in tact during TinyGo transformation passes.
+	for _, name := range functionsUsedInTransforms {
+		fn := c.mod.NamedFunction(name)
+		if fn.IsNil() {
+			continue
+		}
+		fn.SetLinkage(llvm.ExternalLinkage)
+	}
 
 	// Load some attributes
 	getAttr := func(attrName string) llvm.Attribute {
@@ -1041,25 +1070,21 @@ func (c *Compiler) parseInstr(frame *Frame, instr ssa.Instruction) {
 		}
 		calleeFn := c.ir.GetFunction(callee)
 
-		// Mark this function as a 'go' invocation and break invalid
-		// interprocedural optimizations. For example, heap-to-stack
-		// transformations are not sound as goroutines can outlive their parent.
-		calleeType := calleeFn.LLVMFn.Type()
-		calleeValue := c.builder.CreatePtrToInt(calleeFn.LLVMFn, c.uintptrType, "")
-		calleeValue = c.createRuntimeCall("makeGoroutine", []llvm.Value{calleeValue}, "")
-		calleeValue = c.builder.CreateIntToPtr(calleeValue, calleeType, "")
-
 		// Get all function parameters to pass to the goroutine.
 		var params []llvm.Value
 		for _, param := range instr.Call.Args {
 			params = append(params, c.getValue(frame, param))
 		}
-		if !calleeFn.IsExported() {
+		if !calleeFn.IsExported() && c.selectScheduler() != "tasks" {
+			// For coroutine scheduling, this is only required when calling an
+			// external function.
+			// For tasks, because all params are stored in a single object, no
+			// unnecessary parameters should be stored anyway.
 			params = append(params, llvm.Undef(c.i8ptrType)) // context parameter
 			params = append(params, llvm.Undef(c.i8ptrType)) // parent coroutine handle
 		}
 
-		c.createCall(calleeValue, params, "")
+		c.emitStartGoroutine(calleeFn.LLVMFn, params)
 	case *ssa.If:
 		cond := c.getValue(frame, instr.Cond)
 		block := instr.Block()

compiler/goroutine-lowering.go

@@ -1,5 +1,16 @@
 package compiler
 
+// This file implements lowering for the goroutine scheduler. There are two
+// scheduler implementations, one based on tasks (like RTOSes and the main Go
+// runtime) and one based on a coroutine compiler transformation. The task based
+// implementation requires very little work from the compiler but is not very
+// portable (in particular, it is very hard if not impossible to support on
+// WebAssembly). The coroutine based one requires a lot of work by the compiler
+// to implement, but can run virtually anywhere with a single scheduler
+// implementation.
+//
+// The below description is for the coroutine based scheduler.
+//
 // This file lowers goroutine pseudo-functions into coroutines scheduled by a
 // scheduler at runtime. It uses coroutine support in LLVM for this
 // transformation: https://llvm.org/docs/Coroutines.html
@@ -62,7 +73,7 @@ package compiler
 //         llvm.suspend(hdl)                       // suspend point
 //         println("some other operation")
 //         var i *int                              // allocate space on the stack for the return value
-//         runtime.setTaskPromisePtr(hdl, &i)      // store return value alloca in our coroutine promise
+//         runtime.setTaskStatePtr(hdl, &i)        // store return value alloca in our coroutine promise
 //         bar(hdl)                                // await, pass a continuation (hdl) to bar
 //         llvm.suspend(hdl)                       // suspend point, wait for the callee to re-activate
 //         println("done", *i)
@@ -106,10 +117,65 @@ type asyncFunc struct {
 	unreachableBlock llvm.BasicBlock
 }
 
-// LowerGoroutines is a pass called during optimization that transforms the IR
-// into one where all blocking functions are turned into goroutines and blocking
-// calls into await calls.
+// LowerGoroutines performs some IR transformations necessary to support
+// goroutines. It does something different based on whether it uses the
+// coroutine or the tasks implementation of goroutines, and whether goroutines
+// are necessary at all.
 func (c *Compiler) LowerGoroutines() error {
+	switch c.selectScheduler() {
+	case "coroutines":
+		return c.lowerCoroutines()
+	case "tasks":
+		return c.lowerTasks()
+	default:
+		panic("unknown scheduler type")
+	}
+}
+
+// lowerTasks starts the main goroutine and then runs the scheduler.
+// This is enough compiler-level transformation for the task-based scheduler.
+func (c *Compiler) lowerTasks() error {
+	uses := getUses(c.mod.NamedFunction("runtime.callMain"))
+	if len(uses) != 1 || uses[0].IsACallInst().IsNil() {
+		panic("expected exactly 1 call of runtime.callMain, check the entry point")
+	}
+	mainCall := uses[0]
+
+	realMain := c.mod.NamedFunction(c.ir.MainPkg().Pkg.Path() + ".main")
+	if len(getUses(c.mod.NamedFunction("runtime.startGoroutine"))) != 0 {
+		// Program needs a scheduler. Start main.main as a goroutine and start
+		// the scheduler.
+		realMainWrapper := c.createGoroutineStartWrapper(realMain)
+		c.builder.SetInsertPointBefore(mainCall)
+		zero := llvm.ConstInt(c.uintptrType, 0, false)
+		c.createRuntimeCall("startGoroutine", []llvm.Value{realMainWrapper, zero}, "")
+		c.createRuntimeCall("scheduler", nil, "")
+	} else {
+		// Program doesn't need a scheduler. Call main.main directly.
+		c.builder.SetInsertPointBefore(mainCall)
+		params := []llvm.Value{
+			llvm.Undef(c.i8ptrType), // unused context parameter
+			llvm.Undef(c.i8ptrType), // unused coroutine handle
+		}
+		c.createCall(realMain, params, "")
+		// runtime.Goexit isn't needed so let it be optimized away by
+		// globalopt.
+		c.mod.NamedFunction("runtime.Goexit").SetLinkage(llvm.InternalLinkage)
+	}
+	mainCall.EraseFromParentAsInstruction()
+
+	// main.main was set to external linkage during IR construction. Set it to
+	// internal linkage to enable interprocedural optimizations.
+	realMain.SetLinkage(llvm.InternalLinkage)
+
+	return nil
+}
+
+// lowerCoroutines transforms the IR into one where all blocking functions are
+// turned into goroutines and blocking calls into await calls. It also makes
+// sure that the first coroutine is started and the coroutine scheduler will be
+// run.
+func (c *Compiler) lowerCoroutines() error {
 	needsScheduler, err := c.markAsyncFunctions()
 	if err != nil {
 		return err
@@ -144,12 +210,6 @@ func (c *Compiler) LowerGoroutines() error {
 	// main.main was set to external linkage during IR construction. Set it to
 	// internal linkage to enable interprocedural optimizations.
 	realMain.SetLinkage(llvm.InternalLinkage)
-	c.mod.NamedFunction("runtime.alloc").SetLinkage(llvm.InternalLinkage)
-	c.mod.NamedFunction("runtime.free").SetLinkage(llvm.InternalLinkage)
-	c.mod.NamedFunction("runtime.sleepTask").SetLinkage(llvm.InternalLinkage)
-	c.mod.NamedFunction("runtime.setTaskPromisePtr").SetLinkage(llvm.InternalLinkage)
-	c.mod.NamedFunction("runtime.getTaskPromisePtr").SetLinkage(llvm.InternalLinkage)
-	c.mod.NamedFunction("runtime.scheduler").SetLinkage(llvm.InternalLinkage)
 
 	return nil
 }
@@ -173,9 +233,9 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 	if !sleep.IsNil() {
 		worklist = append(worklist, sleep)
 	}
-	deadlockStub := c.mod.NamedFunction("runtime.deadlockStub")
-	if !deadlockStub.IsNil() {
-		worklist = append(worklist, deadlockStub)
+	deadlock := c.mod.NamedFunction("runtime.deadlock")
+	if !deadlock.IsNil() {
+		worklist = append(worklist, deadlock)
 	}
 	chanSend := c.mod.NamedFunction("runtime.chanSend")
 	if !chanSend.IsNil() {
@@ -300,7 +360,7 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 
 	// Transform all async functions into coroutines.
 	for _, f := range asyncList {
-		if f == sleep || f == deadlockStub || f == chanSend || f == chanRecv {
+		if f == sleep || f == deadlock || f == chanSend || f == chanRecv {
 			continue
 		}
 
@@ -317,7 +377,7 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 			for inst := bb.FirstInstruction(); !inst.IsNil(); inst = llvm.NextInstruction(inst) {
 				if !inst.IsACallInst().IsNil() {
 					callee := inst.CalledValue()
-					if _, ok := asyncFuncs[callee]; !ok || callee == sleep || callee == deadlockStub || callee == chanSend || callee == chanRecv {
+					if _, ok := asyncFuncs[callee]; !ok || callee == sleep || callee == deadlock || callee == chanSend || callee == chanRecv {
 						continue
 					}
 					asyncCalls = append(asyncCalls, inst)
@@ -365,7 +425,7 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 				retvalAlloca = c.builder.CreateAlloca(inst.Type(), "coro.retvalAlloca")
 				c.builder.SetInsertPointBefore(inst)
 				data := c.builder.CreateBitCast(retvalAlloca, c.i8ptrType, "")
-				c.createRuntimeCall("setTaskPromisePtr", []llvm.Value{frame.taskHandle, data}, "")
+				c.createRuntimeCall("setTaskStatePtr", []llvm.Value{frame.taskHandle, data}, "")
 			}
 
 			// Suspend.
@@ -403,7 +463,7 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 			var parentHandle llvm.Value
 			if f.Linkage() == llvm.ExternalLinkage {
 				// Exported function.
-				// Note that getTaskPromisePtr will panic if it is called with
+				// Note that getTaskStatePtr will panic if it is called with
 				// a nil pointer, so blocking exported functions that try to
 				// return anything will not work.
 				parentHandle = llvm.ConstPointerNull(c.i8ptrType)
@@ -423,7 +483,7 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 				// Return this value by writing to the pointer stored in the
 				// parent handle. The parent coroutine has made an alloca that
 				// we can write to to store our return value.
-				returnValuePtr := c.createRuntimeCall("getTaskPromisePtr", []llvm.Value{parentHandle}, "coro.parentData")
+				returnValuePtr := c.createRuntimeCall("getTaskStatePtr", []llvm.Value{parentHandle}, "coro.parentData")
 				alloca := c.builder.CreateBitCast(returnValuePtr, llvm.PointerType(inst.Operand(0).Type(), 0), "coro.parentAlloca")
 				c.builder.CreateStore(inst.Operand(0), alloca)
 			default:
@@ -502,9 +562,9 @@ func (c *Compiler) markAsyncFunctions() (needsScheduler bool, err error) {
 		sleepCall.EraseFromParentAsInstruction()
 	}
 
-	// Transform calls to runtime.deadlockStub into coroutine suspends (without
+	// Transform calls to runtime.deadlock into coroutine suspends (without
 	// resume).
-	for _, deadlockCall := range getUses(deadlockStub) {
+	for _, deadlockCall := range getUses(deadlock) {
 		// deadlockCall must be a call instruction.
 		frame := asyncFuncs[deadlockCall.InstructionParent().Parent()]