From a3ffb0d9eb948409c0898c6b1803401c9bc68ed4 Mon Sep 17 00:00:00 2001
From: Austin Clements <austin@google.com>
Date: Wed, 16 Oct 2019 19:10:06 -0400
Subject: [PATCH] runtime: asynchronous preemption function for x86

This adds asynchronous preemption function for amd64 and 386. These
functions spill and restore all register state that can be used by
user Go code.

For the moment we stub out the other arches.

For #10958, #24543.

Change-Id: I6f93fabe9875f4834922a5712362e79045c00aca
Reviewed-on: https://go-review.googlesource.com/c/go/+/201759
Run-TryBot: Austin Clements <austin@google.com>
Reviewed-by: Cherry Zhang <cherryyz@google.com>
---
 src/cmd/compile/internal/ssa/gen/386Ops.go   |   2 +
 src/cmd/compile/internal/ssa/gen/AMD64Ops.go |   2 +
 src/runtime/mkpreempt.go                     | 266 +++++++++++++++++++
 src/runtime/preempt.go                       |  15 ++
 src/runtime/preempt_386.s                    |  52 ++++
 src/runtime/preempt_amd64.s                  |  79 ++++++
 src/runtime/preempt_arm.s                    |   8 +
 src/runtime/preempt_arm64.s                  |   8 +
 src/runtime/preempt_mips64x.s                |  10 +
 src/runtime/preempt_mipsx.s                  |  10 +
 src/runtime/preempt_ppc64x.s                 |  10 +
 src/runtime/preempt_s390x.s                  |   8 +
 src/runtime/preempt_wasm.s                   |   8 +
 13 files changed, 478 insertions(+)
 create mode 100644 src/runtime/mkpreempt.go
 create mode 100644 src/runtime/preempt_386.s
 create mode 100644 src/runtime/preempt_amd64.s
 create mode 100644 src/runtime/preempt_arm.s
 create mode 100644 src/runtime/preempt_arm64.s
 create mode 100644 src/runtime/preempt_mips64x.s
 create mode 100644 src/runtime/preempt_mipsx.s
 create mode 100644 src/runtime/preempt_ppc64x.s
 create mode 100644 src/runtime/preempt_s390x.s
 create mode 100644 src/runtime/preempt_wasm.s

diff --git a/src/cmd/compile/internal/ssa/gen/386Ops.go b/src/cmd/compile/internal/ssa/gen/386Ops.go
index 4fb61adfeb20e..426fe48c2a938 100644
--- a/src/cmd/compile/internal/ssa/gen/386Ops.go
+++ b/src/cmd/compile/internal/ssa/gen/386Ops.go
@@ -45,6 +45,8 @@ var regNames386 = []string{
 	"X6",
 	"X7",
 
+	// If you add registers, update asyncPreempt in runtime
+
 	// pseudo-registers
 	"SB",
 }
diff --git a/src/cmd/compile/internal/ssa/gen/AMD64Ops.go b/src/cmd/compile/internal/ssa/gen/AMD64Ops.go
index cd2d0d61d19e0..420d0de9acfa8 100644
--- a/src/cmd/compile/internal/ssa/gen/AMD64Ops.go
+++ b/src/cmd/compile/internal/ssa/gen/AMD64Ops.go
@@ -63,6 +63,8 @@ var regNamesAMD64 = []string{
 	"X14",
 	"X15",
 
+	// If you add registers, update asyncPreempt in runtime
+
 	// pseudo-registers
 	"SB",
 }
diff --git a/src/runtime/mkpreempt.go b/src/runtime/mkpreempt.go
new file mode 100644
index 0000000000000..c28f89581dc73
--- /dev/null
+++ b/src/runtime/mkpreempt.go
@@ -0,0 +1,266 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build ignore
+
+// mkpreempt generates the asyncPreempt functions for each
+// architecture.
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"strings"
+)
+
+// Copied from cmd/compile/internal/ssa/gen/*Ops.go
+
+var regNames386 = []string{
+	"AX",
+	"CX",
+	"DX",
+	"BX",
+	"SP",
+	"BP",
+	"SI",
+	"DI",
+	"X0",
+	"X1",
+	"X2",
+	"X3",
+	"X4",
+	"X5",
+	"X6",
+	"X7",
+}
+
+var regNamesAMD64 = []string{
+	"AX",
+	"CX",
+	"DX",
+	"BX",
+	"SP",
+	"BP",
+	"SI",
+	"DI",
+	"R8",
+	"R9",
+	"R10",
+	"R11",
+	"R12",
+	"R13",
+	"R14",
+	"R15",
+	"X0",
+	"X1",
+	"X2",
+	"X3",
+	"X4",
+	"X5",
+	"X6",
+	"X7",
+	"X8",
+	"X9",
+	"X10",
+	"X11",
+	"X12",
+	"X13",
+	"X14",
+	"X15",
+}
+
+var out io.Writer
+
+var arches = map[string]func(){
+	"386":     gen386,
+	"amd64":   genAMD64,
+	"arm":     notImplemented,
+	"arm64":   notImplemented,
+	"mips64x": notImplemented,
+	"mipsx":   notImplemented,
+	"ppc64x":  notImplemented,
+	"s390x":   notImplemented,
+	"wasm":    genWasm,
+}
+var beLe = map[string]bool{"mips64x": true, "mipsx": true, "ppc64x": true}
+
+func main() {
+	flag.Parse()
+	if flag.NArg() > 0 {
+		out = os.Stdout
+		for _, arch := range flag.Args() {
+			gen, ok := arches[arch]
+			if !ok {
+				log.Fatalf("unknown arch %s", arch)
+			}
+			header(arch)
+			gen()
+		}
+		return
+	}
+
+	for arch, gen := range arches {
+		f, err := os.Create(fmt.Sprintf("preempt_%s.s", arch))
+		if err != nil {
+			log.Fatal(err)
+		}
+		out = f
+		header(arch)
+		gen()
+		if err := f.Close(); err != nil {
+			log.Fatal(err)
+		}
+	}
+}
+
+func header(arch string) {
+	fmt.Fprintf(out, "// Code generated by mkpreempt.go; DO NOT EDIT.\n\n")
+	if beLe[arch] {
+		base := arch[:len(arch)-1]
+		fmt.Fprintf(out, "// +build %s %sle\n\n", base, base)
+	}
+	fmt.Fprintf(out, "#include \"go_asm.h\"\n")
+	fmt.Fprintf(out, "#include \"textflag.h\"\n\n")
+	fmt.Fprintf(out, "TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0\n")
+}
+
+func p(f string, args ...interface{}) {
+	fmted := fmt.Sprintf(f, args...)
+	fmt.Fprintf(out, "\t%s\n", strings.Replace(fmted, "\n", "\n\t", -1))
+}
+
+type layout struct {
+	stack int
+	regs  []regPos
+}
+
+type regPos struct {
+	pos int
+
+	op  string
+	reg string
+
+	// If this register requires special save and restore, these
+	// give those operations with a %d placeholder for the stack
+	// offset.
+	save, restore string
+}
+
+func (l *layout) add(op, reg string, size int) {
+	l.regs = append(l.regs, regPos{op: op, reg: reg, pos: l.stack})
+	l.stack += size
+}
+
+func (l *layout) addSpecial(save, restore string, size int) {
+	l.regs = append(l.regs, regPos{save: save, restore: restore, pos: l.stack})
+	l.stack += size
+}
+
+func (l *layout) save() {
+	for _, reg := range l.regs {
+		if reg.save != "" {
+			p(reg.save, reg.pos)
+		} else {
+			p("%s %s, %d(SP)", reg.op, reg.reg, reg.pos)
+		}
+	}
+}
+
+func (l *layout) restore() {
+	for i := len(l.regs) - 1; i >= 0; i-- {
+		reg := l.regs[i]
+		if reg.restore != "" {
+			p(reg.restore, reg.pos)
+		} else {
+			p("%s %d(SP), %s", reg.op, reg.pos, reg.reg)
+		}
+	}
+}
+
+func gen386() {
+	p("PUSHFL")
+
+	// Save general purpose registers.
+	var l layout
+	for _, reg := range regNames386 {
+		if reg == "SP" || strings.HasPrefix(reg, "X") {
+			continue
+		}
+		l.add("MOVL", reg, 4)
+	}
+
+	// Save the 387 state.
+	l.addSpecial(
+		"FSAVE %d(SP)\nFLDCW runtime·controlWord64(SB)",
+		"FRSTOR %d(SP)",
+		108)
+
+	// Save SSE state only if supported.
+	lSSE := layout{stack: l.stack}
+	for i := 0; i < 8; i++ {
+		lSSE.add("MOVUPS", fmt.Sprintf("X%d", i), 16)
+	}
+
+	p("ADJSP $%d", lSSE.stack)
+	p("NOP SP")
+	l.save()
+	p("CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1\nJNE nosse")
+	lSSE.save()
+	p("nosse:")
+	p("CALL ·asyncPreempt2(SB)")
+	p("CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1\nJNE nosse2")
+	lSSE.restore()
+	p("nosse2:")
+	l.restore()
+	p("ADJSP $%d", -lSSE.stack)
+
+	p("POPFL")
+	p("RET")
+}
+
+func genAMD64() {
+	// Assign stack offsets.
+	var l layout
+	for _, reg := range regNamesAMD64 {
+		if reg == "SP" || reg == "BP" {
+			continue
+		}
+		if strings.HasPrefix(reg, "X") {
+			l.add("MOVUPS", reg, 16)
+		} else {
+			l.add("MOVQ", reg, 8)
+		}
+	}
+
+	// TODO: MXCSR register?
+
+	p("PUSHQ BP")
+	p("MOVQ SP, BP")
+	p("// Save flags before clobbering them")
+	p("PUSHFQ")
+	p("// obj doesn't understand ADD/SUB on SP, but does understand ADJSP")
+	p("ADJSP $%d", l.stack)
+	p("// But vet doesn't know ADJSP, so suppress vet stack checking")
+	p("NOP SP")
+	l.save()
+	p("CALL ·asyncPreempt2(SB)")
+	l.restore()
+	p("ADJSP $%d", -l.stack)
+	p("POPFQ")
+	p("POPQ BP")
+	p("RET")
+}
+
+func genWasm() {
+	p("// No async preemption on wasm")
+	p("UNDEF")
+}
+
+func notImplemented() {
+	p("// Not implemented yet")
+	p("JMP ·abort(SB)")
+}
diff --git a/src/runtime/preempt.go b/src/runtime/preempt.go
index 96eaa3488b80f..57ec493b8d046 100644
--- a/src/runtime/preempt.go
+++ b/src/runtime/preempt.go
@@ -232,3 +232,18 @@ func resumeG(state suspendGState) {
 func canPreemptM(mp *m) bool {
 	return mp.locks == 0 && mp.mallocing == 0 && mp.preemptoff == "" && mp.p.ptr().status == _Prunning
 }
+
+//go:generate go run mkpreempt.go
+
+// asyncPreempt saves all user registers and calls asyncPreempt2.
+//
+// When stack scanning encounters an asyncPreempt frame, it scans that
+// frame and its parent frame conservatively.
+//
+// asyncPreempt is implemented in assembly.
+func asyncPreempt()
+
+//go:nosplit
+func asyncPreempt2() {
+	// TODO: Enter scheduler
+}
diff --git a/src/runtime/preempt_386.s b/src/runtime/preempt_386.s
new file mode 100644
index 0000000000000..a7961e02ce6ea
--- /dev/null
+++ b/src/runtime/preempt_386.s
@@ -0,0 +1,52 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	PUSHFL
+	ADJSP $264
+	NOP SP
+	MOVL AX, 0(SP)
+	MOVL CX, 4(SP)
+	MOVL DX, 8(SP)
+	MOVL BX, 12(SP)
+	MOVL BP, 16(SP)
+	MOVL SI, 20(SP)
+	MOVL DI, 24(SP)
+	FSAVE 28(SP)
+	FLDCW runtime·controlWord64(SB)
+	CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1
+	JNE nosse
+	MOVUPS X0, 136(SP)
+	MOVUPS X1, 152(SP)
+	MOVUPS X2, 168(SP)
+	MOVUPS X3, 184(SP)
+	MOVUPS X4, 200(SP)
+	MOVUPS X5, 216(SP)
+	MOVUPS X6, 232(SP)
+	MOVUPS X7, 248(SP)
+	nosse:
+	CALL ·asyncPreempt2(SB)
+	CMPB internal∕cpu·X86+const_offsetX86HasSSE2(SB), $1
+	JNE nosse2
+	MOVUPS 248(SP), X7
+	MOVUPS 232(SP), X6
+	MOVUPS 216(SP), X5
+	MOVUPS 200(SP), X4
+	MOVUPS 184(SP), X3
+	MOVUPS 168(SP), X2
+	MOVUPS 152(SP), X1
+	MOVUPS 136(SP), X0
+	nosse2:
+	FRSTOR 28(SP)
+	MOVL 24(SP), DI
+	MOVL 20(SP), SI
+	MOVL 16(SP), BP
+	MOVL 12(SP), BX
+	MOVL 8(SP), DX
+	MOVL 4(SP), CX
+	MOVL 0(SP), AX
+	ADJSP $-264
+	POPFL
+	RET
diff --git a/src/runtime/preempt_amd64.s b/src/runtime/preempt_amd64.s
new file mode 100644
index 0000000000000..d50c2f3a5169c
--- /dev/null
+++ b/src/runtime/preempt_amd64.s
@@ -0,0 +1,79 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	PUSHQ BP
+	MOVQ SP, BP
+	// Save flags before clobbering them
+	PUSHFQ
+	// obj doesn't understand ADD/SUB on SP, but does understand ADJSP
+	ADJSP $368
+	// But vet doesn't know ADJSP, so suppress vet stack checking
+	NOP SP
+	MOVQ AX, 0(SP)
+	MOVQ CX, 8(SP)
+	MOVQ DX, 16(SP)
+	MOVQ BX, 24(SP)
+	MOVQ SI, 32(SP)
+	MOVQ DI, 40(SP)
+	MOVQ R8, 48(SP)
+	MOVQ R9, 56(SP)
+	MOVQ R10, 64(SP)
+	MOVQ R11, 72(SP)
+	MOVQ R12, 80(SP)
+	MOVQ R13, 88(SP)
+	MOVQ R14, 96(SP)
+	MOVQ R15, 104(SP)
+	MOVUPS X0, 112(SP)
+	MOVUPS X1, 128(SP)
+	MOVUPS X2, 144(SP)
+	MOVUPS X3, 160(SP)
+	MOVUPS X4, 176(SP)
+	MOVUPS X5, 192(SP)
+	MOVUPS X6, 208(SP)
+	MOVUPS X7, 224(SP)
+	MOVUPS X8, 240(SP)
+	MOVUPS X9, 256(SP)
+	MOVUPS X10, 272(SP)
+	MOVUPS X11, 288(SP)
+	MOVUPS X12, 304(SP)
+	MOVUPS X13, 320(SP)
+	MOVUPS X14, 336(SP)
+	MOVUPS X15, 352(SP)
+	CALL ·asyncPreempt2(SB)
+	MOVUPS 352(SP), X15
+	MOVUPS 336(SP), X14
+	MOVUPS 320(SP), X13
+	MOVUPS 304(SP), X12
+	MOVUPS 288(SP), X11
+	MOVUPS 272(SP), X10
+	MOVUPS 256(SP), X9
+	MOVUPS 240(SP), X8
+	MOVUPS 224(SP), X7
+	MOVUPS 208(SP), X6
+	MOVUPS 192(SP), X5
+	MOVUPS 176(SP), X4
+	MOVUPS 160(SP), X3
+	MOVUPS 144(SP), X2
+	MOVUPS 128(SP), X1
+	MOVUPS 112(SP), X0
+	MOVQ 104(SP), R15
+	MOVQ 96(SP), R14
+	MOVQ 88(SP), R13
+	MOVQ 80(SP), R12
+	MOVQ 72(SP), R11
+	MOVQ 64(SP), R10
+	MOVQ 56(SP), R9
+	MOVQ 48(SP), R8
+	MOVQ 40(SP), DI
+	MOVQ 32(SP), SI
+	MOVQ 24(SP), BX
+	MOVQ 16(SP), DX
+	MOVQ 8(SP), CX
+	MOVQ 0(SP), AX
+	ADJSP $-368
+	POPFQ
+	POPQ BP
+	RET
diff --git a/src/runtime/preempt_arm.s b/src/runtime/preempt_arm.s
new file mode 100644
index 0000000000000..5697268ce1612
--- /dev/null
+++ b/src/runtime/preempt_arm.s
@@ -0,0 +1,8 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// Not implemented yet
+	JMP ·abort(SB)
diff --git a/src/runtime/preempt_arm64.s b/src/runtime/preempt_arm64.s
new file mode 100644
index 0000000000000..5697268ce1612
--- /dev/null
+++ b/src/runtime/preempt_arm64.s
@@ -0,0 +1,8 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// Not implemented yet
+	JMP ·abort(SB)
diff --git a/src/runtime/preempt_mips64x.s b/src/runtime/preempt_mips64x.s
new file mode 100644
index 0000000000000..713c074abf568
--- /dev/null
+++ b/src/runtime/preempt_mips64x.s
@@ -0,0 +1,10 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+// +build mips64 mips64le
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// Not implemented yet
+	JMP ·abort(SB)
diff --git a/src/runtime/preempt_mipsx.s b/src/runtime/preempt_mipsx.s
new file mode 100644
index 0000000000000..2538a2ee004c9
--- /dev/null
+++ b/src/runtime/preempt_mipsx.s
@@ -0,0 +1,10 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+// +build mips mipsle
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// Not implemented yet
+	JMP ·abort(SB)
diff --git a/src/runtime/preempt_ppc64x.s b/src/runtime/preempt_ppc64x.s
new file mode 100644
index 0000000000000..7e4315a37f779
--- /dev/null
+++ b/src/runtime/preempt_ppc64x.s
@@ -0,0 +1,10 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+// +build ppc64 ppc64le
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// Not implemented yet
+	JMP ·abort(SB)
diff --git a/src/runtime/preempt_s390x.s b/src/runtime/preempt_s390x.s
new file mode 100644
index 0000000000000..5697268ce1612
--- /dev/null
+++ b/src/runtime/preempt_s390x.s
@@ -0,0 +1,8 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// Not implemented yet
+	JMP ·abort(SB)
diff --git a/src/runtime/preempt_wasm.s b/src/runtime/preempt_wasm.s
new file mode 100644
index 0000000000000..0cf57d3d2265b
--- /dev/null
+++ b/src/runtime/preempt_wasm.s
@@ -0,0 +1,8 @@
+// Code generated by mkpreempt.go; DO NOT EDIT.
+
+#include "go_asm.h"
+#include "textflag.h"
+
+TEXT ·asyncPreempt(SB),NOSPLIT|NOFRAME,$0-0
+	// No async preemption on wasm
+	UNDEF