Skip to content

Commit

Permalink
Add an Allocatable trait to the runner interface (#179)
Browse files Browse the repository at this point in the history
This PR adds an `Allocatable` trait to the runner interface. This
generally makes it cleaner to write generic functions over Carton types
without needing to reference implementation details in trait bounds
(e.g. `where InlineAllocator: TypedAlloc<T, Output =
InlineTensorStorage>`).

Note: this PR does touch files inside `do_not_modify`, but it does so in
a way that does not affect the wire protocol.
  • Loading branch information
VivekPanyam authored Dec 27, 2023
1 parent 60afa98 commit 4786140
Show file tree
Hide file tree
Showing 7 changed files with 73 additions and 52 deletions.
13 changes: 11 additions & 2 deletions source/carton-runner-interface/src/do_not_modify/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,17 @@ pub trait AsPtr<T> {
fn as_mut_ptr(&mut self) -> *mut T;
}

pub trait TypedAlloc<T> {
type Output: AsPtr<T>;
pub trait Allocator {
type Output;
}

pub trait TypedAlloc<T>: Allocator
where
Self::Output: AsPtr<T>,
{
fn alloc(&self, numel: usize) -> Self::Output;
}

pub trait AllocatableBy<A: Allocator>: Sized {
fn alloc(allocator: &A, numel: usize) -> A::Output;
}
46 changes: 23 additions & 23 deletions source/carton-runner-interface/src/do_not_modify/alloc_inline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};

use super::{
alloc::{AsPtr, NumericTensorType, TypedAlloc},
alloc::{AllocatableBy, Allocator, AsPtr, NumericTensorType, TypedAlloc},
alloc_pool::{PoolAllocator, PoolItem},
storage::TensorStorage,
};
Expand Down Expand Up @@ -79,12 +79,14 @@ impl<T> AsPtr<T> for InlineTensorStorage {
}
}

impl Allocator for InlineAllocator {
type Output = InlineTensorStorage;
}

for_each_numeric_carton_type! {
$(
/// We're using a macro here instead of a generic impl because rust gives misleading error messages otherwise.
impl TypedAlloc<$RustType> for InlineAllocator {
type Output = InlineTensorStorage;

fn alloc(&self, numel: usize) -> Self::Output {
// We need to convert to size_bytes since we always use a Vec<u8>
let size_bytes = numel * std::mem::size_of::<$RustType>();
Expand All @@ -101,8 +103,6 @@ for_each_numeric_carton_type! {
}

impl TypedAlloc<String> for InlineAllocator {
type Output = InlineTensorStorage;

fn alloc(&self, numel: usize) -> Self::Output {
let out = if !self.use_pool {
vec![String::default(); numel].into()
Expand All @@ -115,10 +115,8 @@ impl TypedAlloc<String> for InlineAllocator {
}

// Copy the data
impl<T: NumericTensorType + Default + Copy> From<ndarray::ArrayViewD<'_, T>>
for TensorStorage<T, InlineTensorStorage>
where
InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
impl<T: NumericTensorType + Default + Copy + AllocatableBy<InlineAllocator>>
From<ndarray::ArrayViewD<'_, T>> for TensorStorage<T, InlineTensorStorage>
{
fn from(view: ndarray::ArrayViewD<'_, T>) -> Self {
// Alloc a tensor
Expand Down Expand Up @@ -152,17 +150,14 @@ impl From<ndarray::ArrayViewD<'_, String>> for TensorStorage<String, InlineTenso

// Allocates a contiguous tensor with a shape and type
#[cfg(feature = "benchmark")]
pub fn alloc_tensor_no_pool<T: Default + Clone>(
pub fn alloc_tensor_no_pool<T: Default + Clone + AllocatableBy<InlineAllocator>>(
shape: Vec<u64>,
) -> TensorStorage<T, InlineTensorStorage>
where
InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
{
) -> TensorStorage<T, InlineTensorStorage> {
static POOL_ALLOCATOR: Lazy<InlineAllocator> = Lazy::new(|| InlineAllocator::without_pool());

let numel = shape.iter().product::<u64>().max(1) as usize;

let data = <InlineAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
let data = T::alloc(&POOL_ALLOCATOR, numel);

TensorStorage {
data,
Expand All @@ -172,15 +167,14 @@ where
}
}

pub fn alloc_tensor<T: Default + Clone>(shape: Vec<u64>) -> TensorStorage<T, InlineTensorStorage>
where
InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
{
pub fn alloc_tensor<T: Default + Clone + AllocatableBy<InlineAllocator>>(
shape: Vec<u64>,
) -> TensorStorage<T, InlineTensorStorage> {
static POOL_ALLOCATOR: Lazy<InlineAllocator> = Lazy::new(|| InlineAllocator::new());

let numel = shape.iter().product::<u64>().max(1) as usize;

let data = <InlineAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
let data = T::alloc(&POOL_ALLOCATOR, numel);

TensorStorage {
data,
Expand All @@ -190,11 +184,17 @@ where
}
}

impl<T: Default + Clone> TensorStorage<T, InlineTensorStorage>
impl<T: Default + Clone + AllocatableBy<InlineAllocator>> TensorStorage<T, InlineTensorStorage> {
pub fn new(shape: Vec<u64>) -> TensorStorage<T, InlineTensorStorage> {
alloc_tensor(shape)
}
}

impl<T> AllocatableBy<InlineAllocator> for T
where
InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
{
pub fn new(shape: Vec<u64>) -> TensorStorage<T, InlineTensorStorage> {
alloc_tensor(shape)
fn alloc(allocator: &InlineAllocator, numel: usize) -> InlineTensorStorage {
<InlineAllocator as TypedAlloc<T>>::alloc(allocator, numel)
}
}
32 changes: 19 additions & 13 deletions source/carton-runner-interface/src/do_not_modify/alloc_shm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use dashmap::DashMap;
use once_cell::sync::Lazy;

use super::{
alloc::{AsPtr, NumericTensorType, TypedAlloc},
alloc::{AsPtr, NumericTensorType, TypedAlloc, Allocator, AllocatableBy},
alloc_pool::{AllocItem, PoolAllocator, PoolItem},
storage::TensorStorage,
};
Expand Down Expand Up @@ -260,12 +260,14 @@ impl SHMAllocator {
}
}

impl Allocator for SHMAllocator {
type Output = SHMTensorStorage;
}

for_each_numeric_carton_type! {
$(
/// We're using a macro here instead of a generic impl because rust gives misleading error messages otherwise.
impl TypedAlloc<$RustType> for SHMAllocator {
type Output = SHMTensorStorage;

fn alloc(&self, numel: usize) -> Self::Output {
// We need to convert to size_bytes
let size_bytes = numel * std::mem::size_of::<$RustType>();
Expand All @@ -282,8 +284,6 @@ for_each_numeric_carton_type! {
}

impl TypedAlloc<String> for SHMAllocator {
type Output = SHMTensorStorage;

fn alloc(&self, numel: usize) -> Self::Output {
let out = if !self.use_pool {
vec![String::default(); numel].into()
Expand Down Expand Up @@ -384,17 +384,15 @@ impl From<ndarray::ArrayViewD<'_, String>> for TensorStorage<String, SHMTensorSt

// Allocates a contiguous tensor with a shape and type
#[cfg(feature = "benchmark")]
pub fn alloc_tensor_no_pool<T: Default + Clone>(
pub fn alloc_tensor_no_pool<T: Default + Clone + AllocatableBy<SHMAllocator>>(
shape: Vec<u64>,
) -> TensorStorage<T, SHMTensorStorage>
where
SHMAllocator: TypedAlloc<T, Output = SHMTensorStorage>,
{
static POOL_ALLOCATOR: Lazy<SHMAllocator> = Lazy::new(|| SHMAllocator::without_pool());

let numel = shape.iter().product::<u64>().max(1) as usize;

let data = <SHMAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
let data = T::alloc(&POOL_ALLOCATOR, numel);

TensorStorage {
data,
Expand All @@ -404,15 +402,13 @@ where
}
}

pub fn alloc_tensor<T: Default + Clone>(shape: Vec<u64>) -> TensorStorage<T, SHMTensorStorage>
where
SHMAllocator: TypedAlloc<T, Output = SHMTensorStorage>,
pub fn alloc_tensor<T: Default + Clone + AllocatableBy<SHMAllocator>>(shape: Vec<u64>) -> TensorStorage<T, SHMTensorStorage>
{
static POOL_ALLOCATOR: Lazy<SHMAllocator> = Lazy::new(|| SHMAllocator::new());

let numel = shape.iter().product::<u64>().max(1) as usize;

let data = <SHMAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
let data = T::alloc(&POOL_ALLOCATOR, numel);

TensorStorage {
data,
Expand All @@ -421,3 +417,13 @@ where
pd: PhantomData,
}
}


impl<T> AllocatableBy<SHMAllocator> for T
where
SHMAllocator: TypedAlloc<T, Output = SHMTensorStorage>,
{
fn alloc(allocator: &SHMAllocator, numel: usize) -> SHMTensorStorage {
<SHMAllocator as TypedAlloc<T>>::alloc(allocator, numel)
}
}
11 changes: 9 additions & 2 deletions source/carton-runner-interface/src/do_not_modify/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,15 @@
// See the License for the specific language governing permissions and
// limitations under the License.

pub use carton_macros::for_each_carton_type;
pub use carton_macros::{for_each_carton_type, for_each_numeric_carton_type};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use super::{alloc_inline::InlineTensorStorage, comms::Comms};
use super::{
alloc::AllocatableBy,
alloc_inline::{InlineAllocator, InlineTensorStorage},
comms::Comms,
};

#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct RPCRequest {
Expand Down Expand Up @@ -247,6 +251,9 @@ for_each_carton_type! {

pub type TensorStorage<T> = super::storage::TensorStorage<T, InlineTensorStorage>;

pub trait Allocatable: AllocatableBy<InlineAllocator> {}
impl<T> Allocatable for T where T: AllocatableBy<InlineAllocator> {}

for_each_carton_type! {
$(
impl From<TensorStorage<$RustType>> for Tensor {
Expand Down
14 changes: 6 additions & 8 deletions source/carton-runner-interface/src/runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,8 @@ use std::{collections::HashMap, sync::Arc};
use crate::{
client::Client,
do_not_modify::comms::OwnedComms,
do_not_modify::{
alloc::TypedAlloc,
alloc_inline::{InlineAllocator, InlineTensorStorage},
types::{Device, RPCRequestData, RPCResponseData, SealHandle, Tensor},
},
types::{Handle, RunnerOpt, TensorStorage},
do_not_modify::types::{Device, RPCRequestData, RPCResponseData, SealHandle, Tensor},
types::{Allocatable, Handle, RunnerOpt, TensorStorage},
};

use futures::Stream;
Expand Down Expand Up @@ -299,9 +295,11 @@ impl Runner {
}
}

pub fn alloc_tensor<T: Clone + Default>(&self, shape: Vec<u64>) -> Result<Tensor, String>
pub fn alloc_tensor<T: Clone + Default + Allocatable>(
&self,
shape: Vec<u64>,
) -> Result<Tensor, String>
where
InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
Tensor: From<TensorStorage<T>>,
{
Ok(TensorStorage::new(shape).into())
Expand Down
4 changes: 2 additions & 2 deletions source/carton-runner-wasm/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ publish = false
exclude = ["tests/test_model", "carton-wasm-interface"]

[dependencies]
carton = { path = "../carton" }
carton-runner-interface = { path = "../carton-runner-interface" }
color-eyre = "0.6.2"
lunchbox = { version = "0.1", default-features = false }
Expand All @@ -28,4 +27,5 @@ semver = "1.0.20"
[dev-dependencies]
escargot = "0.5.8"
paste = "1.0.14"
tempfile = "3.8.0"
tempfile = "3.8.0"
carton = { path = "../carton" }
5 changes: 3 additions & 2 deletions source/carton-runner-wasm/src/types.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
use color_eyre::eyre::{ensure, eyre};
use color_eyre::{Report, Result};

use carton::types::for_each_numeric_carton_type;
use carton_runner_interface::types::{Tensor as CartonTensor, TensorStorage as CartonStorage};
use carton_runner_interface::types::{
for_each_numeric_carton_type, Tensor as CartonTensor, TensorStorage as CartonStorage,
};

use crate::component::{Dtype, Tensor as WasmTensor, TensorNumeric, TensorString};

Expand Down

0 comments on commit 4786140

Please sign in to comment.