cmd/compile/internal/escape: speed up analyzing some functions with many closures

Escape analysis examines functions in batches. In some cases, closures
are in the same batch as their parent function, but in other cases,
the closures are in different batches. This can mean the per-batch
ir.ReassignOracle cache is not as effective.

For example, #74615 has 4,000 closures in a single function that
are all in different batches. For that example, these caches
had an ~80% hit rate.

This CL makes the ir.ReassignOracle cache more broadly scoped, instead
of per batch.

This speeds up escape analysis when a function has many closures
that end up in different batches, including this resolves #74615.
For that example, this cache now has a ~100% hit rate.

In addition, in (*batch).rewriteWithLiterals, we also slightly delay
checking the ir.ReassignOracle cache, which is more natural to do now
compared to when rewriteWithLiterals was first merged. This means we can
avoid consulting or populating the cache in more cases. (We also leave
a new type-related TODO there. If we were to also implement that TODO, a
quick test suggests we could independently resolve the specific example
in #74615 even without making the cache more broadly scoped,
though other conceivable examples would not be helped; the scoping of
the cache is the more fundamental improvement.)

If we look at cumulative time spent via pprof for the #74615 example
using this CL, the work of ir.ReassignOracle escape analysis
cache now typically shows zero cpu samples.

This CL passes "go build -toolexec 'toolstash -cmp' -a std cmd".

Fixes #74615

Change-Id: I3c17c527fbb546ffb8a4fa52cd61e41ff3cdb869
Reviewed-on: https://go-review.googlesource.com/c/go/+/688075
Auto-Submit: Keith Randall <[email protected]>
Reviewed-by: Cherry Mui <[email protected]>
Reviewed-by: Keith Randall <[email protected]>
Reviewed-by: Keith Randall <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: thepudds

src/cmd/compile/internal/escape/escape.go

@@ -122,17 +122,24 @@ type escape struct {
 }
 
 func Funcs(all []*ir.Func) {
-	ir.VisitFuncsBottomUp(all, Batch)
+	// Make a cache of ir.ReassignOracles. The cache is lazily populated.
+	// TODO(thepudds): consider adding a field on ir.Func instead. We might also be able
+	// to use that field elsewhere, like in walk. See discussion in https://go.dev/cl/688075.
+	reassignOracles := make(map[*ir.Func]*ir.ReassignOracle)
+
+	ir.VisitFuncsBottomUp(all, func(list []*ir.Func, recursive bool) {
+		Batch(list, reassignOracles)
+	})
 }
 
 // Batch performs escape analysis on a minimal batch of
 // functions.
-func Batch(fns []*ir.Func, recursive bool) {
+func Batch(fns []*ir.Func, reassignOracles map[*ir.Func]*ir.ReassignOracle) {
 	var b batch
 	b.heapLoc.attrs = attrEscapes | attrPersists | attrMutates | attrCalls
 	b.mutatorLoc.attrs = attrMutates
 	b.calleeLoc.attrs = attrCalls
-	b.reassignOracles = make(map[*ir.Func]*ir.ReassignOracle)
+	b.reassignOracles = reassignOracles
 
 	// Construct data-flow graph from syntax trees.
 	for _, fn := range fns {
@@ -531,15 +538,6 @@ func (b *batch) rewriteWithLiterals(n ir.Node, fn *ir.Func) {
 	if n == nil || fn == nil {
 		return
 	}
-	if n.Op() != ir.OMAKESLICE && n.Op() != ir.OCONVIFACE {
-		return
-	}
-
-	// Look up a cached ReassignOracle for the function, lazily computing one if needed.
-	ro := b.reassignOracle(fn)
-	if ro == nil {
-		base.Fatalf("no ReassignOracle for function %v with closure parent %v", fn, fn.ClosureParent)
-	}
 
 	assignTemp := func(n ir.Node, init *ir.Nodes) {
 		// Preserve any side effects of n by assigning it to an otherwise unused temp.
@@ -561,6 +559,11 @@ func (b *batch) rewriteWithLiterals(n ir.Node, fn *ir.Func) {
 		}
 
 		if (*r).Op() != ir.OLITERAL {
+			// Look up a cached ReassignOracle for the function, lazily computing one if needed.
+			ro := b.reassignOracle(fn)
+			if ro == nil {
+				base.Fatalf("no ReassignOracle for function %v with closure parent %v", fn, fn.ClosureParent)
+			}
 			if s := ro.StaticValue(*r); s.Op() == ir.OLITERAL {
 				lit, ok := s.(*ir.BasicLit)
 				if !ok || lit.Val().Kind() != constant.Int {
@@ -582,6 +585,12 @@ func (b *batch) rewriteWithLiterals(n ir.Node, fn *ir.Func) {
 		// a literal to avoid heap allocating the underlying interface value.
 		conv := n.(*ir.ConvExpr)
 		if conv.X.Op() != ir.OLITERAL && !conv.X.Type().IsInterface() {
+			// TODO(thepudds): likely could avoid some work by tightening the check of conv.X's type.
+			// Look up a cached ReassignOracle for the function, lazily computing one if needed.
+			ro := b.reassignOracle(fn)
+			if ro == nil {
+				base.Fatalf("no ReassignOracle for function %v with closure parent %v", fn, fn.ClosureParent)
+			}
 			v := ro.StaticValue(conv.X)
 			if v != nil && v.Op() == ir.OLITERAL && ir.ValidTypeForConst(conv.X.Type(), v.Val()) {
 				if !base.LiteralAllocHash.MatchPos(n.Pos(), nil) {