luma_tensor/ops/
matmul.rs

1use crate::{DTypeKind, Device, Float, Int, Layout, Shape, Storage, Tensor, TensorMeta};
2
3pub trait MatmulDTypeKind<D: Device>: DTypeKind<D> + Sized {
4    fn matmul_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: &Self::Storage, rhs_l: &Layout) -> crate::Result<(Self::Storage, Shape)>;
5}
6
7impl<D: Device> MatmulDTypeKind<D> for Float {
8    #[inline]
9    fn matmul_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: &Self::Storage, rhs_l: &Layout) -> crate::Result<(Self::Storage, Shape)> {
10        D::f_matmul(lhs, lhs_l, rhs, rhs_l)
11    }
12}
13
14impl<D: Device> MatmulDTypeKind<D> for Int {
15    #[inline]
16    fn matmul_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: &Self::Storage, rhs_l: &Layout) -> crate::Result<(Self::Storage, Shape)> {
17        D::i_matmul(lhs, lhs_l, rhs, rhs_l)
18    }
19}
20
21impl<D: Device, K: MatmulDTypeKind<D>> Tensor<D, K> {
22    pub fn matmul(&self, rhs: &Self) -> crate::Result<Self> {
23        let (storage, shape) = K::matmul_dispatch(&*self.storage_read()?, self.layout(), &*rhs.storage_read()?, rhs.layout())?;
24        let meta = K::Meta::on_matmul(self, rhs);
25        assert_eq!(self.dtype(), storage.dtype());
26        Ok(Self::from_storage(storage, shape, meta))
27    }
28}
luma_tensor/ops/matmul.rs

luma_tensor/ops/
matmul.rs
