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"):
@@ -121,6 +123,20 @@ func (b *builder) createKeepAliveImpl() {
 	b.CreateRetVoid()
 }
 
+// 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.