cmd/compile: missing escape analysis on generic func impls

[rsc]

Here is a simple program in two packages, checked into rsc.io/tmp/escbug.

package p

import (
	"runtime"

	"rsc.io/tmp/escbug/q"
)

func P() {
	var buf [24]byte
	runtime.KeepAlive(q.Q(buf[:0], 1))
}

package q

func nop(x int) int { return x }

func Q(dst []byte, x int64) []byte {
	return q(dst, x)
}

func q[T int32 | int64](dst []byte, x T) []byte {
	return append(dst, byte(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(nop(int(x))))))))))))))))))))))))
}

p.P calls q.Q with a stack buffer, which the compiler incorrectly promotes to a heap allocation.

q.Q calls q.q, which returns the result of appending to the buffer. The result of the analysis should be that the argument dst leaks to the result but not to the heap. And in fact if you build this, you do get that result when analyzing q.q and q.Q:

% go build -gcflags=./p=-m=2 -gcflags=./q=-m=2 ./p
...
q/q.go:9:25: parameter dst leaks to ~r0 for q[go.shape.int64] with derefs=0:
q/q.go:9:25:   flow: {temp} ← dst:
q/q.go:9:25:     from append(dst, byte(~r0)) (call parameter) at q/q.go:10:15
q/q.go:9:25:   flow: ~r0 ← {temp}:
q/q.go:9:25:     from return append(dst, byte(~r0)) (return) at q/q.go:10:2
...
q/q.go:5:8: parameter dst leaks to ~r0 for Q with derefs=0:
q/q.go:5:8:   flow: {temp} ← dst:
q/q.go:5:8:     from q[go.shape.int64](&.dict.q[int64], dst, x) (call parameter) at q/q.go:6:10
q/q.go:5:8:   flow: ~r0 ← {temp}:
q/q.go:5:8:     from return q[go.shape.int64](&.dict.q[int64], dst, x) (return) at q/q.go:6:2
q/q.go:5:8: leaking param: dst to result ~r0 level=0
...
q/q.go:9:25: parameter dst leaks to ~r0 for q[int64] with derefs=0:
q/q.go:9:25:   flow: {temp} ← dst:
q/q.go:9:25:     from q[go.shape.int64](&.dict.q[int64], dst, x) (call parameter) at q/q.go:9:6
q/q.go:9:25:   flow: ~r0 ← {temp}:
q/q.go:9:25:     from return q[go.shape.int64](&.dict.q[int64], dst, x) (return) at q/q.go:9:6
...

I made this change to the compiler to print the export data related to escape analysis, and assuming I got all the possible places that need prints, the only export data printed in q is for q.Q and q.nop. Note that there is no escape info for q.q in any form (the above analysis considered both q[go.shape.int64] and q[int64]).

q/q.go:5:6: escape export data: Q: dst:"esc:\x00\x01\x00\x01" x:""
q/q.go:3:6: escape export data: nop: x:""

Now let's look at the output from building p, which I elided above. This is the full output related to p and its calls:

...
q/q.go:3:6: can inline nop with cost 2 as: func(int) int { return x }
q/q.go:9:6: cannot inline q[go.shape.int64]: function too complex: cost 89 exceeds budget 80
q/q.go:5:6: can inline Q with cost 64 as: func([]byte, int64) []byte { return q[go.shape.int64](&.dict.q[int64], dst, x) }
q/q.go:9:6: can inline q[int64] with cost 64 as: func([]byte, int64) []byte { return q[go.shape.int64](&.dict.q[int64], dst, x) }
...
p/p.go:9:6: can inline P with cost 78 as: func() { buf = <nil>; runtime.KeepAlive(q.Q(buf[:0], int64(1))) }
p/p.go:11:23: inlining call to q.Q
p/p.go:10:6: buf escapes to heap in P:
p/p.go:10:6:   flow: q.dst ← &buf:
p/p.go:10:6:     from buf (address-of) at p/p.go:11:27
p/p.go:10:6:     from buf[:0] (slice) at p/p.go:11:27
p/p.go:10:6:     from q.dst, q.x := buf[:0], int64(1) (assign-pair) at p/p.go:11:23
p/p.go:10:6:   flow: {heap} ← q.dst:
p/p.go:10:6:     from q.q[go.shape.int64](&q..dict.q[int64], q.dst, q.x) (call parameter) at p/p.go:11:23
p/p.go:10:6: moved to heap: buf
p/p.go:11:23: ~r0 does not escape
p/p.go:9:6: escape export data: P:

As I understand it, p.P calls q.Q, which on the page calls q.q, which the compiler treats as a direct call to q.q[go.shape.int64]. The first call is inlinable, so that the actual code being compiled is p.P calling q.q[go.shape.int64] directly—the many nop calls in q.q make the actual implementation not inlinable. (The wrapper implementation q.q[int64] is not in play; it is only used when a function value is needed.)

Then the escape analysis runs over p.P. It decides that q.dst flows to the heap and therefore p.P's buf escapes and needs to be heap allocated. As far as I can tell, this decision is based on q.q[go.shape.int64] having no escape annotations available at all. The escape analysis must assume the worst.

A variant that still demonstrates the problem is to remove the nop calls and just mark q noinline:

//go:noinline
func q[T int32 | int64](dst []byte, x T) []byte {
	return append(dst, 1)
}

Interestingly, if I add

func q64(dst []byte, x int64) []byte {
	return append(dst, 1)
}

then I still don't see q64 listed on an "escape export data" line. But if I make Q call it:

func Q(dst []byte, x int64) []byte {
	return q64(dst, x)
}

Now I get the export data for q64 and the moved-to-heap goes away.

q/q.go:5:6: escape export data: Q: dst:"esc:\x00\x01\x00\x01" x:""
q/q.go:14:6: escape export data: q64: dst:"esc:\x00\x01\x00\x01" x:""
# rsc.io/tmp/escbug/p
p/p.go:9:6: can inline P with cost 23 as: func() { buf = <nil>; runtime.KeepAlive(q.Q(buf[:0], int64(1))) }
p/p.go:11:23: inlining call to q.Q
p/p.go:11:23: inlining call to q.q64
p/p.go:11:23: ~r0 does not escape
p/p.go:11:23: append does not escape
p/p.go:9:6: escape export data: P:

So perhaps in the generic version, the problem is that when exporting q.Q we don't correctly see the q.q call in a way that marks it as needing detailed export information.

Leaving for someone with more compiler expertise (and compiler responsibility). It does seem like in relatively plausible cases we are assuming the worst about the escape behavior of generic code.

[gabyhelp]

Related Issues

• cmd/compile: generic function argument causes escape to heap #48849
• cmd/compile: wrapping funcs in interfaces causes them to escape #74356
• cmd/compile: calling a method causes interface with non-pointer receiver to escape #78336
• cmd/internal/gc: incorrect escape analysis #10466 (closed)
• cmd/compile: argument escapes when it should not #18001 (closed)
• cmd/compile: escape analysis does't treat reflect header correctly #39497 (closed)
• cmd/compile: devirtualized type switch causes temporary allocation to escape #74383
• cmd/compile, bytes: bootstrap array causes bytes.Buffer to always be heap-allocated #7921
• cmd/compile: unclear escape analysis #36504 (closed)
• cmd/compile: escape analysis report mismatch with allocations #36083 (closed)

_{(Emoji vote if this was helpful or unhelpful; more detailed feedback welcome in this discussion.)}

[thepudds]

In very prelim look, this seems to get this particular example so that it no longer escapes:

diff --git a/src/cmd/compile/internal/noder/unified.go b/src/cmd/compile/internal/noder/unified.go
index ca601ef9cc..53fa300d33 100644
--- a/src/cmd/compile/internal/noder/unified.go
+++ b/src/cmd/compile/internal/noder/unified.go
@@ -315,9 +315,12 @@ func readBodies(target *ir.Package, duringInlining bool) {
                inline.CanInlineFuncs(inlDecls, nil)
                base.Flag.LowerM = oldLowerM
 
-               for _, fn := range inlDecls {
-                       fn.Body = nil // free memory
-               }
+               target.Funcs = append(target.Funcs, inlDecls...)
        }
 }

Not saying that is the correct fix (or more precisely, it is probably wrong, including it might end up in target.Funcs 2x with that? Not sure).

Mostly just poking at it to understand a little better.

(I had looked at a possibly tangentially related issue a few weeks ago related to how the inliner treats generic functions, including how the inliner can remove func bodies after concluding a function does not inline, but the func can still later be found via the IR).

[thepudds]

Looks like it was introduced in Go 1.20.

Seemed to bisect to this commit:

9944ba7 cmd/compile: fix transitive inlining of generic functions

CC @golang/compiler

[gopherbot]

Change https://go.dev/cl/781342 mentions this issue: cmd/compile/internal/escape: fix escape analysis for generic functions instantiated across packages

[thepudds]

I think I understand the escape analysis side of this, and I think I understand where the unified IR is going wrong, though it's not 100% clear to me what the proper fix is.

In any event, I sent https://go.dev/cl/781342 to help move the discussion forward. Seems to pass trybots at least.

[thepudds]

On other bit of background is that in 2022, Matthew wrote in #56280 (comment):

Instantiated bodies are not included in unified IR's export data format. It only exports generic bodies, and then reinstantiates them as appropriate. Similarly, it also reanalyzes for inlining and escape analysis.

So I think that was the expectation. I suspect the problem though was that when analyzing the reinstantiation, escape analysis saw a nil body and could not reanalyze properly.

[cherrymui]

Not saying that is the correct fix (or more precisely, it is probably wrong, including it might end up in target.Funcs 2x with that? Not sure).

I have not looked into the details, but I agree with you that adding more functions to the list of functions to be compiled is probably not the right approach. We are already compiling too many instantiated functions unnecessarily (#56718). I don't think we want to do more. I think the information should be there somewhere without recompiling. Maybe we need to revisit what are included in the export data for instantiated functions (which may also address #56718).

[thepudds]

Hi Cherry,

I agree with you that adding more functions to the list of functions to be compiled is probably not the right approach.

Agreed. That was just the first thing I tried within an hour of Russ opening this issue to see if it would get the proper escape analysis results to appear, which it did.

The second thing I tried and mailed that same day in the WIP https://go.dev/cl/781342 got the same "good" escape analysis result by doing the escape analysis work, but it did so without adding to the work done downstream in walk, SSA, and so on. (I don't know if that second approach is a great approach either, but it at least helped me understand the problem better, which I now understand better than I did initially 😅).

It could well be time to expand how Unified IR exports instantiated functions as you suggest.

Finally one bit of data I gathered while poking at this -- this problem seems to affect ~1% of *ir.Func seen when compiling std and ~2% for vitess.io/vitess (though probably best to consider this data preliminary):

By instrumenting the compiler, it seems roughly 1% of all the *ir.Func seen when compiling the standard library had essentially no actually escape analysis done, with all their parameters then tagged as escaping. Spot checking Vitess as a large external package that makes heavy-ish use of generics, it was ~2% of functions defined in its primary vitess.io/vitess module, though that drops to ~1.2% when the universe is expanded to include all functions in its build (that is, including the stdlib and other dependencies).

(Quote taken from the WIP https://go.dev/cl/781342).