cmd/compile: add ast nodes and basic block counters

Current patch adds the counters to the AST and SSA nodes. The counters
are loaded from the pprof file, no profile format changes needed.

To use basic block counters you should add an option -pgobb:

go build -a -pgobb -pgo=file.prof

Fixes golang#65466

Change-Id: I5d6be7d87f384625259a9ba794744a652060de4e

Author: alexanius

src/cmd/compile/internal/base/flag.go

@@ -126,6 +126,7 @@ type CmdFlags struct {
 	TrimPath           string       "help:\"remove `prefix` from recorded source file paths\""
 	WB                 bool         "help:\"enable write barrier\"" // TODO: remove
 	PgoProfile         string       "help:\"read profile or pre-process profile from `file`\""
+	PgoBbProfile       bool         "help:\"use basic block counters in pgo mode\""
 	ErrorURL           bool         "help:\"print explanatory URL with error message if applicable\""
 
 	// Configuration derived from flags; not a flag itself.

src/cmd/compile/internal/gc/compile.go

@@ -180,6 +180,10 @@ func compileFunctions(profile *pgoir.Profile) {
 		for _, fn := range fns {
 			fn := fn
 			queue(func(worker int) {
+				if base.Flag.PgoBbProfile && profile != nil {
+					pgoir.CorrectProfileAfterInline(profile.Prof, fn)
+				}
+
 				ssagen.Compile(fn, worker, profile)
 				compile(fn.Closures)
 				wg.Done()

src/cmd/compile/internal/ir/fmt.go

@@ -133,7 +133,7 @@ func fmtNode(n Node, s fmt.State, verb rune) {
 	// %+v prints Dump.
 	// Otherwise we print Go syntax.
 	if s.Flag('+') && verb == 'v' {
-		dumpNode(s, n, 1)
+		dumpNode(s, n, 1, nil)
 		return
 	}
 
@@ -872,7 +872,7 @@ func ellipsisIf(b bool) string {
 func (l Nodes) Format(s fmt.State, verb rune) {
 	if s.Flag('+') && verb == 'v' {
 		// %+v is DumpList output
-		dumpNodes(s, l, 1)
+		dumpNodes(s, l, 1, nil)
 		return
 	}
 
@@ -904,20 +904,23 @@ func Dump(s string, n Node) {
 // DumpList prints the message s followed by a debug dump of each node in the list.
 func DumpList(s string, list Nodes) {
 	var buf bytes.Buffer
-	FDumpList(&buf, s, list)
+	FDumpList(&buf, s, list, nil)
 	os.Stdout.Write(buf.Bytes())
 }
 
 // FDumpList prints to w the message s followed by a debug dump of each node in the list.
-func FDumpList(w io.Writer, s string, list Nodes) {
+func FDumpList(w io.Writer, s string, list Nodes, f *Func) {
 	io.WriteString(w, s)
-	dumpNodes(w, list, 1)
+	dumpNodes(w, list, 1, f)
 	io.WriteString(w, "\n")
 }
 
 // indent prints indentation to w.
-func indent(w io.Writer, depth int) {
+func indent(w io.Writer, depth int, counter int64) {
 	fmt.Fprint(w, "\n")
+	if base.Flag.PgoBbProfile {
+		fmt.Fprintf(w, "%d ", counter)
+	}
 	for i := 0; i < depth; i++ {
 		fmt.Fprint(w, ".   ")
 	}
@@ -1048,8 +1051,8 @@ func dumpNodeHeader(w io.Writer, n Node) {
 	}
 }
 
-func dumpNode(w io.Writer, n Node, depth int) {
-	indent(w, depth)
+func dumpNode(w io.Writer, n Node, depth int, f *Func) {
+	indent(w, depth, GetCounter(f, n))
 	if depth > 40 {
 		fmt.Fprint(w, "...")
 		return
@@ -1062,8 +1065,8 @@ func dumpNode(w io.Writer, n Node, depth int) {
 
 	if len(n.Init()) != 0 {
 		fmt.Fprintf(w, "%+v-init", n.Op())
-		dumpNodes(w, n.Init(), depth+1)
-		indent(w, depth)
+		dumpNodes(w, n.Init(), depth+1, f)
+		indent(w, depth, GetCounter(f, n))
 	}
 
 	switch n.Op() {
@@ -1116,23 +1119,23 @@ func dumpNode(w io.Writer, n Node, depth int) {
 		dumpNodeHeader(w, n)
 		fn := n
 		if len(fn.Dcl) > 0 {
-			indent(w, depth)
+			indent(w, depth, GetCounter(f, n))
 			fmt.Fprintf(w, "%+v-Dcl", n.Op())
 			for _, dcl := range n.Dcl {
-				dumpNode(w, dcl, depth+1)
+				dumpNode(w, dcl, depth+1, f)
 			}
 		}
 		if len(fn.ClosureVars) > 0 {
-			indent(w, depth)
+			indent(w, depth, GetCounter(f, n))
 			fmt.Fprintf(w, "%+v-ClosureVars", n.Op())
 			for _, cv := range fn.ClosureVars {
-				dumpNode(w, cv, depth+1)
+				dumpNode(w, cv, depth+1, f)
 			}
 		}
 		if len(fn.Body) > 0 {
-			indent(w, depth)
+			indent(w, depth, GetCounter(f, n))
 			fmt.Fprintf(w, "%+v-body", n.Op())
-			dumpNodes(w, fn.Body, depth+1)
+			dumpNodes(w, fn.Body, depth+1, f)
 		}
 		return
 	}
@@ -1164,30 +1167,34 @@ func dumpNode(w io.Writer, n Node, depth int) {
 		switch val := vf.Interface().(type) {
 		case Node:
 			if name != "" {
-				indent(w, depth)
+				indent(w, depth, GetCounter(f, n))
 				fmt.Fprintf(w, "%+v-%s", n.Op(), name)
 			}
-			dumpNode(w, val, depth+1)
+			dumpNode(w, val, depth+1, f)
 		case Nodes:
 			if len(val) == 0 {
 				continue
 			}
 			if name != "" {
-				indent(w, depth)
+				c := GetCounter(f, n)
+				if len(val) > 0 {
+					c = GetCounter(f, val[0])
+				}
+				indent(w, depth, c)
 				fmt.Fprintf(w, "%+v-%s", n.Op(), name)
 			}
-			dumpNodes(w, val, depth+1)
+			dumpNodes(w, val, depth+1, f)
 		default:
 			if vf.Kind() == reflect.Slice && vf.Type().Elem().Implements(nodeType) {
 				if vf.Len() == 0 {
 					continue
 				}
 				if name != "" {
-					indent(w, depth)
+					indent(w, depth, GetCounter(f, n))
 					fmt.Fprintf(w, "%+v-%s", n.Op(), name)
 				}
 				for i, n := 0, vf.Len(); i < n; i++ {
-					dumpNode(w, vf.Index(i).Interface().(Node), depth+1)
+					dumpNode(w, vf.Index(i).Interface().(Node), depth+1, f)
 				}
 			}
 		}
@@ -1196,13 +1203,13 @@ func dumpNode(w io.Writer, n Node, depth int) {
 
 var nodeType = reflect.TypeOf((*Node)(nil)).Elem()
 
-func dumpNodes(w io.Writer, list Nodes, depth int) {
+func dumpNodes(w io.Writer, list Nodes, depth int, f *Func) {
 	if len(list) == 0 {
 		fmt.Fprintf(w, " <nil>")
 		return
 	}
 
 	for _, n := range list {
-		dumpNode(w, n, depth)
+		dumpNode(w, n, depth, f)
 	}
 }

src/cmd/compile/internal/ir/func.go

@@ -142,6 +142,56 @@ type Func struct {
 	// WasmImport is used by the //go:wasmimport directive to store info about
 	// a WebAssembly function import.
 	WasmImport *WasmImport
+
+	ProfTable NodeProfTable
+}
+
+// The type of a counter in node
+type Counter = int64
+
+// The type of index in the function counter table
+type counterIndex = string
+
+// The type of a table with counters
+type NodeProfTable = map[counterIndex]Counter
+
+// ShouldSetCounter returns true if this node type should have a counter
+func ShouldSetCounter(n Node) bool {
+	op := n.Op()
+	return op != ONAME && op != OLITERAL
+}
+
+// Set the counter c to the node n in the function fn
+func SetCounter(fn *Func, n Node, c Counter) {
+	if !ShouldSetCounter(n) {
+		return
+	}
+
+	idx := fmt.Sprintf("%d:%d", n.Pos().FileIndex(), n.Pos().Line())
+	t := fn.ProfTable
+	if t == nil {
+		// TODO why
+		return
+	}
+	t[idx] = c
+}
+
+// Get the counter c to the node n in the function fn
+func GetCounter(fn *Func, n Node) Counter {
+	if !ShouldSetCounter(n) {
+		return 0
+	}
+
+	idx := fmt.Sprintf("%d:%d", n.Pos().FileIndex(), n.Pos().Line())
+	t := fn.ProfTable
+	return t[idx]
+}
+
+// Get the counter c to the node n in the function fn
+func GetCounterByPos(fn *Func, p src.XPos) Counter {
+	idx := fmt.Sprintf("%d:%d", p.FileIndex(), p.Line())
+	t := fn.ProfTable
+	return t[idx]
 }
 
 // WasmImport stores metadata associated with the //go:wasmimport pragma.

src/cmd/compile/internal/ir/sizeof_test.go

@@ -20,8 +20,8 @@ func TestSizeof(t *testing.T) {
 		_32bit uintptr     // size on 32bit platforms
 		_64bit uintptr     // size on 64bit platforms
 	}{
-		{Func{}, 176, 296},
-		{Name{}, 96, 168},
+		{Func{}, 176, 304},
+		{Name{}, 104, 168},
 	}
 
 	for _, tt := range tests {