MLShapedArray, MLMultiArray, memory allocations and memory layout

Question

thomastapeit OP

Created Nov ’22

Replies 1

Boosts 0

Views 1.6k

Participants 2

I have two questions on MLShapedArray.

1. Memory Layout

I'm using MLMultiArray to get data into and out of a CoreML model. I need to preprocess the data before feeding it to the model, and right now I store it in a ContiguousArray since I know that I can safely pass this to vDSP methods.

I'm wondering if I could use an MLShapedArray instead. Is an MLShapedArray guaranteed to have a contiguous memory buffer underneath?

2. Memory Allocations

MLShapedArray and MLMultiArray have initializers that allow converting between them. Is data copied, or is the underlying buffer reused?

I'd love fo the buffer to be reused to avoid malloc calls. In my ML pipeline, I'd like to allocate all buffers at the start and then just reuse them as I do my processing.

Boost

Answer 1

Frameworks Engineer OP

Apple

Nov ’22

Accepted Answer

Memory Layout

Neither MLShapedArray nor MLMultiArray guarantees the memory layout of the backing buffer when they allocate it.

The initializer init(bytesNoCopy:shape:strides:deallocator:) lets client specify a backing buffer and its memory layout (i.e. strides). With this, you can allocate a contiguous buffer, preprocess it with vDSP, and wrap it with MLShapedArray before sending to a model.

// Allocate a contiguous buffer
let rawBuffer = UnsafeMutableRawBufferPointer.allocate(byteCount: 3 * 2 * MemoryLayout<Float32>.size, alignment: 1)

// Preprocess the buffer ...
rawBuffer.withMemoryRebound(to: Float32.self) { buffer in
    buffer.assign(repeating: 3.14)
}

// Wrap the buffer with MLShapedArray
var shapedArray = MLShapedArray<Float32>(bytesNoCopy: rawBuffer.baseAddress!, shape: [3, 2], strides: [2, 1], deallocator: .custom({ _, _ in withExtendedLifetime(rawBuffer) {} }))

// Send to CoreML

Memory Allocation

Conversion between MLShapedArray and MLMultiArray is zero-copy.

1