From c0e817263fc1d494dc37426183c31724cd08a224 Mon Sep 17 00:00:00 2001 From: Vivek Panyam Date: Wed, 18 Oct 2023 23:58:58 -0400 Subject: [PATCH] Add an `Allocatable` trait to the runner interface --- .../src/do_not_modify/alloc.rs | 13 +++++- .../src/do_not_modify/alloc_inline.rs | 46 +++++++++---------- .../src/do_not_modify/alloc_shm.rs | 32 +++++++------ .../src/do_not_modify/types.rs | 11 ++++- source/carton-runner-interface/src/runner.rs | 14 +++--- source/carton-runner-wasm/Cargo.toml | 4 +- source/carton-runner-wasm/src/types.rs | 5 +- 7 files changed, 73 insertions(+), 52 deletions(-) diff --git a/source/carton-runner-interface/src/do_not_modify/alloc.rs b/source/carton-runner-interface/src/do_not_modify/alloc.rs index d740a7b..e2c23de 100644 --- a/source/carton-runner-interface/src/do_not_modify/alloc.rs +++ b/source/carton-runner-interface/src/do_not_modify/alloc.rs @@ -31,8 +31,17 @@ pub trait AsPtr { fn as_mut_ptr(&mut self) -> *mut T; } -pub trait TypedAlloc { - type Output: AsPtr; +pub trait Allocator { + type Output; +} +pub trait TypedAlloc: Allocator +where + Self::Output: AsPtr, +{ fn alloc(&self, numel: usize) -> Self::Output; } + +pub trait AllocatableBy: Sized { + fn alloc(allocator: &A, numel: usize) -> A::Output; +} diff --git a/source/carton-runner-interface/src/do_not_modify/alloc_inline.rs b/source/carton-runner-interface/src/do_not_modify/alloc_inline.rs index fb1edc6..89e1881 100644 --- a/source/carton-runner-interface/src/do_not_modify/alloc_inline.rs +++ b/source/carton-runner-interface/src/do_not_modify/alloc_inline.rs @@ -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, }; @@ -79,12 +79,14 @@ impl AsPtr 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 let size_bytes = numel * std::mem::size_of::<$RustType>(); @@ -101,8 +103,6 @@ for_each_numeric_carton_type! { } impl TypedAlloc for InlineAllocator { - type Output = InlineTensorStorage; - fn alloc(&self, numel: usize) -> Self::Output { let out = if !self.use_pool { vec![String::default(); numel].into() @@ -115,10 +115,8 @@ impl TypedAlloc for InlineAllocator { } // Copy the data -impl From> - for TensorStorage -where - InlineAllocator: TypedAlloc, +impl> + From> for TensorStorage { fn from(view: ndarray::ArrayViewD<'_, T>) -> Self { // Alloc a tensor @@ -152,17 +150,14 @@ impl From> for TensorStorage( +pub fn alloc_tensor_no_pool>( shape: Vec, -) -> TensorStorage -where - InlineAllocator: TypedAlloc, -{ +) -> TensorStorage { static POOL_ALLOCATOR: Lazy = Lazy::new(|| InlineAllocator::without_pool()); let numel = shape.iter().product::().max(1) as usize; - let data = >::alloc(&POOL_ALLOCATOR, numel); + let data = T::alloc(&POOL_ALLOCATOR, numel); TensorStorage { data, @@ -172,15 +167,14 @@ where } } -pub fn alloc_tensor(shape: Vec) -> TensorStorage -where - InlineAllocator: TypedAlloc, -{ +pub fn alloc_tensor>( + shape: Vec, +) -> TensorStorage { static POOL_ALLOCATOR: Lazy = Lazy::new(|| InlineAllocator::new()); let numel = shape.iter().product::().max(1) as usize; - let data = >::alloc(&POOL_ALLOCATOR, numel); + let data = T::alloc(&POOL_ALLOCATOR, numel); TensorStorage { data, @@ -190,11 +184,17 @@ where } } -impl TensorStorage +impl> TensorStorage { + pub fn new(shape: Vec) -> TensorStorage { + alloc_tensor(shape) + } +} + +impl AllocatableBy for T where InlineAllocator: TypedAlloc, { - pub fn new(shape: Vec) -> TensorStorage { - alloc_tensor(shape) + fn alloc(allocator: &InlineAllocator, numel: usize) -> InlineTensorStorage { + >::alloc(allocator, numel) } } diff --git a/source/carton-runner-interface/src/do_not_modify/alloc_shm.rs b/source/carton-runner-interface/src/do_not_modify/alloc_shm.rs index 3aa749b..041e45e 100644 --- a/source/carton-runner-interface/src/do_not_modify/alloc_shm.rs +++ b/source/carton-runner-interface/src/do_not_modify/alloc_shm.rs @@ -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, }; @@ -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>(); @@ -282,8 +284,6 @@ for_each_numeric_carton_type! { } impl TypedAlloc for SHMAllocator { - type Output = SHMTensorStorage; - fn alloc(&self, numel: usize) -> Self::Output { let out = if !self.use_pool { vec![String::default(); numel].into() @@ -384,17 +384,15 @@ impl From> for TensorStorage( +pub fn alloc_tensor_no_pool>( shape: Vec, ) -> TensorStorage -where - SHMAllocator: TypedAlloc, { static POOL_ALLOCATOR: Lazy = Lazy::new(|| SHMAllocator::without_pool()); let numel = shape.iter().product::().max(1) as usize; - let data = >::alloc(&POOL_ALLOCATOR, numel); + let data = T::alloc(&POOL_ALLOCATOR, numel); TensorStorage { data, @@ -404,15 +402,13 @@ where } } -pub fn alloc_tensor(shape: Vec) -> TensorStorage -where - SHMAllocator: TypedAlloc, +pub fn alloc_tensor>(shape: Vec) -> TensorStorage { static POOL_ALLOCATOR: Lazy = Lazy::new(|| SHMAllocator::new()); let numel = shape.iter().product::().max(1) as usize; - let data = >::alloc(&POOL_ALLOCATOR, numel); + let data = T::alloc(&POOL_ALLOCATOR, numel); TensorStorage { data, @@ -421,3 +417,13 @@ where pd: PhantomData, } } + + +impl AllocatableBy for T +where + SHMAllocator: TypedAlloc, +{ + fn alloc(allocator: &SHMAllocator, numel: usize) -> SHMTensorStorage { + >::alloc(allocator, numel) + } +} diff --git a/source/carton-runner-interface/src/do_not_modify/types.rs b/source/carton-runner-interface/src/do_not_modify/types.rs index d20977b..92e56f6 100644 --- a/source/carton-runner-interface/src/do_not_modify/types.rs +++ b/source/carton-runner-interface/src/do_not_modify/types.rs @@ -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 { @@ -247,6 +251,9 @@ for_each_carton_type! { pub type TensorStorage = super::storage::TensorStorage; +pub trait Allocatable: AllocatableBy {} +impl Allocatable for T where T: AllocatableBy {} + for_each_carton_type! { $( impl From> for Tensor { diff --git a/source/carton-runner-interface/src/runner.rs b/source/carton-runner-interface/src/runner.rs index 1b33482..e4eabb3 100644 --- a/source/carton-runner-interface/src/runner.rs +++ b/source/carton-runner-interface/src/runner.rs @@ -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; @@ -299,9 +295,11 @@ impl Runner { } } - pub fn alloc_tensor(&self, shape: Vec) -> Result + pub fn alloc_tensor( + &self, + shape: Vec, + ) -> Result where - InlineAllocator: TypedAlloc, Tensor: From>, { Ok(TensorStorage::new(shape).into()) diff --git a/source/carton-runner-wasm/Cargo.toml b/source/carton-runner-wasm/Cargo.toml index 6183a72..7549baa 100644 --- a/source/carton-runner-wasm/Cargo.toml +++ b/source/carton-runner-wasm/Cargo.toml @@ -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 } @@ -28,4 +27,5 @@ semver = "1.0.20" [dev-dependencies] escargot = "0.5.8" paste = "1.0.14" -tempfile = "3.8.0" \ No newline at end of file +tempfile = "3.8.0" +carton = { path = "../carton" } \ No newline at end of file diff --git a/source/carton-runner-wasm/src/types.rs b/source/carton-runner-wasm/src/types.rs index 47d6ed1..3ad4109 100644 --- a/source/carton-runner-wasm/src/types.rs +++ b/source/carton-runner-wasm/src/types.rs @@ -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};