runtime: add memprofile attribution support

cl/compile.go

@@ -152,23 +152,25 @@ type pkgInfo struct {
 type none = struct{}
 
 type context struct {
-	prog        llssa.Program
-	pkg         llssa.Package
-	fn          llssa.Function
-	goFn        *ssa.Function
-	fset        *token.FileSet
-	goProg      *ssa.Program
-	goTyps      *types.Package
-	goPkg       *ssa.Package
-	pyMod       string
-	skips       map[string]none
-	loaded      map[*types.Package]*pkgInfo // loaded packages
-	bvals       map[ssa.Value]llssa.Expr    // block values
-	vargs       map[*ssa.Alloc][]llssa.Expr // varargs
-	funcs       map[*ssa.Function]llssa.Function
-	stackDefers map[*ssa.Function]bool
-	anonDefers  map[*ssa.Function]bool
-	paramDIVars map[*types.Var]llssa.DIVar
+	prog                  llssa.Program
+	pkg                   llssa.Package
+	fn                    llssa.Function
+	goFn                  *ssa.Function
+	fset                  *token.FileSet
+	goProg                *ssa.Program
+	goTyps                *types.Package
+	goPkg                 *ssa.Package
+	pyMod                 string
+	skips                 map[string]none
+	loaded                map[*types.Package]*pkgInfo // loaded packages
+	bvals                 map[ssa.Value]llssa.Expr    // block values
+	vargs                 map[*ssa.Alloc][]llssa.Expr // varargs
+	funcs                 map[*ssa.Function]llssa.Function
+	stackDefers           map[*ssa.Function]bool
+	anonDefers            map[*ssa.Function]bool
+	paramDIVars           map[*types.Var]llssa.DIVar
+	noInlineForMemProfile bool
+	memProfileInstrument  bool
 
 	patches  Patches
 	blkInfos []blocks.Info
@@ -358,6 +360,23 @@ func isInstance(f *ssa.Function) bool {
 	return false
 }
 
+func (p *context) applyNoInline(fn llssa.Function) {
+	if disableInline || p.noInlineForMemProfile {
+		fn.Inline(llssa.NoInline)
+	}
+	if p.noInlineForMemProfile {
+		fn.DisableTailCalls()
+	}
+}
+
+func memProfileFunctionName(name string) string {
+	const commandLineArguments = "command-line-arguments."
+	if strings.HasPrefix(name, commandLineArguments) {
+		return "main." + name[len(commandLineArguments):]
+	}
+	return name
+}
+
 func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Function, llssa.PyObjRef, int) {
 	pkgTypes, name, ftype := p.funcName(f)
 	if ftype != goFunc {
@@ -395,10 +414,8 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 	}
 	if fn == nil {
 		fn = pkg.NewFuncEx(name, sig, llssa.Background(ftype), hasCtx, isInstance(f))
-		if disableInline {
-			fn.Inline(llssa.NoInline)
-		}
 	}
+	p.applyNoInline(fn)
 	p.funcs[f] = fn
 	isCgo := isCgoExternSymbol(f)
 	if nblk := len(f.Blocks); nblk > 0 {
@@ -425,13 +442,17 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 		}
 		dbgEnabled := enableDbg && (f == nil || f.Origin() == nil)
 		dbgSymsEnabled := enableDbgSyms && (f == nil || f.Origin() == nil)
+		instrumentMemProfile := p.noInlineForMemProfile && !isCgo
 		p.inits = append(p.inits, func() {
 			oldFn, oldGoFn := p.fn, p.goFn
+			oldMemProfileInstrument := p.memProfileInstrument
 			p.fn = fn
 			p.goFn = f
+			p.memProfileInstrument = instrumentMemProfile
 			p.state = state // restore pkgState when compiling funcBody
 			defer func() {
 				p.fn, p.goFn = oldFn, oldGoFn
+				p.memProfileInstrument = oldMemProfileInstrument
 			}()
 			p.phis = nil
 			if dbgSymsEnabled {
@@ -552,6 +573,9 @@ func (p *context) compileBlock(b llssa.Builder, block *ssa.BasicBlock, n int, do
 	var instrs = block.Instrs[n:]
 	var ret = fn.Block(block.Index)
 	b.SetBlock(ret)
+	if block.Index == 0 && p.memProfileInstrument {
+		b.MemProfileEnter(memProfileFunctionName(fn.Name()))
+	}
 	if block.Index == 0 && enableCallTracing && !strings.HasPrefix(fn.Name(), "github.com/goplus/llgo/runtime/internal/runtime.Print") {
 		b.Printf("call " + fn.Name() + "\n\x00")
 	}
@@ -1122,6 +1146,9 @@ func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 		if p.returnNeedsImplicitRunDefers(v) {
 			b.RunDefers()
 		}
+		if p.memProfileInstrument {
+			b.MemProfileExit()
+		}
 		b.Return(results...)
 	case *ssa.If:
 		fn := p.fn
@@ -1410,6 +1437,52 @@ func NewPackageEx(prog llssa.Program, patches Patches, rewrites map[string]strin
 	return newPackageEx(prog, patches, rewrites, pkg, files, nil)
 }
 
+func packageUsesRuntimeMemProfile(files []*ast.File) bool {
+	for _, file := range files {
+		runtimeNames := make(map[string]none)
+		for _, imp := range file.Imports {
+			path, err := strconv.Unquote(imp.Path.Value)
+			if err != nil || path != "runtime" {
+				continue
+			}
+			if imp.Name != nil {
+				if imp.Name.Name == "_" || imp.Name.Name == "." {
+					continue
+				}
+				runtimeNames[imp.Name.Name] = none{}
+				continue
+			}
+			runtimeNames["runtime"] = none{}
+		}
+		if len(runtimeNames) == 0 {
+			continue
+		}
+		found := false
+		ast.Inspect(file, func(n ast.Node) bool {
+			sel, ok := n.(*ast.SelectorExpr)
+			if !ok {
+				return true
+			}
+			if sel.Sel.Name != "MemProfile" && sel.Sel.Name != "MemProfileRate" {
+				return true
+			}
+			x, ok := sel.X.(*ast.Ident)
+			if !ok {
+				return true
+			}
+			if _, ok := runtimeNames[x.Name]; !ok {
+				return true
+			}
+			found = true
+			return false
+		})
+		if found {
+			return true
+		}
+	}
+	return false
+}
+
 // NewPackageExWithEmbed compiles a package using pre-loaded go:embed metadata.
 //
 // This avoids re-scanning directives when the caller already loaded them.
@@ -1437,16 +1510,17 @@ func newPackageEx(prog llssa.Program, patches Patches, rewrites map[string]strin
 	}
 
 	ctx := &context{
-		prog:    prog,
-		pkg:     ret,
-		fset:    pkgProg.Fset,
-		goProg:  pkgProg,
-		goTyps:  pkgTypes,
-		goPkg:   pkg,
-		patches: patches,
-		skips:   make(map[string]none),
-		vargs:   make(map[*ssa.Alloc][]llssa.Expr),
-		funcs:   make(map[*ssa.Function]llssa.Function),
+		prog:                  prog,
+		pkg:                   ret,
+		fset:                  pkgProg.Fset,
+		goProg:                pkgProg,
+		goTyps:                pkgTypes,
+		goPkg:                 pkg,
+		patches:               patches,
+		noInlineForMemProfile: packageUsesRuntimeMemProfile(files),
+		skips:                 make(map[string]none),
+		vargs:                 make(map[*ssa.Alloc][]llssa.Expr),
+		funcs:                 make(map[*ssa.Function]llssa.Function),
 		loaded: map[*types.Package]*pkgInfo{
 			types.Unsafe: {kind: PkgDeclOnly}, // TODO(xsw): PkgNoInit or PkgDeclOnly?
 		},

cl/instr.go

@@ -644,9 +644,7 @@ func (p *context) funcOf(fn *ssa.Function) (aFn llssa.Function, pyFn llssa.PyObj
 			}
 			sig := p.patchType(fn.Signature).(*types.Signature)
 			aFn = pkg.NewFuncEx(name, sig, llssa.Background(ftype), false, fn.Origin() != nil)
-			if disableInline {
-				aFn.Inline(llssa.NoInline)
-			}
+			p.applyNoInline(aFn)
 		}
 	}
 	return

runtime/internal/lib/runtime/debug.go

@@ -1,5 +1,7 @@
 package runtime
 
+import llrt "github.com/goplus/llgo/runtime/internal/runtime"
+
 func NumCPU() int {
 	return int(c_maxprocs())
 }
@@ -9,4 +11,5 @@ func Breakpoint() {
 }
 
 func Gosched() {
+	llrt.SetMemProfileRate(MemProfileRate)
 }

runtime/internal/lib/runtime/pprof_runtime_stub_llgo.go

@@ -2,18 +2,27 @@
 
 package runtime
 
+import llrt "github.com/goplus/llgo/runtime/internal/runtime"
+
 // StackRecord is a minimal placeholder for runtime/pprof.
 type StackRecord struct {
-	Stack []uintptr
+	Stack0 [32]uintptr
+}
+
+func (r *StackRecord) Stack() []uintptr {
+	for i, pc := range r.Stack0 {
+		if pc == 0 {
+			return r.Stack0[:i]
+		}
+	}
+	return r.Stack0[:]
 }
 
 // MemProfileRecord is a minimal placeholder for runtime/pprof.
 type MemProfileRecord struct {
-	AllocBytes   int64
-	FreeBytes    int64
-	AllocObjects int64
-	FreeObjects  int64
-	Stack        []uintptr
+	AllocBytes, FreeBytes     int64
+	AllocObjects, FreeObjects int64
+	Stack0                    [32]uintptr
 }
 
 func (r *MemProfileRecord) InUseBytes() int64 {
@@ -24,15 +33,80 @@ func (r *MemProfileRecord) InUseObjects() int64 {
 	return r.AllocObjects - r.FreeObjects
 }
 
+func (r *MemProfileRecord) Stack() []uintptr {
+	for i, pc := range r.Stack0 {
+		if pc == 0 {
+			return r.Stack0[:i]
+		}
+	}
+	return r.Stack0[:]
+}
+
 // BlockProfileRecord is a minimal placeholder for runtime/pprof.
 type BlockProfileRecord struct {
 	Count  int64
 	Cycles int64
-	Stack  []uintptr
+	StackRecord
 }
 
 func MemProfile(p []MemProfileRecord, inuseZero bool) (n int, ok bool) {
-	return 0, false
+	llrt.SetMemProfileRate(MemProfileRate)
+	records := llrt.MemProfileSnapshot()
+	n = len(records)
+	if len(p) < n {
+		return n, false
+	}
+	for i, r := range records {
+		allocObjects, allocBytes := scaleMemProfileSample(r.AllocObjects, r.AllocBytes, MemProfileRate)
+		p[i] = MemProfileRecord{
+			AllocBytes:   allocBytes,
+			AllocObjects: allocObjects,
+			Stack0:       r.Stack0,
+		}
+	}
+	return n, true
+}
+
+func scaleMemProfileSample(objects, bytes int64, rate int) (int64, int64) {
+	if objects <= 0 || bytes <= 0 {
+		return 0, 0
+	}
+	if rate <= 1 {
+		return objects, bytes
+	}
+	avgSize := float64(bytes) / float64(objects)
+	prob := 1 - expNeg(avgSize/float64(rate))
+	if prob <= 0 {
+		return 1, bytes / objects
+	}
+	sampledObjects := int64(float64(objects)*prob + 0.5)
+	if sampledObjects < 1 {
+		sampledObjects = 1
+	}
+	return sampledObjects, sampledObjects * (bytes / objects)
+}
+
+func expNeg(x float64) float64 {
+	if x <= 0 {
+		return 1
+	}
+	if x > 0.5 {
+		y := expNeg(x / 2)
+		return y * y
+	}
+	term := 1.0
+	sum := 1.0
+	for i := 1; i <= 12; i++ {
+		term *= -x / float64(i)
+		sum += term
+	}
+	if sum < 0 {
+		return 0
+	}
+	if sum > 1 {
+		return 1
+	}
+	return sum
 }
 
 func BlockProfile(p []BlockProfileRecord) (n int, ok bool) {

runtime/internal/lib/runtime/runtime_gc.go

@@ -3,21 +3,23 @@
 package runtime
 
 import (
-	"runtime"
+	goruntime "runtime"
 
 	"github.com/goplus/llgo/runtime/internal/clite/bdwgc"
+	llrt "github.com/goplus/llgo/runtime/internal/runtime"
 )
 
-func ReadMemStats(m *runtime.MemStats) {
+func ReadMemStats(m *goruntime.MemStats) {
 	if m == nil {
 		return
 	}
 	// LLGo currently doesn't provide accurate allocation statistics when using BDWGC.
 	// Populate a zeroed snapshot so stdlib callers like testing.AllocsPerRun can run.
-	*m = runtime.MemStats{}
+	*m = goruntime.MemStats{}
 }
 
 func GC() {
+	llrt.SetMemProfileRate(MemProfileRate)
 	bdwgc.Gcollect()
 	// BDW finalizers are observed on a subsequent collection cycle.
 	// Run one extra cycle so weak-pointer cleanup hooks (unique/weak) see

runtime/internal/lib/runtime/symtab.go

@@ -9,6 +9,7 @@ import (
 
 	c "github.com/goplus/llgo/runtime/internal/clite"
 	clitedebug "github.com/goplus/llgo/runtime/internal/clite/debug"
+	llrt "github.com/goplus/llgo/runtime/internal/runtime"
 )
 
 // Frames may be used to get function/file/line information for a
@@ -119,6 +120,16 @@ func (ci *Frames) Next() (frame Frame, more bool) {
 		} else {
 			pc, ci.callers = ci.callers[0], ci.callers[1:]
 		}
+		if fn, line, ok := llrt.MemProfileSyntheticFrame(pc); ok {
+			ci.frames = append(ci.frames, Frame{
+				PC:        pc,
+				Function:  fn,
+				Line:      line,
+				startLine: line,
+				Entry:     pc,
+			})
+			continue
+		}
 		info := &clitedebug.Info{}
 		if clitedebug.Addrinfo(unsafe.Pointer(pc), info) == 0 {
 			ci.frames = append(ci.frames, Frame{