machine: make sure DMA buffers do not escape unnecessarily

Writing the pointer of a buffer to memory-mapped I/O will normally cause
it to escape, which forces the compiler to heap-allocate the buffer. But
we do know how long the value stays alive, so we can tell the compiler
to keep it alive exactly until it is not needed anymore - and tell it to
not treat the pointer-to-uintptr cast as escaping.

Author: aykevl

compiler/intrinsics.go

@@ -27,6 +27,8 @@ func (b *builder) defineIntrinsicFunction() {
 		b.createStackSaveImpl()
 	case name == "runtime.KeepAlive":
 		b.createKeepAliveImpl()
+	case name == "machine.keepAliveNoEscape":
+		b.createMachineKeepAliveImpl()
 	case strings.HasPrefix(name, "runtime/volatile.Load"):
 		b.createVolatileLoad()
 	case strings.HasPrefix(name, "runtime/volatile.Store"):
@@ -144,6 +146,20 @@ func (b *builder) createAbiEscapeImpl() {
 	b.CreateRet(result)
 }
 
+// Implement machine.keepAliveNoEscape, which makes sure the compiler keeps the
+// pointer parameter alive until this point (for GC).
+func (b *builder) createMachineKeepAliveImpl() {
+	b.createFunctionStart(true)
+	pointerValue := b.getValue(b.fn.Params[0], getPos(b.fn))
+
+	// See createKeepAliveImpl for details.
+	asmType := llvm.FunctionType(b.ctx.VoidType(), []llvm.Type{b.dataPtrType}, false)
+	asmFn := llvm.InlineAsm(asmType, "", "r", true, false, 0, false)
+	b.createCall(asmType, asmFn, []llvm.Value{pointerValue}, "")
+
+	b.CreateRetVoid()
+}
+
 var mathToLLVMMapping = map[string]string{
 	"math.Ceil":  "llvm.ceil.f64",
 	"math.Exp":   "llvm.exp.f64",

compiler/symbol.go

@@ -154,6 +154,8 @@ func (c *compilerContext) getFunction(fn *ssa.Function) (llvm.Type, llvm.Value)
 		llvmFn.AddFunctionAttr(c.ctx.CreateEnumAttribute(llvm.AttributeKindID("noreturn"), 0))
 	case "internal/abi.NoEscape":
 		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":
 		// 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.

src/machine/machine.go

@@ -1,6 +1,9 @@
 package machine
 
-import "errors"
+import (
+	"errors"
+	"unsafe"
+)
 
 var (
 	ErrTimeoutRNG         = errors.New("machine: RNG Timeout")
@@ -62,3 +65,30 @@ func (p Pin) Low() {
 type ADC struct {
 	Pin Pin
 }
+
+// Convert the pointer to a uintptr, to be used for memory I/O (DMA for
+// example). It also means the pointer is "gone" as far as the compiler is
+// concerned, and a GC cycle might deallocate the object. To prevent this from
+// happening, also call keepAliveNoEscape at a point after the address isn't
+// accessed anymore by the hardware.
+// The only exception is if the pointer is accessed later in a volatile way
+// (volatile read/write), which also forces the value to stay alive until that
+// point.
+//
+// This function is treated specially by the compiler to mark the 'ptr'
+// parameter as not escaping.
+//
+// TODO: this function should eventually be replaced with the proposed ptrtoaddr
+// instruction in LLVM. See:
+// https://discourse.llvm.org/t/clarifiying-the-semantics-of-ptrtoint/83987/10
+// https://github.com/llvm/llvm-project/pull/139357
+func unsafeNoEscape(ptr unsafe.Pointer) uintptr {
+	return uintptr(ptr)
+}
+
+// Make sure the given pointer stays alive until this point. This is similar to
+// runtime.KeepAlive, with the difference that it won't let the pointer escape.
+// This is typically used together with unsafeNoEscape.
+//
+// This is a compiler intrinsic.
+func keepAliveNoEscape(ptr unsafe.Pointer)

src/machine/machine_nrf528xx.go

@@ -49,7 +49,7 @@ func (i2c *I2C) Tx(addr uint16, w, r []byte) (err error) {
 
 	// Configure for a single shot to perform both write and read (as applicable)
 	if len(w) != 0 {
-		i2c.Bus.TXD.PTR.Set(uint32(uintptr(unsafe.Pointer(&w[0]))))
+		i2c.Bus.TXD.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&w[0]))))
 		i2c.Bus.TXD.MAXCNT.Set(uint32(len(w)))
 
 		// If no read, immediately signal stop after TX
@@ -58,7 +58,7 @@ func (i2c *I2C) Tx(addr uint16, w, r []byte) (err error) {
 		}
 	}
 	if len(r) != 0 {
-		i2c.Bus.RXD.PTR.Set(uint32(uintptr(unsafe.Pointer(&r[0]))))
+		i2c.Bus.RXD.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&r[0]))))
 		i2c.Bus.RXD.MAXCNT.Set(uint32(len(r)))
 
 		// Auto-start Rx after Tx and Stop after Rx
@@ -89,6 +89,15 @@ func (i2c *I2C) Tx(addr uint16, w, r []byte) (err error) {
 		}
 	}
 
+	// Make sure the w and r buffers stay alive until this point, so they won't
+	// be garbage collected while the buffers are used by the hardware.
+	if len(w) > 0 {
+		keepAliveNoEscape(unsafe.Pointer(&w[0]))
+	}
+	if len(r) > 0 {
+		keepAliveNoEscape(unsafe.Pointer(&r[0]))
+	}
+
 	return
 }
 
@@ -117,7 +126,7 @@ func (i2c *I2C) Listen(addr uint8) error {
 //
 // For request events, the caller MUST call `Reply` to avoid hanging the i2c bus indefinitely.
 func (i2c *I2C) WaitForEvent(buf []byte) (evt I2CTargetEvent, count int, err error) {
-	i2c.BusT.RXD.PTR.Set(uint32(uintptr(unsafe.Pointer(&buf[0]))))
+	i2c.BusT.RXD.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&buf[0]))))
 	i2c.BusT.RXD.MAXCNT.Set(uint32(len(buf)))
 
 	i2c.BusT.TASKS_PREPARERX.Set(nrf.TWIS_TASKS_PREPARERX_TASKS_PREPARERX_Trigger)
@@ -134,6 +143,10 @@ func (i2c *I2C) WaitForEvent(buf []byte) (evt I2CTargetEvent, count int, err err
 		}
 	}
 
+	// Make sure buf stays alive until this point, so it won't be garbage
+	// collected while it is used by the hardware.
+	keepAliveNoEscape(unsafe.Pointer(&buf[0]))
+
 	count = 0
 	evt = I2CFinish
 	err = nil
@@ -163,7 +176,7 @@ func (i2c *I2C) WaitForEvent(buf []byte) (evt I2CTargetEvent, count int, err err
 
 // Reply supplies the response data the controller.
 func (i2c *I2C) Reply(buf []byte) error {
-	i2c.BusT.TXD.PTR.Set(uint32(uintptr(unsafe.Pointer(&buf[0]))))
+	i2c.BusT.TXD.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&buf[0]))))
 	i2c.BusT.TXD.MAXCNT.Set(uint32(len(buf)))
 
 	i2c.BusT.EVENTS_STOPPED.Set(0)
@@ -180,6 +193,10 @@ func (i2c *I2C) Reply(buf []byte) error {
 		}
 	}
 
+	// Make sure the buffer stays alive until this point, so it won't be garbage
+	// collected while it is used by the hardware.
+	keepAliveNoEscape(unsafe.Pointer(&buf[0]))
+
 	i2c.BusT.EVENTS_STOPPED.Set(0)
 
 	return nil

src/machine/machine_nrf52xxx.go

@@ -145,7 +145,9 @@ func (a *ADC) Get() uint16 {
 	nrf.SAADC.CH[0].PSELP.Set(pwmPin)
 
 	// Destination for sample result.
-	nrf.SAADC.RESULT.PTR.Set(uint32(uintptr(unsafe.Pointer(&rawValue))))
+	// Note: rawValue doesn't need to be kept alive for the GC, since the
+	// volatile read later will force it to stay alive.
+	nrf.SAADC.RESULT.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&rawValue))))
 	nrf.SAADC.RESULT.MAXCNT.Set(1) // One sample
 
 	// Start tasks.
@@ -314,7 +316,7 @@ func (spi *SPI) Tx(w, r []byte) error {
 			if nr > 255 {
 				nr = 255
 			}
-			spi.Bus.RXD.PTR.Set(uint32(uintptr(unsafe.Pointer(&r[0]))))
+			spi.Bus.RXD.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&r[0]))))
 			r = r[nr:]
 		}
 		spi.Bus.RXD.MAXCNT.Set(nr)
@@ -325,7 +327,7 @@ func (spi *SPI) Tx(w, r []byte) error {
 			if nw > 255 {
 				nw = 255
 			}
-			spi.Bus.TXD.PTR.Set(uint32(uintptr(unsafe.Pointer(&w[0]))))
+			spi.Bus.TXD.PTR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&w[0]))))
 			w = w[nw:]
 		}
 		spi.Bus.TXD.MAXCNT.Set(nw)
@@ -339,6 +341,15 @@ func (spi *SPI) Tx(w, r []byte) error {
 		spi.Bus.EVENTS_END.Set(0)
 	}
 
+	// Make sure the w and r buffers stay alive for the GC until this point,
+	// since they are used by the hardware but not otherwise visible.
+	if len(r) != 0 {
+		keepAliveNoEscape(unsafe.Pointer(&r[0]))
+	}
+	if len(w) != 0 {
+		keepAliveNoEscape(unsafe.Pointer(&w[0]))
+	}
+
 	return nil
 }
 

src/machine/machine_rp2_spi.go

@@ -309,7 +309,7 @@ func (spi *SPI) tx(tx []byte) error {
 	//   - set data size to single bytes
 	//   - set the DREQ so that the DMA will fill the SPI FIFO as needed
 	//   - start the transfer
-	ch.READ_ADDR.Set(uint32(uintptr(unsafe.Pointer(&tx[0]))))
+	ch.READ_ADDR.Set(uint32(unsafeNoEscape(unsafe.Pointer(&tx[0]))))
 	ch.WRITE_ADDR.Set(uint32(uintptr(unsafe.Pointer(&spi.Bus.SSPDR))))
 	ch.TRANS_COUNT.Set(uint32(len(tx)))
 	ch.CTRL_TRIG.Set(rp.DMA_CH0_CTRL_TRIG_INCR_READ |
@@ -328,6 +328,11 @@ func (spi *SPI) tx(tx []byte) error {
 	for ch.CTRL_TRIG.Get()&rp.DMA_CH0_CTRL_TRIG_BUSY != 0 {
 	}
 
+	// Make sure the read buffer stays alive until this point (in the unlikely
+	// case the tx slice wasn't read after this function returns and a GC cycle
+	// happened inbetween).
+	keepAliveNoEscape(unsafe.Pointer(&tx[0]))
+
 	// We didn't read any result values, which means the RX FIFO has likely
 	// overflown. We have to clean up this mess now.
 

transform/testdata/allocs2.go

@@ -1,11 +1,17 @@
 package main
 
+import (
+	"runtime/volatile"
+	"unsafe"
+	_ "unsafe"
+)
+
 func main() {
 	n1 := 5
 	derefInt(&n1)
 
 	// This should eventually be modified to not escape.
-	n2 := 6 // OUT: object allocated on the heap: escapes at line 9
+	n2 := 6 // OUT: object allocated on the heap: escapes at line 15
 	returnIntPtr(&n2)
 
 	s1 := make([]int, 3)
@@ -15,36 +21,36 @@ func main() {
 	readIntSlice(s2[:])
 
 	// This should also be modified to not escape.
-	s3 := make([]int, 3) // OUT: object allocated on the heap: escapes at line 19
+	s3 := make([]int, 3) // OUT: object allocated on the heap: escapes at line 25
 	returnIntSlice(s3)
 
 	useSlice(make([]int, getUnknownNumber())) // OUT: object allocated on the heap: size is not constant
 
 	s4 := make([]byte, 300) // OUT: object allocated on the heap: object size 300 exceeds maximum stack allocation size 256
 	readByteSlice(s4)
 
-	s5 := make([]int, 4) // OUT: object allocated on the heap: escapes at line 27
+	s5 := make([]int, 4) // OUT: object allocated on the heap: escapes at line 33
 	_ = append(s5, 5)
 
 	s6 := make([]int, 3)
 	s7 := []int{1, 2, 3}
 	copySlice(s6, s7)
 
-	c1 := getComplex128() // OUT: object allocated on the heap: escapes at line 34
+	c1 := getComplex128() // OUT: object allocated on the heap: escapes at line 40
 	useInterface(c1)
 
 	n3 := 5
 	func() int {
 		return n3
 	}()
 
-	callVariadic(3, 5, 8) // OUT: object allocated on the heap: escapes at line 41
+	callVariadic(3, 5, 8) // OUT: object allocated on the heap: escapes at line 47
 
-	s8 := []int{3, 5, 8} // OUT: object allocated on the heap: escapes at line 44
+	s8 := []int{3, 5, 8} // OUT: object allocated on the heap: escapes at line 50
 	callVariadic(s8...)
 
-	n4 := 3 // OUT: object allocated on the heap: escapes at line 48
-	n5 := 7 // OUT: object allocated on the heap: escapes at line 48
+	n4 := 3 // OUT: object allocated on the heap: escapes at line 54
+	n5 := 7 // OUT: object allocated on the heap: escapes at line 54
 	func() {
 		n4 = n5
 	}()
@@ -58,6 +64,19 @@ func main() {
 	var rbuf [5]rune
 	s = string(rbuf[:])
 	println(s)
+
+	// Unsafe usage of DMA buffers: the compiler thinks this buffer won't be
+	// used anymore after the volatile store.
+	var dmaBuf1 [4]byte
+	pseudoVolatile.Set(uint32(unsafeNoEscape(unsafe.Pointer(&dmaBuf1[0]))))
+
+	// Safe usage of DMA buffers: keep the buffer alive until it is no longer
+	// needed, but don't mark it as needing to be heap allocated. The compiler
+	// will keep the buffer stack allocated if possible.
+	var dmaBuf2 [4]byte
+	pseudoVolatile.Set(uint32(unsafeNoEscape(unsafe.Pointer(&dmaBuf2[0]))))
+	// ...use the buffer in the DMA peripheral
+	keepAliveNoEscape(unsafe.Pointer(&dmaBuf2[0]))
 }
 
 func derefInt(x *int) int {
@@ -93,3 +112,13 @@ func useInterface(interface{})
 func callVariadic(...int)
 
 func useSlice([]int)
+
+// See the function with the same name in the machine package.
+//
+//go:linkname unsafeNoEscape machine.unsafeNoEscape
+func unsafeNoEscape(ptr unsafe.Pointer) uintptr
+
+//go:linkname keepAliveNoEscape machine.keepAliveNoEscape
+func keepAliveNoEscape(ptr unsafe.Pointer)
+
+var pseudoVolatile volatile.Register32