cmd/compile: initial function value devirtualization

Today, PGO-based devirtualization only applies to interface calls. This
CL extends initial support to function values (i.e., function/closure
pointers passed as arguments or stored in a struct).

This CL is a minimal implementation with several limitations.

* Export data lookup of function value callees not implemented
  (equivalent of CL 497175; done in CL 540258).
* Callees must be standard static functions. Callees that are closures
  (requiring closure context) are not supported.

For #61577.

Change-Id: I7d328859035249e176294cd0d9885b2d08c853f6
Reviewed-on: https://go-review.googlesource.com/c/go/+/539699
Reviewed-by: Matthew Dempsky <mdempsky@google.com>
Reviewed-by: Cherry Mui <cherryyz@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

5 files changed, 336 insertions, 12 deletions

diff --git a/src/cmd/compile/internal/test/pgo_devirtualize_test.go b/src/cmd/compile/internal/test/pgo_devirtualize_test.go
index fbee8dedfd..3e264a3f41 100644
--- a/src/cmd/compile/internal/test/pgo_devirtualize_test.go
+++ b/src/cmd/compile/internal/test/pgo_devirtualize_test.go
@@ -29,11 +29,21 @@ go 1.19
 		t.Fatalf("error writing go.mod: %v", err)
 	}
 
+	// Run the test without PGO to ensure that the test assertions are
+	// correct even in the non-optimized version.
+	cmd := testenv.CleanCmdEnv(testenv.Command(t, testenv.GoToolPath(t), "test", "."))
+	cmd.Dir = dir
+	b, err := cmd.CombinedOutput()
+	t.Logf("Test without PGO:\n%s", b)
+	if err != nil {
+		t.Fatalf("Test failed without PGO: %v", err)
+	}
+
 	// Build the test with the profile.
 	pprof := filepath.Join(dir, "devirt.pprof")
 	gcflag := fmt.Sprintf("-gcflags=-m=2 -pgoprofile=%s -d=pgodebug=3", pprof)
 	out := filepath.Join(dir, "test.exe")
-	cmd := testenv.CleanCmdEnv(testenv.Command(t, testenv.GoToolPath(t), "build", "-o", out, gcflag, "."))
+	cmd = testenv.CleanCmdEnv(testenv.Command(t, testenv.GoToolPath(t), "test", "-o", out, gcflag, "."))
 	cmd.Dir = dir
 
 	pr, pw, err := os.Pipe()
@@ -56,19 +66,50 @@ go 1.19
 	}
 
 	want := []devirtualization{
+		// ExerciseIface
 		{
-			pos:    "./devirt.go:66:21",
+			pos:    "./devirt.go:101:20",
 			callee: "mult.Mult.Multiply",
 		},
 		{
-			pos:    "./devirt.go:66:31",
+			pos:    "./devirt.go:101:39",
 			callee: "Add.Add",
 		},
+		// ExerciseFuncConcrete
+		{
+			pos:    "./devirt.go:178:18",
+			callee: "AddFn",
+		},
+		// TODO(prattmic): Export data lookup for function value callees not implemented.
+		//{
+		//	pos:    "./devirt.go:179:15",
+		//	callee: "mult.MultFn",
+		//},
+		// ExerciseFuncField
+		{
+			pos:    "./devirt.go:218:13",
+			callee: "AddFn",
+		},
+		// TODO(prattmic): Export data lookup for function value callees not implemented.
+		//{
+		//	pos:    "./devirt.go:219:19",
+		//	callee: "mult.MultFn",
+		//},
+		// ExerciseFuncClosure
+		// TODO(prattmic): Closure callees not implemented.
+		//{
+		//	pos:    "./devirt.go:266:9",
+		//	callee: "AddClosure.func1",
+		//},
+		//{
+		//	pos:    "./devirt.go:267:15",
+		//	callee: "mult.MultClosure.func1",
+		//},
 	}
 
 	got := make(map[devirtualization]struct{})
 
-	devirtualizedLine := regexp.MustCompile(`(.*): PGO devirtualizing .* to (.*)`)
+	devirtualizedLine := regexp.MustCompile(`(.*): PGO devirtualizing \w+ call .* to (.*)`)
 
 	scanner := bufio.NewScanner(pr)
 	for scanner.Scan() {
@@ -102,6 +143,15 @@ go 1.19
 		}
 		t.Errorf("devirtualization %v missing; got %v", w, got)
 	}
+
+	// Run test with PGO to ensure the assertions are still true.
+	cmd = testenv.CleanCmdEnv(testenv.Command(t, out))
+	cmd.Dir = dir
+	b, err = cmd.CombinedOutput()
+	t.Logf("Test with PGO:\n%s", b)
+	if err != nil {
+		t.Fatalf("Test failed without PGO: %v", err)
+	}
 }
 
 // TestPGODevirtualize tests that specific functions are devirtualized when PGO
diff --git a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.go b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.go
index 4748e19e10..63de3d3c3f 100644
--- a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.go
+++ b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.go
@@ -16,7 +16,11 @@ package devirt
 //
 // Dots in the last package path component are escaped in symbol names. Use one
 // to ensure the escaping doesn't break lookup.
-import "example.com/pgo/devirtualize/mult.pkg"
+import (
+	"fmt"
+
+	"example.com/pgo/devirtualize/mult.pkg"
+)
 
 var sink int
 
@@ -42,15 +46,46 @@ func (Sub) Add(a, b int) int {
 	return a - b
 }
 
-// Exercise calls mostly a1 and m1.
+// ExerciseIface calls mostly a1 and m1.
 //
 //go:noinline
-func Exercise(iter int, a1, a2 Adder, m1, m2 mult.Multiplier) {
+func ExerciseIface(iter int, a1, a2 Adder, m1, m2 mult.Multiplier) int {
+	// The call below must evaluate selectA() to determine the receiver to
+	// use. This should happen exactly once per iteration. Assert that is
+	// the case to ensure the IR manipulation does not result in over- or
+	// under-evaluation.
+	selectI := 0
+	selectA := func(gotI int) Adder {
+		if gotI != selectI {
+			panic(fmt.Sprintf("selectA not called once per iteration; got i %d want %d", gotI, selectI))
+		}
+		selectI++
+
+		if gotI%10 == 0 {
+			return a2
+		}
+		return a1
+	}
+	oneI := 0
+	one := func(gotI int) int {
+		if gotI != oneI {
+			panic(fmt.Sprintf("one not called once per iteration; got i %d want %d", gotI, oneI))
+		}
+		oneI++
+
+		// The function value must be evaluated before arguments, so
+		// selectI must have been incremented already.
+		if selectI != oneI {
+			panic(fmt.Sprintf("selectA not called before not called before one; got i %d want %d", selectI, oneI))
+		}
+
+		return 1
+	}
+
+	val := 0
 	for i := 0; i < iter; i++ {
-		a := a1
 		m := m1
 		if i%10 == 0 {
-			a = a2
 			m = m2
 		}
 
@@ -63,6 +98,173 @@ func Exercise(iter int, a1, a2 Adder, m1, m2 mult.Multiplier) {
 		// If they were not mutually exclusive (for example, two Add
 		// calls), then we could not definitively select the correct
 		// callee.
-		sink += m.Multiply(42, a.Add(1, 2))
+		val += m.Multiply(42, selectA(i).Add(one(i), 2))
+	}
+	return val
+}
+
+type AddFunc func(int, int) int
+
+func AddFn(a, b int) int {
+	for i := 0; i < 1000; i++ {
+		sink++
+	}
+	return a + b
+}
+
+func SubFn(a, b int) int {
+	for i := 0; i < 1000; i++ {
+		sink++
+	}
+	return a - b
+}
+
+// ExerciseFuncConcrete calls mostly a1 and m1.
+//
+//go:noinline
+func ExerciseFuncConcrete(iter int, a1, a2 AddFunc, m1, m2 mult.MultFunc) int {
+	// The call below must evaluate selectA() to determine the function to
+	// call. This should happen exactly once per iteration. Assert that is
+	// the case to ensure the IR manipulation does not result in over- or
+	// under-evaluation.
+	selectI := 0
+	selectA := func(gotI int) AddFunc {
+		if gotI != selectI {
+			panic(fmt.Sprintf("selectA not called once per iteration; got i %d want %d", gotI, selectI))
+		}
+		selectI++
+
+		if gotI%10 == 0 {
+			return a2
+		}
+		return a1
+	}
+	oneI := 0
+	one := func(gotI int) int {
+		if gotI != oneI {
+			panic(fmt.Sprintf("one not called once per iteration; got i %d want %d", gotI, oneI))
+		}
+		oneI++
+
+		// The function value must be evaluated before arguments, so
+		// selectI must have been incremented already.
+		if selectI != oneI {
+			panic(fmt.Sprintf("selectA not called before not called before one; got i %d want %d", selectI, oneI))
+		}
+
+		return 1
+	}
+
+	val := 0
+	for i := 0; i < iter; i++ {
+		m := m1
+		if i%10 == 0 {
+			m = m2
+		}
+
+		// N.B. Profiles only distinguish calls on a per-line level,
+		// making the two calls ambiguous. However because the
+		// function types are mutually exclusive, devirtualization can
+		// still select the correct callee for each.
+		//
+		// If they were not mutually exclusive (for example, two
+		// AddFunc calls), then we could not definitively select the
+		// correct callee.
+		//
+		// TODO(prattmic): Export data lookup for function value
+		// callees not implemented, meaning the type is unavailable.
+		//sink += int(m(42, int64(a(1, 2))))
+
+		v := selectA(i)(one(i), 2)
+		val += int(m(42, int64(v)))
+	}
+	return val
+}
+
+// ExerciseFuncField calls mostly a1 and m1.
+//
+// This is a simplified version of ExerciseFuncConcrete, but accessing the
+// function values via a struct field.
+//
+//go:noinline
+func ExerciseFuncField(iter int, a1, a2 AddFunc, m1, m2 mult.MultFunc) int {
+	ops := struct {
+		a AddFunc
+		m mult.MultFunc
+	}{}
+
+	val := 0
+	for i := 0; i < iter; i++ {
+		ops.a = a1
+		ops.m = m1
+		if i%10 == 0 {
+			ops.a = a2
+			ops.m = m2
+		}
+
+		// N.B. Profiles only distinguish calls on a per-line level,
+		// making the two calls ambiguous. However because the
+		// function types are mutually exclusive, devirtualization can
+		// still select the correct callee for each.
+		//
+		// If they were not mutually exclusive (for example, two
+		// AddFunc calls), then we could not definitively select the
+		// correct callee.
+		//
+		// TODO(prattmic): Export data lookup for function value
+		// callees not implemented, meaning the type is unavailable.
+		//sink += int(ops.m(42, int64(ops.a(1, 2))))
+
+		v := ops.a(1, 2)
+		val += int(ops.m(42, int64(v)))
+	}
+	return val
+}
+
+//go:noinline
+func AddClosure() AddFunc {
+	// Implicit closure by capturing the receiver.
+	var a Add
+	return a.Add
+}
+
+//go:noinline
+func SubClosure() AddFunc {
+	var s Sub
+	return s.Add
+}
+
+// ExerciseFuncClosure calls mostly a1 and m1.
+//
+// This is a simplified version of ExerciseFuncConcrete, but we need two
+// distinct call sites to test two different types of function values.
+//
+//go:noinline
+func ExerciseFuncClosure(iter int, a1, a2 AddFunc, m1, m2 mult.MultFunc) int {
+	val := 0
+	for i := 0; i < iter; i++ {
+		a := a1
+		m := m1
+		if i%10 == 0 {
+			a = a2
+			m = m2
+		}
+
+		// N.B. Profiles only distinguish calls on a per-line level,
+		// making the two calls ambiguous. However because the
+		// function types are mutually exclusive, devirtualization can
+		// still select the correct callee for each.
+		//
+		// If they were not mutually exclusive (for example, two
+		// AddFunc calls), then we could not definitively select the
+		// correct callee.
+		//
+		// TODO(prattmic): Export data lookup for function value
+		// callees not implemented, meaning the type is unavailable.
+		//sink += int(m(42, int64(a(1, 2))))
+
+		v := a(1, 2)
+		val += int(m(42, int64(v)))
 	}
+	return val
 }
diff --git a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.pprof b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.pprof
index 87e7b62736..de064582ff 100644
--- a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.pprof
+++ b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt.pprof
diff --git a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt_test.go b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt_test.go
index ef637a876b..59b565d77f 100644
--- a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt_test.go
+++ b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/devirt_test.go
@@ -17,7 +17,7 @@ import (
 	"example.com/pgo/devirtualize/mult.pkg"
 )
 
-func BenchmarkDevirt(b *testing.B) {
+func BenchmarkDevirtIface(b *testing.B) {
 	var (
 		a1 Add
 		a2 Sub
@@ -25,5 +25,49 @@ func BenchmarkDevirt(b *testing.B) {
 		m2 mult.NegMult
 	)
 
-	Exercise(b.N, a1, a2, m1, m2)
+	ExerciseIface(b.N, a1, a2, m1, m2)
+}
+
+// Verify that devirtualization doesn't result in calls or side effects applying more than once.
+func TestDevirtIface(t *testing.T) {
+	var (
+		a1 Add
+		a2 Sub
+		m1 mult.Mult
+		m2 mult.NegMult
+	)
+
+	if v := ExerciseIface(10, a1, a2, m1, m2); v != 1176 {
+		t.Errorf("ExerciseIface(10) got %d want 1176", v)
+	}
+}
+
+func BenchmarkDevirtFuncConcrete(b *testing.B) {
+	ExerciseFuncConcrete(b.N, AddFn, SubFn, mult.MultFn, mult.NegMultFn)
+}
+
+func TestDevirtFuncConcrete(t *testing.T) {
+	if v := ExerciseFuncConcrete(10, AddFn, SubFn, mult.MultFn, mult.NegMultFn); v != 1176 {
+		t.Errorf("ExerciseFuncConcrete(10) got %d want 1176", v)
+	}
+}
+
+func BenchmarkDevirtFuncField(b *testing.B) {
+	ExerciseFuncField(b.N, AddFn, SubFn, mult.MultFn, mult.NegMultFn)
+}
+
+func TestDevirtFuncField(t *testing.T) {
+	if v := ExerciseFuncField(10, AddFn, SubFn, mult.MultFn, mult.NegMultFn); v != 1176 {
+		t.Errorf("ExerciseFuncField(10) got %d want 1176", v)
+	}
+}
+
+func BenchmarkDevirtFuncClosure(b *testing.B) {
+	ExerciseFuncClosure(b.N, AddClosure(), SubClosure(), mult.MultClosure(), mult.NegMultClosure())
+}
+
+func TestDevirtFuncClosure(t *testing.T) {
+	if v := ExerciseFuncClosure(10, AddClosure(), SubClosure(), mult.MultClosure(), mult.NegMultClosure()); v != 1176 {
+		t.Errorf("ExerciseFuncClosure(10) got %d want 1176", v)
+	}
 }
diff --git a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/mult.pkg/mult.go b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/mult.pkg/mult.go
index 8a026a52f5..64f405ff9e 100644
--- a/src/cmd/compile/internal/test/testdata/pgo/devirtualize/mult.pkg/mult.go
+++ b/src/cmd/compile/internal/test/testdata/pgo/devirtualize/mult.pkg/mult.go
@@ -30,3 +30,31 @@ func (NegMult) Multiply(a, b int) int {
 	}
 	return -1 * a * b
 }
+
+// N.B. Different types than AddFunc to test intra-line disambiguation.
+type MultFunc func(int64, int64) int64
+
+func MultFn(a, b int64) int64 {
+	return a * b
+}
+
+func NegMultFn(a, b int64) int64 {
+	return -1 * a * b
+}
+
+//go:noinline
+func MultClosure() MultFunc {
+	// Explicit closure to differentiate from AddClosure.
+	c := 1
+	return func(a, b int64) int64 {
+		return a * b * int64(c)
+	}
+}
+
+//go:noinline
+func NegMultClosure() MultFunc {
+	c := 1
+	return func(a, b int64) int64 {
+		return -1 * a * b * int64(c)
+	}
+}