hotg-ai / rune

Rune provides containers to encapsulate and deploy edgeML pipelines and applications
Apache License 2.0
136 stars 15 forks source link

Remove Copies and Allocations from the Tensor Definition #287

Open Michael-F-Bryan opened 3 years ago

Michael-F-Bryan commented 3 years ago

Quoting directly from https://github.com/hotg-ai/rune/issues/242#issuecomment-906305787


Remove Copies from the Tensor Definition

The current definition for a "tensor" inside a Rune contains a reference-counted buffer (on the heap) with methods that ensure copy-on-write semantics and do the appropriate indexing math. We allocate a new set of buffers on every run through the pipeline.

struct Tensor<T> {
  dimensions: vec![usize],
  elements: Arc<[T]>,
}

(src)

To avoid excess allocations and give us finer control over allocations/mutation (with copy-on-write you don't necessarily know when it'll make a copy or when it'll edit in place) I think we should change the definition to

struct Tensor<'a, T> {
  dimensions: &'a [usize],
  elements: &'a [T],
}

struct TensorMut<'a, T> {
  dimensions: &'a [usize],
  elements: &'a mut [T],
}

That way we can move the backing arrays to static variables.

Michael-F-Bryan commented 3 years ago

Part of this ticket is moving all tensor buffers to static variables.

I imagine we'll need to update the hotg-rune-compiler crate to generate something like this:

static mut BUFFERS: Buffers = ...;
static mut PIPELINE: Option<Pipeline> = None;

struct Buffers {
  buffer_0: [f32; 1*256*256*3],
  buffer_1: [u8: 1960],
}

struct Pipeline {
  image: Capability<u8>,
  model: Model<f32, f32>,
}

impl Pipeline {
  pub fn call(&mut self, buffers: &mut Buffers) { 
    self.image.generate(
      OutputTensor::new(&mut buffers.buffer_0, &[1, 256, 256, 3]),
    );
    self.mode.infer(
      InputTensor::new(&buffers.buffer_0, &[1, 256, 256, 3]),
      OutputTensor::new(&mut buffers.buffer_1, &[1960]),
    );
  }
}

pub extern "C" _manifest() {
  let image = ...;
  let model = ...;

  unsafe {
    PIPELINE = Some(Pipeline { image, model });
  }
}

pub extern "C" _call() {
  unsafe {
    let pipeline = PIPELINE.as_mut().expect("The pipeline wasn't initialized");
    pipeline.call(&mut BUFFERS);
  }
}

(note: I'm using the names InputTensor and OutputTensor instead of Tensor and TensorMut because it more accurately represents what they are for)