Integrating custos (device: CPU) #234

What are the advantages of using custos' device let device = CPU::new(); in allocation compared to usual allocation?

• This keeps the doors open for supporting other accelerators (e.g. GPUs) at an advanced state of the operations implementer (in this case, probably linalg).
  For instance, someone could just start with a CPU implementation, and if needed, add a CUDA implementation later:

let device = CPU::new();
let a = Matrix::from((&device, 2, 3, [1., 2., 3., 4., 5., 6.]));
let b = Matrix::from((&device, 2, 3, [1., 2., 3., 4., 5., 6.]));
let c_cpu = a + b;
let device = CUDA::new(0)?;
let a = Matrix::from((&device, 2, 3, [1., 2., 3., 4., 5., 6.]));
let b = Matrix::from((&device, 2, 3, [1., 2., 3., 4., 5., 6.]));
let c_cuda = a + b;
assert_eq!(c_cpu.read(), c_cuda.read());

It should be possibly to create a struct that returns a device depending on the compilation target and activated feature-set.

• The specific device keeps track on allocations, which effectively creates a graph.
  New nodes are added by calling let c = device.retrieve(len, (&a, &b));. (This allocates new memory, or returns previously allocated memory, if custos was configured to 'cache' old allocations). The node is stored in Buffer.

• It is technically possible to allocate Buffer or Matrix via the typical way as a device can be stored in a static variable.
  Look here.

let buf = buf![2, 3, 6].to_gpu();

How can we use custos Buffer for arrays? what would be the advantages?

• Buffer is designed to be 'modular'. A Buffer can be allocated with any struct that implements the Device trait. This is also beneficial if accelerator support is necessary at some point.

pub struct Buffer<'a, T = f32, D: Device = CPU, const N: usize = 0> {
 pub ptr: D::Ptr<T, N>,
 pub len: usize,
 pub device: Option<&'a D>,
 pub flag: BufFlag,
 pub node: Node,
}

The const generic optionally enables a Buffer to be allocated on the stack.

Thanks, this is really helpful.

I see a lot of interesting principles, we currently have a way of providing bindings for external libraries. For example the trait smartcore::linalg::arrays:Array could be implemented for Buffer. This way all the basic methods can work with a buffer type.
The idea would be to have:

let buf = buf![2, 3, 6];
let el = buf.get(0);
println!("{}", el);
>>> 2
let a = buf!([1, 2, 3]);
let v: Vec<i32> = a.iterator(0).map(|&v| v).collect();
assert_eq!(v, vec!(1, 2, 3));
let mut a = buf!([1, 2, 3]);
a.iterator_mut(0).for_each(|v| *v = 1);
assert_eq!(a, buf!([1, 1, 1]));

This may be possible with something like:

impl<T: Debug + Display + Copy + Sized> Array<T, (usize, usize)>
 for Buffer<T, D>
where
 D = Device
{
 fn get(&self, pos: (usize, usize)) -> &T {
 ... return element at position ...
 }
 fn shape(&self) -> (usize, usize) {
 ... return rows, columns shape ...
 }
 fn is_empty(&self) -> bool {
 ...
 }
 fn iterator<'b>(&'b self, axis: u8) -> Box<dyn Iterator<Item = &'b T> + 'b> {
 assert!(
 axis == 1 || axis == 0,
 "For two dimensional array `axis` should be either 0 or 1"
 );
 match axis {
 0 => Box::new(self.iter()),
 _ => Box::new(
 (0..self.ncols()).flat_map(move |c| (0..self.nrows()).map(move |r| &self[[r, c]])),
 ),
 }
 }
}

This is just a generic (probably not working) example, the same can be done for custos::Matrix with smartcore::linalg::arrays::Array2.

The part I cannot see at the moment is how to instantiate a Device "automatically" without the end user needing to instantiate explicitly. let device = CPU::new(); should be instantiated under-the-hood in a prelude or something like that? What is the cost of instantiating a device? does it makes sense to consider instantiating a device for particular operations? For example, in the case of expensive operations:

x.dot(&y)

it would make sense or provide advantages to have a "scoped" device? Like:

/// this is an high-level public API called by a client or end user
fn dot(&self, x: Buffer) { // assuming self is a Buffer
 let device = CPU::new_with((self, x)); // "load" a device with existing buffers 
 device.dot().apply() // this returns the result of a dot product between the buffers in the device
 // device is dropped
}

This may look trivial with dot-product but maybe for operations on trees could do to load a tree in a device and use it only for that tree?
I am trying to consider the possibility of pluggable features from custos.

We really want to develop things consistently from public interface to allocation level, sometimes there is tension between ergonomicity and "programmability" so better think things ahead.

Please follow up with your ideas.

The part I cannot see at the moment is how to instantiate a Device "automatically" without the end user needing to instantiate explicitly
The only general way I could come up with is storing the device in a static variable and this is definitely not ideal. static-api

This should work too, however, this probably isn't what you want:

pub struct SomeMLAlgo<T, D: Device = custos::CPU> {
 device: D,
 a: Matrix<T, D>, // or whatever
 ...
}
impl<T, D: Device> SomeMLAlgo<T, D> {
 pub fn fit(a: &[T], ...) -> Self {
 let device = D::new();
 let a = Matrix::from((&device, ..., a));
 
 Self {
 device,
 a
 }
 }
}

What is the cost of instantiating a device?
Instantiating a CPU is cheap. Instantiating, for example, a CUDA device is expensive.

/// this is an high-level public API called by a client or end user
fn dot(&self, x: Buffer) { // assuming self is a Buffer
 let device = CPU::new_with((self, x)); // "load" a device with existing buffers 
 device.dot().apply() // this returns the result of a dot product between the buffers in the device
 // device is dropped
}

In this case I would probably do something like this:

fn dot(&self, x: Buffer) {
 self.device().dot(self, x)
}
``

About GPUs: are you interested in doing something with triton? Do you see something like that to be a good fit for a Rust integration? If yes, how would you integrate kernels written in Triton into a Rust library?

cc: @morenol @Steboss89 please contribute to the conversation if you have any idea

@Mec-iS triton is really interesting, I love the idea, but IMHO it may be an overkill at the moment. In which cases are we going to use triton?