compiler,runtime: implement stack-based scheduler

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()]