#tensor #gpu #compute

tensor_compute

A Tensor computing library for Rust

1 unstable release

0.1.0 Aug 19, 2020

#93 in #tensor

MIT/Apache

2MB
2.5K SLoC

Rust 2.5K SLoC // 0.1% comments GLSL 278 SLoC // 0.1% comments

Contains (Mach-o exe, 8.5MB) src/glslc

Rust Tensor Compute

This is a personal project to teach myself WebGPU computing, focused
on Machine Learning application.

Features for now:

  • Select which GPU to use (if more than 1 in system)
  • Clone
  • (Batch) Matmul
  • Relu
  • Transpose
  • Fill
  • Compare
  • Make Contiguous
  • Slice (Needs more design/work)
  • Index (Needs more design/work)
  • Create Views Tensor

In progress:

Element Wise Tensor - Tensor Ops:

  • Add
  • Sub
  • Mul

Scalar - Tensor Ops:

  • Add
  • Sub
  • Mul

Working Example:

fn main() {
    println!("Running in {:?}", GpuStore::get_default().info());
    let ma = Tensor::from_data_and_shape(vec![1., 2., 3., 4., 5., 6., 7., 8.], vec![2, 2, 2]);
    let mb = Tensor::from_data_and_shape(vec![2., 3., 4., 5.], vec![2, 2]);
    let result = ma.matmul(&mb);
    println!("{:?}", result);
    /*
    Running in AdapterInfo { name: "AMD Radeon Pro 560", vendor: 0, device: 0, device_type: DiscreteGpu, backend: Metal }
    Shape: [2, 2, 2]
    [[[ 10  13 ]
      [ 22  29 ]]

     [[ 34  45 ]
      [ 46  61 ]]]
    */
}

Short-term goals

  • How should indexing and assignment work?
  • Finish public API docs
  • More Tests

Dependencies

~10–18MB
~337K SLoC