varunsh-xilinx commented 1 year ago

Summary of Changes

Allow using custom memory allocation schemes

Closes #162

Motivation

Some workers use custom classes to hold incoming data for inference. By having the batcher store the data in the right format, we can make an assumption in all workers that the incoming data will be in a compatible and valid format that's ready to use.

Implementation

The memory pool now expects an input tensor rather than a flat numeric size to indicate how much memory to allocate. The basic allocator uses the tensor to get the flat memory space still (this logic used to be on the batcher side) but if an allocator needs the raw information about the tensor, then it can use it.

The VartTensorBuffer is back in slightly modified form and it serves the same purpose as before but now it's backed by the allocator rather than owning the memory itself.

Notes

One gotcha is that since the allocators may be creating, freeing and merging blocks of memory, the state of the memory needs to be kept in lists so that these modifications don't invalidate existing pointers to other parts of the structure that are in use elsewhere.

gbuildx commented 1 year ago

Build failed!