Get the cpu backend to compile.

This commit is contained in:
laurent
2023-06-28 14:12:38 +01:00
parent 54a6c40f27
commit 14449ff80c
5 changed files with 44 additions and 59 deletions

View File

@ -1,5 +1,5 @@
use crate::op::{BinaryOp, UnaryOp};
use crate::{DType, Error, Layout, Result, Shape, StridedIndex};
use crate::{DType, Error, Layout, Result, Shape};
use gemm::{gemm, Parallelism};
use half::{bf16, f16};
@ -81,14 +81,13 @@ fn unary_map<T: Copy, U: Copy, F: FnMut(T) -> U>(vs: &[T], layout: &Layout, mut
// This function maps over two strided index sequences.
fn binary_map<T: Copy, F: FnMut(T, T) -> T>(
shape: &Shape,
lhs_layout: &Layout,
rhs_layout: &Layout,
lhs: &[T],
rhs: &[T],
mut f: F,
) -> Vec<T> {
let dims = shape.dims();
let shape = lhs_layout.shape();
if lhs_layout.is_contiguous() && rhs_layout.is_contiguous() {
(0..shape.elem_count()).map(|i| f(lhs[i], rhs[i])).collect()
} else {
@ -148,17 +147,19 @@ fn copy_strided_src_<T: Copy + std::fmt::Display>(
}
}
fn matmul_impl<T: 'static + num_traits::Num + Copy>(
fn matmul<T: 'static + num_traits::Num + Copy>(
lhs: &[T],
rhs: &[T],
(b, m, n, k): (usize, usize, usize, usize),
lhs_stride: &[usize],
rhs_stride: &[usize],
lhs_layout: &Layout,
rhs_layout: &Layout,
) -> Result<Vec<T>> {
let a_skip: usize = m * k;
let b_skip: usize = n * k;
let c_skip: usize = m * n;
let lhs_stride = lhs_layout.stride();
let rhs_stride = rhs_layout.stride();
let rank = lhs_stride.len();
let lhs_cs = lhs_stride[rank - 1];
let lhs_rs = lhs_stride[rank - 2];
@ -512,29 +513,28 @@ impl CpuStorage {
pub(crate) fn binary_impl<B: BinaryOp>(
&self,
rhs: &Self,
shape: &Shape,
lhs_layout: &Layout,
rhs_layout: &Layout,
) -> Result<Self> {
match (self, rhs) {
(Self::BF16(lhs), Self::BF16(rhs)) => {
let data = binary_map(shape, lhs_layout, rhs_layout, lhs, rhs, B::bf16);
let data = binary_map(lhs_layout, rhs_layout, lhs, rhs, B::bf16);
Ok(Self::BF16(data))
}
(Self::F16(lhs), Self::F16(rhs)) => {
let data = binary_map(shape, lhs_layout, rhs_layout, lhs, rhs, B::f16);
let data = binary_map(lhs_layout, rhs_layout, lhs, rhs, B::f16);
Ok(Self::F16(data))
}
(Self::F32(lhs), Self::F32(rhs)) => {
let data = binary_map(shape, lhs_layout, rhs_layout, lhs, rhs, B::f32);
let data = binary_map(lhs_layout, rhs_layout, lhs, rhs, B::f32);
Ok(Self::F32(data))
}
(Self::F64(lhs), Self::F64(rhs)) => {
let data = binary_map(shape, lhs_layout, rhs_layout, lhs, rhs, B::f64);
let data = binary_map(lhs_layout, rhs_layout, lhs, rhs, B::f64);
Ok(Self::F64(data))
}
(Self::U32(lhs), Self::U32(rhs)) => {
let data = binary_map(shape, lhs_layout, rhs_layout, lhs, rhs, B::u32);
let data = binary_map(lhs_layout, rhs_layout, lhs, rhs, B::u32);
Ok(Self::U32(data))
}
_ => {
@ -622,24 +622,24 @@ impl CpuStorage {
map1!(vs, take_impl1, ids, layout, vocab_size, hidden_size)
}
pub(crate) fn matmul_impl(
pub(crate) fn matmul(
&self,
rhs: &Self,
bmnk: (usize, usize, usize, usize),
lhs_stride: &[usize],
rhs_stride: &[usize],
lhs_layout: &Layout,
rhs_layout: &Layout,
) -> Result<Self> {
match (self, rhs) {
(CpuStorage::F16(lhs), CpuStorage::F16(rhs)) => {
let dst = matmul_impl(lhs, rhs, bmnk, lhs_stride, rhs_stride)?;
let dst = matmul(lhs, rhs, bmnk, lhs_layout, rhs_layout)?;
Ok(Self::F16(dst))
}
(CpuStorage::F32(lhs), CpuStorage::F32(rhs)) => {
let dst = matmul_impl(lhs, rhs, bmnk, lhs_stride, rhs_stride)?;
let dst = matmul(lhs, rhs, bmnk, lhs_layout, rhs_layout)?;
Ok(Self::F32(dst))
}
(CpuStorage::F64(lhs), CpuStorage::F64(rhs)) => {
let dst = matmul_impl(lhs, rhs, bmnk, lhs_stride, rhs_stride)?;
let dst = matmul(lhs, rhs, bmnk, lhs_layout, rhs_layout)?;
Ok(Self::F64(dst))
}
_ => Err(Error::DTypeMismatchBinaryOp {

View File

@ -1,5 +1,5 @@
#![allow(dead_code)]
use crate::{CpuStorage, DType, Error, Result, Shape};
use crate::{CpuStorage, DType, Error, Layout, Result, Shape};
#[derive(thiserror::Error, Debug)]
pub enum DummyError {}
@ -60,11 +60,11 @@ impl CudaStorage {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn affine_impl(&self, _: &Shape, _: &[usize], _: f64, _: f64) -> Result<Self> {
pub(crate) fn affine(&self, _: &Layout, _: f64, _: f64) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn sum(&self, _: &Shape, _: &[usize], _: &[usize]) -> Result<Self> {
pub(crate) fn sum(&self, _: &Layout, _: &[usize]) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
@ -72,65 +72,49 @@ impl CudaStorage {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn to_dtype(&self, _: &Shape, _: &[usize], _: DType) -> Result<Self> {
pub(crate) fn to_dtype(&self, _: &Layout, _: DType) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn unary_impl<B: crate::op::UnaryOp>(&self, _: &Shape, _: &[usize]) -> Result<Self> {
pub(crate) fn unary_impl<B: crate::op::UnaryOp>(&self, _: &Layout) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn binary_impl<B: crate::op::BinaryOp>(
&self,
_: &Self,
_: &Shape,
_: &[usize],
_: &[usize],
_: &Layout,
_: &Layout,
) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn where_cond(
&self,
_: &Shape,
_: &[usize],
_: &Layout,
_: &Self,
_: &[usize],
_: &Layout,
_: &Self,
_: &[usize],
_: &Layout,
) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn embedding_impl(
&self,
_: &Shape,
_: &[usize],
_: &Self,
_: usize,
_: usize,
) -> Result<Self> {
pub(crate) fn embedding(&self, _: &Layout, _: &Self, _: usize, _: usize) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn matmul_impl(
pub(crate) fn matmul(
&self,
_: &Self,
_: (usize, usize, usize, usize),
_: &[usize],
_: &[usize],
_: &Layout,
_: &Layout,
) -> Result<Self> {
Err(Error::NotCompiledWithCudaSupport)
}
pub(crate) fn copy_strided_src(
&self,
_: &mut Self,
_: usize,
_: &Shape,
_: &[usize],
_: usize,
) -> Result<()> {
pub(crate) fn copy_strided_src(&self, _: &mut Self, _: usize, _: &Layout) -> Result<()> {
Err(Error::NotCompiledWithCudaSupport)
}
}

View File

@ -106,7 +106,7 @@ impl Layout {
if shape.rank() < self.shape().rank() {
Err(Error::BroadcastIncompatibleShapes {
src_shape: self.shape().clone(),
dst_shape: shape,
dst_shape: shape.clone(),
})?
}
let added_dims = shape.rank() - self.shape().rank();
@ -135,6 +135,6 @@ impl Layout {
}
pub(crate) fn strided_index(&self) -> crate::StridedIndex {
crate::StridedIndex::new(&self)
crate::StridedIndex::new(self)
}
}

View File

@ -79,6 +79,7 @@ impl Storage {
}
}
// This assumes a contiguous layout and no offset.
pub(crate) fn divide_by_sum_over_dim(&mut self, shape: &Shape, dim: usize) -> Result<()> {
match self {
Storage::Cpu(storage) => storage.divide_by_sum_over_dim(shape, dim)?,
@ -196,22 +197,22 @@ impl Storage {
}
}
pub(crate) fn matmul_impl(
pub(crate) fn matmul(
&self,
rhs: &Self,
bmnk: (usize, usize, usize, usize),
lhs_stride: &[usize],
rhs_stride: &[usize],
lhs_layout: &Layout,
rhs_layout: &Layout,
) -> Result<Self> {
self.same_device(rhs, "matmul")?;
self.same_dtype(rhs, "matmul")?;
match (self, rhs) {
(Self::Cpu(lhs), Self::Cpu(rhs)) => {
let storage = lhs.matmul_impl(rhs, bmnk, lhs_stride, rhs_stride)?;
let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
Ok(Self::Cpu(storage))
}
(Self::Cuda(lhs), Self::Cuda(rhs)) => {
let storage = lhs.matmul_impl(rhs, bmnk, lhs_stride, rhs_stride)?;
let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
Ok(Self::Cuda(storage))
}
(lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {

View File

@ -432,7 +432,7 @@ impl Tensor {
let c_shape = Shape::from(&a_dims[..dim - 2]).extend(&[m, n]);
let batching: usize = a_dims[..dim - 2].iter().product();
let storage = self.storage.matmul_impl(
let storage = self.storage.matmul(
&rhs.storage,
(batching, m, n, k),
self.layout(),
@ -587,7 +587,7 @@ impl Tensor {
}
pub fn shape(&self) -> &Shape {
&self.layout().shape()
self.layout().shape()
}
pub fn dims(&self) -> &[usize] {
@ -600,7 +600,7 @@ impl Tensor {
// TODO: Rename to `stride` once the PR that introduced the layout has been merged.
pub fn stride_tmp(&self) -> &[usize] {
&self.layout.stride()
self.layout.stride()
}
pub fn rank(&self) -> usize {