cmd/compile: lookup indirect callees from export data for devirtualization

Today, the PGO IR graph only contains entries for ir.Func loaded into
the package. This can include functions from transitive dependencies,
but only if they happen to be referenced by something in the current
package. If they are not referenced, noder never bothers to load them.

This leads to a deficiency in PGO devirtualization: some callee methods
are available in transitive dependencies but do not devirtualize because
they happen to not get loaded from export data.

Resolve this by adding an explicit lookup from export data of callees
mentioned in the profile.

I have chosen to do this during loading of the profile for simplicity:
the PGO IR graph always contains all of the functions we might need.
That said, it isn't strictly necessary. PGO devirtualization could do
the lookup lazily if it decides it actually needs a method. This saves
work at the expense of a bit more complexity, but I've chosen the
simpler approach for now as I measured the cost of this as significantly
less than the rest of PGO loading.

For golang#61577.

Change-Id: Ieafb2a549510587027270ee6b4c3aefd149a901f
Reviewed-on: https://go-review.googlesource.com/c/go/+/497175
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Matthew Dempsky <[email protected]>
Reviewed-by: Cherry Mui <[email protected]>

Author: prattmic

src/cmd/compile/internal/inline/inl.go

@@ -30,7 +30,6 @@ import (
 	"fmt"
 	"go/constant"
 	"internal/goexperiment"
-	"sort"
 	"strconv"
 
 	"cmd/compile/internal/base"
@@ -122,38 +121,18 @@ func pgoInlinePrologue(p *pgo.Profile, funcs []*ir.Func) {
 // comparing with the threshold may not accurately reflect which nodes are
 // considiered hot).
 func hotNodesFromCDF(p *pgo.Profile) (float64, []pgo.NodeMapKey) {
-	nodes := make([]pgo.NodeMapKey, len(p.NodeMap))
-	i := 0
-	for n := range p.NodeMap {
-		nodes[i] = n
-		i++
-	}
-	sort.Slice(nodes, func(i, j int) bool {
-		ni, nj := nodes[i], nodes[j]
-		if wi, wj := p.NodeMap[ni].EWeight, p.NodeMap[nj].EWeight; wi != wj {
-			return wi > wj // want larger weight first
-		}
-		// same weight, order by name/line number
-		if ni.CallerName != nj.CallerName {
-			return ni.CallerName < nj.CallerName
-		}
-		if ni.CalleeName != nj.CalleeName {
-			return ni.CalleeName < nj.CalleeName
-		}
-		return ni.CallSiteOffset < nj.CallSiteOffset
-	})
 	cum := int64(0)
-	for i, n := range nodes {
+	for i, n := range p.NodesByWeight {
 		w := p.NodeMap[n].EWeight
 		cum += w
 		if pgo.WeightInPercentage(cum, p.TotalEdgeWeight) > inlineCDFHotCallSiteThresholdPercent {
 			// nodes[:i+1] to include the very last node that makes it to go over the threshold.
 			// (Say, if the CDF threshold is 50% and one hot node takes 60% of weight, we want to
 			// include that node instead of excluding it.)
-			return pgo.WeightInPercentage(w, p.TotalEdgeWeight), nodes[:i+1]
+			return pgo.WeightInPercentage(w, p.TotalEdgeWeight), p.NodesByWeight[:i+1]
 		}
 	}
-	return 0, nodes
+	return 0, p.NodesByWeight
 }
 
 // InlinePackage finds functions that can be inlined and clones them before walk expands them.

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

@@ -1189,6 +1189,51 @@ func MethodSymSuffix(recv *types.Type, msym *types.Sym, suffix string) *types.Sy
 	return rpkg.LookupBytes(b.Bytes())
 }
 
+// LookupMethodSelector returns the types.Sym of the selector for a method
+// named in local symbol name, as well as the types.Sym of the receiver.
+//
+// TODO(prattmic): this does not attempt to handle method suffixes (wrappers).
+func LookupMethodSelector(pkg *types.Pkg, name string) (typ, meth *types.Sym, err error) {
+	typeName, methName := splitType(name)
+	if typeName == "" {
+		return nil, nil, fmt.Errorf("%s doesn't contain type split", name)
+	}
+
+	if len(typeName) > 3 && typeName[:2] == "(*" && typeName[len(typeName)-1] == ')' {
+		// Symbol name is for a pointer receiver method. We just want
+		// the base type name.
+		typeName = typeName[2 : len(typeName)-1]
+	}
+
+	typ = pkg.Lookup(typeName)
+	meth = pkg.Selector(methName)
+	return typ, meth, nil
+}
+
+// splitType splits a local symbol name into type and method (fn). If this a
+// free function, typ == "".
+//
+// N.B. closures and methods can be ambiguous (e.g., bar.func1). These cases
+// are returned as methods.
+func splitType(name string) (typ, fn string) {
+	// Types are split on the first dot, ignoring everything inside
+	// brackets (instantiation of type parameter, usually including
+	// "go.shape").
+	bracket := 0
+	for i, r := range name {
+		if r == '.' && bracket == 0 {
+			return name[:i], name[i+1:]
+		}
+		if r == '[' {
+			bracket++
+		}
+		if r == ']' {
+			bracket--
+		}
+	}
+	return "", name
+}
+
 // MethodExprName returns the ONAME representing the method
 // referenced by expression n, which must be a method selector,
 // method expression, or method value.

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

@@ -12,6 +12,7 @@ import (
 	"cmd/internal/src"
 	"fmt"
 	"strings"
+	"unicode/utf8"
 )
 
 // A Func corresponds to a single function in a Go program
@@ -312,6 +313,67 @@ func LinkFuncName(f *Func) string {
 	return objabi.PathToPrefix(pkg.Path) + "." + s.Name
 }
 
+// ParseLinkFuncName parsers a symbol name (as returned from LinkFuncName) back
+// to the package path and local symbol name.
+func ParseLinkFuncName(name string) (pkg, sym string, err error) {
+	pkg, sym = splitPkg(name)
+	if pkg == "" {
+		return "", "", fmt.Errorf("no package path in name")
+	}
+
+	pkg, err = objabi.PrefixToPath(pkg) // unescape
+	if err != nil {
+		return "", "", fmt.Errorf("malformed package path: %v", err)
+	}
+
+	return pkg, sym, nil
+}
+
+// Borrowed from x/mod.
+func modPathOK(r rune) bool {
+	if r < utf8.RuneSelf {
+		return r == '-' || r == '.' || r == '_' || r == '~' ||
+			'0' <= r && r <= '9' ||
+			'A' <= r && r <= 'Z' ||
+			'a' <= r && r <= 'z'
+	}
+	return false
+}
+
+func escapedImportPathOK(r rune) bool {
+	return modPathOK(r) || r == '+' || r == '/' || r == '%'
+}
+
+// splitPkg splits the full linker symbol name into package and local symbol
+// name.
+func splitPkg(name string) (pkgpath, sym string) {
+	// package-sym split is at first dot after last the / that comes before
+	// any characters illegal in a package path.
+
+	lastSlashIdx := 0
+	for i, r := range name {
+		// Catches cases like:
+		// * example.foo[sync/atomic.Uint64].
+		// * example%2ecom.foo[sync/atomic.Uint64].
+		//
+		// Note that name is still escaped; unescape occurs after splitPkg.
+		if !escapedImportPathOK(r) {
+			break
+		}
+		if r == '/' {
+			lastSlashIdx = i
+		}
+	}
+	for i := lastSlashIdx; i < len(name); i++ {
+		r := name[i]
+		if r == '.' {
+			return name[:i], name[i+1:]
+		}
+	}
+
+	return "", name
+}
+
 var CurFunc *Func
 
 // WithFunc invokes do with CurFunc and base.Pos set to curfn and

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

@@ -0,0 +1,82 @@
+// Copyright 2023 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.
+
+package ir
+
+import (
+	"testing"
+)
+
+func TestSplitPkg(t *testing.T) {
+	tests := []struct {
+		in  string
+		pkg string
+		sym string
+	}{
+		{
+			in:  "foo.Bar",
+			pkg: "foo",
+			sym: "Bar",
+		},
+		{
+			in:  "foo/bar.Baz",
+			pkg: "foo/bar",
+			sym: "Baz",
+		},
+		{
+			in:  "memeqbody",
+			pkg: "",
+			sym: "memeqbody",
+		},
+		{
+			in:  `example%2ecom.Bar`,
+			pkg: `example%2ecom`,
+			sym: "Bar",
+		},
+		{
+			// Not a real generated symbol name, but easier to catch the general parameter form.
+			in:  `foo.Bar[sync/atomic.Uint64]`,
+			pkg: `foo`,
+			sym: "Bar[sync/atomic.Uint64]",
+		},
+		{
+			in:  `example%2ecom.Bar[sync/atomic.Uint64]`,
+			pkg: `example%2ecom`,
+			sym: "Bar[sync/atomic.Uint64]",
+		},
+		{
+			in:  `gopkg.in/yaml%2ev3.Bar[sync/atomic.Uint64]`,
+			pkg: `gopkg.in/yaml%2ev3`,
+			sym: "Bar[sync/atomic.Uint64]",
+		},
+		{
+			// This one is a real symbol name.
+			in:  `foo.Bar[go.shape.struct { sync/atomic._ sync/atomic.noCopy; sync/atomic._ sync/atomic.align64; sync/atomic.v uint64 }]`,
+			pkg: `foo`,
+			sym: "Bar[go.shape.struct { sync/atomic._ sync/atomic.noCopy; sync/atomic._ sync/atomic.align64; sync/atomic.v uint64 }]",
+		},
+		{
+			in:  `example%2ecom.Bar[go.shape.struct { sync/atomic._ sync/atomic.noCopy; sync/atomic._ sync/atomic.align64; sync/atomic.v uint64 }]`,
+			pkg: `example%2ecom`,
+			sym: "Bar[go.shape.struct { sync/atomic._ sync/atomic.noCopy; sync/atomic._ sync/atomic.align64; sync/atomic.v uint64 }]",
+		},
+		{
+			in:  `gopkg.in/yaml%2ev3.Bar[go.shape.struct { sync/atomic._ sync/atomic.noCopy; sync/atomic._ sync/atomic.align64; sync/atomic.v uint64 }]`,
+			pkg: `gopkg.in/yaml%2ev3`,
+			sym: "Bar[go.shape.struct { sync/atomic._ sync/atomic.noCopy; sync/atomic._ sync/atomic.align64; sync/atomic.v uint64 }]",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.in, func(t *testing.T) {
+			pkg, sym := splitPkg(tc.in)
+			if pkg != tc.pkg {
+				t.Errorf("splitPkg(%q) got pkg %q want %q", tc.in, pkg, tc.pkg)
+			}
+			if sym != tc.sym {
+				t.Errorf("splitPkg(%q) got sym %q want %q", tc.in, sym, tc.sym)
+			}
+		})
+	}
+}

src/cmd/compile/internal/noder/unified.go

@@ -5,6 +5,7 @@
 package noder
 
 import (
+	"fmt"
 	"internal/pkgbits"
 	"io"
 	"runtime"
@@ -14,6 +15,7 @@ import (
 	"cmd/compile/internal/base"
 	"cmd/compile/internal/inline"
 	"cmd/compile/internal/ir"
+	"cmd/compile/internal/pgo"
 	"cmd/compile/internal/typecheck"
 	"cmd/compile/internal/types"
 	"cmd/compile/internal/types2"
@@ -25,6 +27,65 @@ import (
 // later.
 var localPkgReader *pkgReader
 
+// LookupMethodFunc returns the ir.Func for an arbitrary full symbol name if
+// that function exists in the set of available export data.
+//
+// This allows lookup of arbitrary methods that aren't otherwise referenced by
+// the local package and thus haven't been read yet.
+//
+// TODO(prattmic): Does not handle instantiation of generic types. Currently
+// profiles don't contain the original type arguments, so we won't be able to
+// create the runtime dictionaries.
+//
+// TODO(prattmic): Hit rate of this function is usually fairly low, and errors
+// are only used when debug logging is enabled. Consider constructing cheaper
+// errors by default.
+func LookupMethodFunc(fullName string) (*ir.Func, error) {
+	pkgPath, symName, err := ir.ParseLinkFuncName(fullName)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing symbol name %q: %v", fullName, err)
+	}
+
+	pkg, ok := types.PkgMap()[pkgPath]
+	if !ok {
+		return nil, fmt.Errorf("pkg %s doesn't exist in %v", pkgPath, types.PkgMap())
+	}
+
+	// N.B. readPackage creates a Sym for every object in the package to
+	// initialize objReader and importBodyReader, even if the object isn't
+	// read.
+	//
+	// However, objReader is only initialized for top-level objects, so we
+	// must first lookup the type and use that to find the method rather
+	// than looking for the method directly.
+	typ, meth, err := ir.LookupMethodSelector(pkg, symName)
+	if err != nil {
+		return nil, fmt.Errorf("error looking up method symbol %q: %v", symName, err)
+	}
+
+	pri, ok := objReader[typ]
+	if !ok {
+		return nil, fmt.Errorf("type sym %v missing objReader", typ)
+	}
+
+	name := pri.pr.objIdx(pri.idx, nil, nil, false).(*ir.Name)
+	if name.Op() != ir.OTYPE {
+		return nil, fmt.Errorf("type sym %v refers to non-type name: %v", typ, name)
+	}
+	if name.Alias() {
+		return nil, fmt.Errorf("type sym %v refers to alias", typ)
+	}
+
+	for _, m := range name.Type().Methods() {
+		if m.Sym == meth {
+			fn := m.Nname.(*ir.Name).Func
+			return fn, nil
+		}
+	}
+
+	return nil, fmt.Errorf("method %s missing from method set of %v", symName, typ)
+}
+
 // unified constructs the local package's Internal Representation (IR)
 // from its syntax tree (AST).
 //
@@ -69,6 +130,7 @@ var localPkgReader *pkgReader
 func unified(m posMap, noders []*noder) {
 	inline.InlineCall = unifiedInlineCall
 	typecheck.HaveInlineBody = unifiedHaveInlineBody
+	pgo.LookupMethodFunc = LookupMethodFunc
 
 	data := writePkgStub(m, noders)