proposal: testing: extend the compiler's treatment of B.Loop to PB.Next #81983

Background

b.Loop was added in #61515. One of the advantages of b.Loop-based benchmark loops is that the compiler recognizes them specially and doesn't perform certain optimizations within these loops:

Within the body of a "for b.Loop() { ... }" loop, arguments to and results from function calls and assigned variables within the loop are kept alive, preventing the compiler from fully optimizing away the loop body. Currently, this is implemented as a compiler transformation that wraps such variables with a runtime.KeepAlive intrinsic call. This applies only to statements syntactically between the curly braces of the loop, and the loop condition must be written exactly as "b.Loop()".

We arrived at this behavior through iterative refinement across several releases (see issues #73137 and #77339), but I think we're in a pretty good place now.

However, none of these compiler specializations apply to benchmarks that use b.RunParallel and PB.Next. All of the same issues that necessitated the compiler special-cases that were discussed in the issues I linked above also apply to parallel benchmarks.

Demo

package pbnext

import (
	"math"
	"testing"
)

func expensive(x float64) float64 {
	x = math.Sqrt(x + 1)
	x = math.Sqrt(x + 1)
	x = math.Sqrt(x + 1)
	x = math.Sqrt(x + 1)
	x = math.Sqrt(x + 1)
	x = math.Sqrt(x + 1)
	return math.Sqrt(x + 1)
}

func BenchmarkLoop(b *testing.B) {
	for i := 0; b.Loop(); i++ {
		expensive(float64(i))
	}
}

func BenchmarkParallel(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		for i := 0; pb.Next(); i++ {
			expensive(float64(i))
		}
	})
}

When I run this, the entire expensive computation gets eliminated from BenchmarkParallel:

▶ go test -bench . -cpu=1
goos: linux
goarch: amd64
pkg: github.com/cespare/misc/pbnext
cpu: AMD Ryzen 9 3900X 12-Core Processor
BenchmarkLoop           79889703                14.52 ns/op
BenchmarkParallel       1000000000               0.4373 ns/op

Proposal

I propose that we apply the exact same rules to pb.Next-based loops as we apply to b.Loop-based loops. Just like b.Loop, it would only apply if the loop condition is exactly pb.Next(), and it would only apply within the curly braces of the loop.

This was also mentioned by @AlekSi in #73137, but he never got a response.

Alternatives

Another proposal from the same time as b.Loop was #61179, which proposes testing.Keep. That one seems to have stalled out at this point (the need for it is much less with the b.Loop addition). But if we added that, then at least people could more easily work around the issue by inserting their own testing.Keep calls.

/cc some bloopers: @thepudds, @aclements, @JunyangShao