compiler: add //go:noheap pragma

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,
 					})
 				}
 			}

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

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)

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),
 		}, "")

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

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)
+}

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 }

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"},

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.
 

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
 

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

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.