From 85d223c5e28a2366d95cadaff25c68e238d5cd15 Mon Sep 17 00:00:00 2001 From: Ayke van Laethem Date: Wed, 20 May 2026 11:01:50 +0200 Subject: [PATCH] compiler: add //go:noheap pragma --- builder/tools.go | 6 +++- compiler/compiler.go | 14 ++++++-- compiler/defer.go | 2 +- compiler/llvm.go | 2 +- compiler/symbol.go | 6 +++- compiler/testdata/pragma.go | 5 +++ compiler/testdata/pragma.ll | 14 ++++++++ errors_test.go | 1 + interp/interpreter.go | 2 +- src/runtime/runtime.go | 5 +++ testdata/errors/noheap.go | 69 +++++++++++++++++++++++++++++++++++++ transform/allocs.go | 19 ++++++++-- 12 files changed, 135 insertions(+), 10 deletions(-) create mode 100644 testdata/errors/noheap.go diff --git a/builder/tools.go b/builder/tools.go index 9d15e4cca..8a5762cb5 100644 --- a/builder/tools.go +++ b/builder/tools.go @@ -122,12 +122,16 @@ func parseLLDErrors(text string) error { parsedError = true line, _ := strconv.Atoi(matches[3]) // TODO: detect common mistakes like -gc=none? + msg := "linker could not find symbol " + symbolName + if symbolName == "runtime.alloc_noheap" { + msg = "object allocated on the heap in //go:noheap function" + } linkErrors = append(linkErrors, scanner.Error{ Pos: token.Position{ Filename: matches[2], Line: line, }, - Msg: "linker could not find symbol " + symbolName, + Msg: msg, }) } } diff --git a/compiler/compiler.go b/compiler/compiler.go index a4dd94859..eaf7b58b1 100644 --- a/compiler/compiler.go +++ b/compiler/compiler.go @@ -192,6 +192,16 @@ func newBuilder(c *compilerContext, irbuilder llvm.Builder, f *ssa.Function) *bu } } +// Return the runtime.alloc function variant. +// This is normally just "alloc", but is "alloc_noheap" if the //go:noheap +// pragma is used. +func (b *builder) allocFunc() string { + if b.info.noheap { + return "alloc_noheap" + } + return "alloc" +} + type blockInfo struct { // entry is the LLVM basic block corresponding to the start of this *ssa.Block. entry llvm.BasicBlock @@ -2173,7 +2183,7 @@ func (b *builder) createExpr(expr ssa.Value) (llvm.Value, error) { } sizeValue := llvm.ConstInt(b.uintptrType, size, false) layoutValue := b.createObjectLayout(typ, expr.Pos()) - buf := b.createRuntimeCall("alloc", []llvm.Value{sizeValue, layoutValue}, expr.Comment) + buf := b.createRuntimeCall(b.allocFunc(), []llvm.Value{sizeValue, layoutValue}, expr.Comment) align := b.targetData.ABITypeAlignment(typ) buf.AddCallSiteAttribute(0, b.ctx.CreateEnumAttribute(llvm.AttributeKindID("align"), uint64(align))) return buf, nil @@ -2405,7 +2415,7 @@ func (b *builder) createExpr(expr ssa.Value) (llvm.Value, error) { } sliceSize := b.CreateBinOp(llvm.Mul, elemSizeValue, sliceCapCast, "makeslice.cap") layoutValue := b.createObjectLayout(llvmElemType, expr.Pos()) - slicePtr := b.createRuntimeCall("alloc", []llvm.Value{sliceSize, layoutValue}, "makeslice.buf") + slicePtr := b.createRuntimeCall(b.allocFunc(), []llvm.Value{sliceSize, layoutValue}, "makeslice.buf") slicePtr.AddCallSiteAttribute(0, b.ctx.CreateEnumAttribute(llvm.AttributeKindID("align"), uint64(elemAlign))) // Extend or truncate if necessary. This is safe as we've already done diff --git a/compiler/defer.go b/compiler/defer.go index 7c95f0acf..ec2bbe00e 100644 --- a/compiler/defer.go +++ b/compiler/defer.go @@ -488,7 +488,7 @@ func (b *builder) createDefer(instr *ssa.Defer) { size := b.targetData.TypeAllocSize(deferredCallType) sizeValue := llvm.ConstInt(b.uintptrType, size, false) nilPtr := llvm.ConstNull(b.dataPtrType) - alloca = b.createRuntimeCall("alloc", []llvm.Value{sizeValue, nilPtr}, "defer.alloc.call") + alloca = b.createRuntimeCall(b.allocFunc(), []llvm.Value{sizeValue, nilPtr}, "defer.alloc.call") } if b.NeedsStackObjects { b.trackPointer(alloca) diff --git a/compiler/llvm.go b/compiler/llvm.go index 5ef656ce5..c25a8d4a3 100644 --- a/compiler/llvm.go +++ b/compiler/llvm.go @@ -129,7 +129,7 @@ func (b *builder) emitPointerPack(values []llvm.Value) llvm.Value { // Packed data is bigger than a pointer, so allocate it on the heap. sizeValue := llvm.ConstInt(b.uintptrType, size, false) align := b.targetData.ABITypeAlignment(packedType) - packedAlloc := b.createRuntimeCall("alloc", []llvm.Value{ + packedAlloc := b.createRuntimeCall(b.allocFunc(), []llvm.Value{ sizeValue, llvm.ConstNull(b.dataPtrType), }, "") diff --git a/compiler/symbol.go b/compiler/symbol.go index 16c5433cf..0e78e2ce3 100644 --- a/compiler/symbol.go +++ b/compiler/symbol.go @@ -34,6 +34,7 @@ type functionInfo struct { interrupt bool // go:interrupt nobounds bool // go:nobounds noescape bool // go:noescape + noheap bool // go:noheap variadic bool // go:variadic (CGo only) inline inlineType // go:inline } @@ -161,7 +162,7 @@ func (c *compilerContext) getFunction(fn *ssa.Function) (llvm.Type, llvm.Value) llvmFn.AddAttributeAtIndex(1, c.ctx.CreateEnumAttribute(llvm.AttributeKindID("nocapture"), 0)) case "machine.keepAliveNoEscape", "machine.unsafeNoEscape": llvmFn.AddAttributeAtIndex(1, c.ctx.CreateEnumAttribute(llvm.AttributeKindID("nocapture"), 0)) - case "runtime.alloc": + case "runtime.alloc", "runtime.alloc_noheap": // Tell the optimizer that runtime.alloc is an allocator, meaning that it // returns values that are never null and never alias to an existing value. for _, attrName := range []string{"noalias", "nonnull"} { @@ -476,6 +477,9 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) { if len(f.Blocks) == 0 { info.noescape = true } + case "//go:noheap": + // Ensure this function does not allocate on the heap. + info.noheap = true case "//go:variadic": // The //go:variadic pragma is emitted by the CGo preprocessing // pass for C variadic functions. This includes both explicit diff --git a/compiler/testdata/pragma.go b/compiler/testdata/pragma.go index 1e6e967f5..ebb3a2a67 100644 --- a/compiler/testdata/pragma.go +++ b/compiler/testdata/pragma.go @@ -115,3 +115,8 @@ func doesNotEscapeParam(a *int, b []int, c chan int, d *[0]byte) //go:noescape func stillEscapes(a *int, b []int, c chan int, d *[0]byte) { } + +//go:noheap +func doesHeapAlloc() *int { + return new(int) +} diff --git a/compiler/testdata/pragma.ll b/compiler/testdata/pragma.ll index 173f2d8c6..c2c06b3d4 100644 --- a/compiler/testdata/pragma.ll +++ b/compiler/testdata/pragma.ll @@ -90,6 +90,18 @@ entry: ret void } +; Function Attrs: nounwind +define hidden ptr @main.doesHeapAlloc(ptr %context) unnamed_addr #1 { +entry: + %stackalloc = alloca i8, align 1 + %new = call align 4 dereferenceable(4) ptr @runtime.alloc_noheap(i32 4, ptr nonnull inttoptr (i32 3 to ptr), ptr undef) #10 + call void @runtime.trackPointer(ptr nonnull %new, ptr nonnull %stackalloc, ptr undef) #10 + ret ptr %new +} + +; Function Attrs: allockind("alloc,zeroed") allocsize(0) +declare noalias nonnull ptr @runtime.alloc_noheap(i32, ptr, ptr) #9 + attributes #0 = { "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" } attributes #1 = { nounwind "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" } attributes #2 = { nounwind "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" "wasm-export-name"="extern_func" } @@ -99,3 +111,5 @@ attributes #5 = { noinline nounwind "target-features"="+bulk-memory,+bulk-memory attributes #6 = { "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" "wasm-import-module"="modulename" "wasm-import-name"="import1" } attributes #7 = { "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" "wasm-import-module"="foobar" "wasm-import-name"="imported" } attributes #8 = { nounwind "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" "wasm-export-name"="exported" } +attributes #9 = { allockind("alloc,zeroed") allocsize(0) "alloc-family"="runtime.alloc" "target-features"="+bulk-memory,+bulk-memory-opt,+call-indirect-overlong,+mutable-globals,+nontrapping-fptoint,+sign-ext,-multivalue,-reference-types" } +attributes #10 = { nounwind } diff --git a/errors_test.go b/errors_test.go index 871dc4c04..69c372bf2 100644 --- a/errors_test.go +++ b/errors_test.go @@ -38,6 +38,7 @@ func TestErrors(t *testing.T) { {name: "loader-invaliddep"}, {name: "loader-invalidpackage"}, {name: "loader-nopackage"}, + {name: "noheap", target: "cortex-m-qemu"}, {name: "optimizer"}, {name: "syntax"}, {name: "types"}, diff --git a/interp/interpreter.go b/interp/interpreter.go index 34bb668da..710f3e813 100644 --- a/interp/interpreter.go +++ b/interp/interpreter.go @@ -299,7 +299,7 @@ func (r *runner) run(fn *function, params []value, parentMem *memoryView, indent // means that monotonic time in the time package is counted from // time.Time{}.Sub(1), which should be fine. locals[inst.localIndex] = literalValue{uint64(0)} - case callFn.name == "runtime.alloc": + case callFn.name == "runtime.alloc" || callFn.name == "runtime.alloc_noheap": // Allocate heap memory. At compile time, this is instead done // by creating a global variable. diff --git a/src/runtime/runtime.go b/src/runtime/runtime.go index c9b095938..872b225b9 100644 --- a/src/runtime/runtime.go +++ b/src/runtime/runtime.go @@ -66,6 +66,11 @@ func llvm_sponentry() unsafe.Pointer //export strlen func strlen(ptr unsafe.Pointer) uintptr +// Special alloc function that should never actually be called. +// It is used instead of normal alloc in //go:noheap functions, and must either +// be optimized away or throw a linker error. +func alloc_noheap(size uintptr, layout unsafe.Pointer) unsafe.Pointer + //export malloc func malloc(size uintptr) unsafe.Pointer diff --git a/testdata/errors/noheap.go b/testdata/errors/noheap.go new file mode 100644 index 000000000..074a6221b --- /dev/null +++ b/testdata/errors/noheap.go @@ -0,0 +1,69 @@ +package main + +//go:noheap +func main() { + // This object is optimized away, and won't cause a linker failure. + var a int + add(&a) + + // This object is not optimized away since its address is captured. + var b int + escape(&b) + + // A simple defer won't cause a heap allocation. + defer func() { + println("once") + }() + + // A defer in a loop does causes a heap allocation. + for range 3 { + defer func() { + println("many") + }() + } + + // This slice should not be heap-allocated. + s1 := make([]int, 5) + sum(s1) + + // But this slice is. + s2 := make([]int, 5) + escape(&s2[0]) + + // Allocated at compile time as a global, so won't escape at runtime. + escape(globalInt) +} + +func add(n *int) { + *n++ +} + +func sum(slice []int) (result int) { + for n := range slice { + result += n + } + return result +} + +var globalInt = globalNewInt() + +var globalAssign *int + +//go:noheap +func globalNewInt() *int { + return new(int) +} + +func escape(n *int) { + // Do some stuff that will definitely force a heap allocation. + // (This might be optimized in the future, in which case we need to change + // the function again). + n2 := func(n *int) *int { + return n + }(n) + println(n2) +} + +// ERROR: noheap.go:10: object allocated on the heap in //go:noheap function +// ERROR: noheap.go:20: object allocated on the heap in //go:noheap function +// ERROR: noheap.go:30: object allocated on the heap in //go:noheap function diff --git a/transform/allocs.go b/transform/allocs.go index d9c7e4709..18b49998c 100644 --- a/transform/allocs.go +++ b/transform/allocs.go @@ -23,8 +23,15 @@ import ( // heap allocation explanation should be printed (why the object can't be stack // allocated). func OptimizeAllocs(mod llvm.Module, printAllocs *regexp.Regexp, maxStackAlloc uint64, logger func(token.Position, string)) { - allocator := mod.NamedFunction("runtime.alloc") - if allocator.IsNil() { + // Find allocator functions. + var allocators []llvm.Value + for _, name := range []string{"runtime.alloc", "runtime.alloc_noheap"} { + allocator := mod.NamedFunction(name) + if !allocator.IsNil() { + allocators = append(allocators, allocator) + } + } + if len(allocators) == 0 { // nothing to optimize return } @@ -43,7 +50,13 @@ func OptimizeAllocs(mod llvm.Module, printAllocs *regexp.Regexp, maxStackAlloc u fmt.Fprintln(os.Stderr, "mode: set") } - for _, heapalloc := range getUses(allocator) { + // Find all allocator calls. + var heapallocs []llvm.Value + for _, allocator := range allocators { + heapallocs = append(heapallocs, getUses(allocator)...) + } + + for _, heapalloc := range heapallocs { logAllocs := printAllocs != nil && printAllocs.MatchString(heapalloc.InstructionParent().Parent().Name()) if heapalloc.Operand(0).IsAConstantInt().IsNil() { // Do not allocate variable length arrays on the stack.