cmd/preprofile: Implement a tool to preprocess the PGO profile.

It fixes the issue #65220.
It also includes  https://go.dev/cl/557458 from Michael.

Change-Id: Ic6109e1b6a9045459ff4a54dea11cbfe732b01e6
Reviewed-on: https://go-review.googlesource.com/c/go/+/557918
Reviewed-by: Michael Pratt <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Cherry Mui <[email protected]>

Author: jinlin-bayarea

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

@@ -124,7 +124,7 @@ type CmdFlags struct {
 	TraceProfile       string       "help:\"write an execution trace to `file`\""
 	TrimPath           string       "help:\"remove `prefix` from recorded source file paths\""
 	WB                 bool         "help:\"enable write barrier\"" // TODO: remove
-	PgoProfile         string       "help:\"read profile from `file`\""
+	PgoProfile         string       "help:\"read profile or pre-process profile from `file`\""
 	ErrorURL           bool         "help:\"print explanatory URL with error message if applicable\""
 
 	// Configuration derived from flags; not a flag itself.

src/cmd/compile/internal/pgo/irgraph.go

@@ -41,16 +41,19 @@
 package pgo
 
 import (
+	"bufio"
 	"cmd/compile/internal/base"
 	"cmd/compile/internal/ir"
-	"cmd/compile/internal/pgo/internal/graph"
 	"cmd/compile/internal/typecheck"
 	"cmd/compile/internal/types"
 	"errors"
 	"fmt"
 	"internal/profile"
+	"io/ioutil"
 	"os"
 	"sort"
+	"strconv"
+	"strings"
 )
 
 // IRGraph is a call graph with nodes pointing to IRs of functions and edges
@@ -105,6 +108,7 @@ type NamedCallEdge struct {
 	CallerName     string
 	CalleeName     string
 	CallSiteOffset int // Line offset from function start line.
+	CallStartLine  int // Start line of the function. Can be 0 which means missing.
 }
 
 // NamedEdgeMap contains all unique call edges in the profile and their
@@ -139,8 +143,47 @@ type Profile struct {
 	WeightedCG *IRGraph
 }
 
-// New generates a profile-graph from the profile.
+var wantHdr = "GO PREPROFILE V1\n"
+
+func isPreProfileFile(filename string) (bool, error) {
+	content, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return false, err
+	}
+
+	/* check the header */
+	fileContent := string(content)
+	if strings.HasPrefix(fileContent, wantHdr) {
+		return true, nil
+	}
+	return false, nil
+}
+
+// New generates a profile-graph from the profile or pre-processed profile.
 func New(profileFile string) (*Profile, error) {
+	var profile *Profile
+	var err error
+	isPreProf, err := isPreProfileFile(profileFile)
+	if err != nil {
+		return nil, fmt.Errorf("error opening profile: %w", err)
+	}
+	if !isPreProf {
+		profile, err = processProto(profileFile)
+		if err != nil {
+			return nil, fmt.Errorf("error processing pprof PGO profile: %w", err)
+		}
+	} else {
+		profile, err = processPreprof(profileFile)
+		if err != nil {
+			return nil, fmt.Errorf("error processing preprocessed PGO profile: %w", err)
+		}
+	}
+	return profile, nil
+
+}
+
+// processProto generates a profile-graph from the profile.
+func processProto(profileFile string) (*Profile, error) {
 	f, err := os.Open(profileFile)
 	if err != nil {
 		return nil, fmt.Errorf("error opening profile: %w", err)
@@ -175,7 +218,7 @@ func New(profileFile string) (*Profile, error) {
 		return nil, fmt.Errorf(`profile does not contain a sample index with value/type "samples/count" or cpu/nanoseconds"`)
 	}
 
-	g := graph.NewGraph(p, &graph.Options{
+	g := profile.NewGraph(p, &profile.Options{
 		SampleValue: func(v []int64) int64 { return v[valueIndex] },
 	})
 
@@ -198,11 +241,134 @@ func New(profileFile string) (*Profile, error) {
 	}, nil
 }
 
+// processPreprof generates a profile-graph from the pre-procesed profile.
+func processPreprof(preprofileFile string) (*Profile, error) {
+	namedEdgeMap, totalWeight, err := createNamedEdgeMapFromPreprocess(preprofileFile)
+	if err != nil {
+		return nil, err
+	}
+
+	if totalWeight == 0 {
+		return nil, nil // accept but ignore profile with no samples.
+	}
+
+	// Create package-level call graph with weights from profile and IR.
+	wg := createIRGraph(namedEdgeMap)
+
+	return &Profile{
+		TotalWeight:  totalWeight,
+		NamedEdgeMap: namedEdgeMap,
+		WeightedCG:   wg,
+	}, nil
+}
+
+func postProcessNamedEdgeMap(weight map[NamedCallEdge]int64, weightVal int64) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
+	if weightVal == 0 {
+		return NamedEdgeMap{}, 0, nil // accept but ignore profile with no samples.
+	}
+	byWeight := make([]NamedCallEdge, 0, len(weight))
+	for namedEdge := range weight {
+		byWeight = append(byWeight, namedEdge)
+	}
+	sort.Slice(byWeight, func(i, j int) bool {
+		ei, ej := byWeight[i], byWeight[j]
+		if wi, wj := weight[ei], weight[ej]; wi != wj {
+			return wi > wj // want larger weight first
+		}
+		// same weight, order by name/line number
+		if ei.CallerName != ej.CallerName {
+			return ei.CallerName < ej.CallerName
+		}
+		if ei.CalleeName != ej.CalleeName {
+			return ei.CalleeName < ej.CalleeName
+		}
+		return ei.CallSiteOffset < ej.CallSiteOffset
+	})
+
+	edgeMap = NamedEdgeMap{
+		Weight:   weight,
+		ByWeight: byWeight,
+	}
+
+	totalWeight = weightVal
+
+	return edgeMap, totalWeight, nil
+}
+
+// restore NodeMap information from a preprocessed profile.
+// The reader can refer to the format of preprocessed profile in cmd/preprofile/main.go.
+func createNamedEdgeMapFromPreprocess(preprofileFile string) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
+	readFile, err := os.Open(preprofileFile)
+	if err != nil {
+		return NamedEdgeMap{}, 0, fmt.Errorf("error opening preprocessed profile: %w", err)
+	}
+	defer readFile.Close()
+
+	fileScanner := bufio.NewScanner(readFile)
+	fileScanner.Split(bufio.ScanLines)
+	weight := make(map[NamedCallEdge]int64)
+
+	if !fileScanner.Scan() {
+		if err := fileScanner.Err(); err != nil {
+			return NamedEdgeMap{}, 0, fmt.Errorf("error reading preprocessed profile: %w", err)
+		}
+		return NamedEdgeMap{}, 0, fmt.Errorf("preprocessed profile missing header")
+	}
+	if gotHdr := fileScanner.Text() + "\n"; gotHdr != wantHdr {
+		return NamedEdgeMap{}, 0, fmt.Errorf("preprocessed profile malformed header; got %q want %q", gotHdr, wantHdr)
+	}
+
+	for fileScanner.Scan() {
+		readStr := fileScanner.Text()
+
+		callerName := readStr
+
+		if !fileScanner.Scan() {
+			if err := fileScanner.Err(); err != nil {
+				return NamedEdgeMap{}, 0, fmt.Errorf("error reading preprocessed profile: %w", err)
+			}
+			return NamedEdgeMap{}, 0, fmt.Errorf("preprocessed profile entry missing callee")
+		}
+		calleeName := fileScanner.Text()
+
+		if !fileScanner.Scan() {
+			if err := fileScanner.Err(); err != nil {
+				return NamedEdgeMap{}, 0, fmt.Errorf("error reading preprocessed profile: %w", err)
+			}
+			return NamedEdgeMap{}, 0, fmt.Errorf("preprocessed profile entry missing weight")
+		}
+		readStr = fileScanner.Text()
+
+		split := strings.Split(readStr, " ")
+
+		if len(split) != 5 {
+			return NamedEdgeMap{}, 0, fmt.Errorf("preprocessed profile entry got %v want 5 fields", split)
+		}
+
+		co, _ := strconv.Atoi(split[0])
+		cs, _ := strconv.Atoi(split[1])
+
+		namedEdge := NamedCallEdge{
+			CallerName:     callerName,
+			CallSiteOffset: co - cs,
+		}
+
+		namedEdge.CalleeName = calleeName
+		EWeight, _ := strconv.ParseInt(split[4], 10, 64)
+
+		weight[namedEdge] += EWeight
+		totalWeight += EWeight
+	}
+
+	return postProcessNamedEdgeMap(weight, totalWeight)
+
+}
+
 // createNamedEdgeMap builds a map of callsite-callee edge weights from the
 // profile-graph.
 //
 // Caller should ignore the profile if totalWeight == 0.
-func createNamedEdgeMap(g *graph.Graph) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
+func createNamedEdgeMap(g *profile.Graph) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
 	seenStartLine := false
 
 	// Process graph and build various node and edge maps which will
@@ -226,42 +392,13 @@ func createNamedEdgeMap(g *graph.Graph) (edgeMap NamedEdgeMap, totalWeight int64
 		}
 	}
 
-	if totalWeight == 0 {
-		return NamedEdgeMap{}, 0, nil // accept but ignore profile with no samples.
-	}
-
 	if !seenStartLine {
 		// TODO(prattmic): If Function.start_line is missing we could
 		// fall back to using absolute line numbers, which is better
 		// than nothing.
 		return NamedEdgeMap{}, 0, fmt.Errorf("profile missing Function.start_line data (Go version of profiled application too old? Go 1.20+ automatically adds this to profiles)")
 	}
-
-	byWeight := make([]NamedCallEdge, 0, len(weight))
-	for namedEdge := range weight {
-		byWeight = append(byWeight, namedEdge)
-	}
-	sort.Slice(byWeight, func(i, j int) bool {
-		ei, ej := byWeight[i], byWeight[j]
-		if wi, wj := weight[ei], weight[ej]; wi != wj {
-			return wi > wj // want larger weight first
-		}
-		// same weight, order by name/line number
-		if ei.CallerName != ej.CallerName {
-			return ei.CallerName < ej.CallerName
-		}
-		if ei.CalleeName != ej.CalleeName {
-			return ei.CalleeName < ej.CalleeName
-		}
-		return ei.CallSiteOffset < ej.CallSiteOffset
-	})
-
-	edgeMap = NamedEdgeMap{
-		Weight:   weight,
-		ByWeight: byWeight,
-	}
-
-	return edgeMap, totalWeight, nil
+	return postProcessNamedEdgeMap(weight, totalWeight)
 }
 
 // initializeIRGraph builds the IRGraph by visiting all the ir.Func in decl list

src/cmd/compile/internal/test/pgo_inl_test.go

@@ -43,7 +43,12 @@ go 1.19
 }
 
 // testPGOIntendedInlining tests that specific functions are inlined.
-func testPGOIntendedInlining(t *testing.T, dir string) {
+func testPGOIntendedInlining(t *testing.T, dir string, preprocessed ...bool) {
+	defaultPGOPackValue := false
+	if len(preprocessed) > 0 {
+		defaultPGOPackValue = preprocessed[0]
+	}
+
 	testenv.MustHaveGoRun(t)
 	t.Parallel()
 
@@ -86,7 +91,12 @@ func testPGOIntendedInlining(t *testing.T, dir string) {
 
 	// Build the test with the profile. Use a smaller threshold to test.
 	// TODO: maybe adjust the test to work with default threshold.
-	pprof := filepath.Join(dir, "inline_hot.pprof")
+	var pprof string
+	if defaultPGOPackValue == false {
+		pprof = filepath.Join(dir, "inline_hot.pprof")
+	} else {
+		pprof = filepath.Join(dir, "inline_hot.pprof.node_map")
+	}
 	gcflag := fmt.Sprintf("-m -m -pgoprofile=%s -d=pgoinlinebudget=160,pgoinlinecdfthreshold=90", pprof)
 	out := buildPGOInliningTest(t, dir, gcflag)
 
@@ -164,6 +174,27 @@ func TestPGOIntendedInlining(t *testing.T) {
 	testPGOIntendedInlining(t, dir)
 }
 
+// TestPGOIntendedInlining tests that specific functions are inlined when PGO
+// is applied to the exact source that was profiled.
+func TestPGOPreprocessInlining(t *testing.T) {
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Fatalf("error getting wd: %v", err)
+	}
+	srcDir := filepath.Join(wd, "testdata/pgo/inline")
+
+	// Copy the module to a scratch location so we can add a go.mod.
+	dir := t.TempDir()
+
+	for _, file := range []string{"inline_hot.go", "inline_hot_test.go", "inline_hot.pprof.node_map"} {
+		if err := copyFile(filepath.Join(dir, file), filepath.Join(srcDir, file)); err != nil {
+			t.Fatalf("error copying %s: %v", file, err)
+		}
+	}
+
+	testPGOIntendedInlining(t, dir, true)
+}
+
 // TestPGOIntendedInlining tests that specific functions are inlined when PGO
 // is applied to the modified source.
 func TestPGOIntendedInliningShiftedLines(t *testing.T) {

src/cmd/compile/internal/test/testdata/pgo/inline/inline_hot.pprof.node_map

@@ -0,0 +1,13 @@
+GO PREPROFILE V1
+example.com/pgo/inline.benchmarkB
+example.com/pgo/inline.A
+18 17 0 1 1
+example.com/pgo/inline.(*BS).NS
+example.com/pgo/inline.T
+13 53 124 129 2
+example.com/pgo/inline.(*BS).NS
+example.com/pgo/inline.T
+8 53 124 129 3
+example.com/pgo/inline.A
+example.com/pgo/inline.(*BS).NS
+7 74 1 130 129