compiler: add //go:noheap pragma

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 }