WIP

Cargo.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+itertools = "0.12.0"
 num = "0.4.1"
 ord_subset = "3.1.1"
 slotmap = "1.0.6"

ideas.md

@@ -21,6 +21,48 @@
 - Tensors all have an implicit ordering of position indexes which follows the significance from left being most significant and right being least significant
 - Iterators between positions, reshaping all operate on this idea
 
+## Algorithms
+
+### Conv2d
+```
+a1: (batches, in_channels, y, x)
+k1: (out_channels, in_channels, k_y, k_x)
+
+---
+
+let a2: (batches, out_channels, in_channels, y, x) = a1.broadcast_splice([out_channels], 1)
+let k2: (batches, out_channels, in_channels, k_y, k_x) = k1.broadcast_splice([batches], 0)
+
+let k_half_len_y = (k_y // 2)
+let k_half_len_x = (k_x // 2)
+
+let start_y = k_half_len_y
+let end_y = y - start_y
+let start_x = k_half_len_x
+let end_x = x - start_x
+
+let mut out_units = 
+
+for curr_y in start_y..end_y {
+    for curr_x in start_x..end_x {
+        let a3: (batches, out_channels, in_channels, k_y, k_x) = a2.slice([:, :, :, curr_y-k_half_len_y:curr_y+k_half_len_y, curr_x-k_half_len_x:curr_x+k_half_len_x])
+
+        let r1: (batches, out_channels, in_channels, k_y, k_x) = a3 * k2
+
+        let r2: (batches, out_channels, in_channels * k_y * k_x) = r1.reshape([batches, out_channels, in_channels, k_y * k_x])
+
+        let r3: (batches, out_channels, 1) = r2.sum()
+
+        // batches, x, y
+        out_units.extend(r3.iter_unit())
+    }
+}
+
+return tensor::from_slice(&out_units): (y, x, batches, out_channels)
+
+
+```
+
 ## TODO
 
 - Refactor comp_graph to improve errors (ones with no nodekey) and reduce repeated code

operations.md

@@ -6,7 +6,7 @@ Check indictates support
     * [x] abs
     * [x] neg
 
-* Pointwise Scalar
+* Pointwise Scalar (broadcast?)
     * [x] add_scalar
     * [x] sub_scalar_lh (s - t)
     * [x] sub_scalar_rh (t - s)
@@ -31,6 +31,6 @@ Check indictates support
     * [x] shape
     * [x] stride
     * [x] iter
-    * [ ] slice
-    * [ ] reshape
+    * [x] slice
+    * [x] reshape
     * [ ] concat

src/comp_graph/mod.rs

@@ -5,7 +5,7 @@ use std::collections::{HashSet, HashMap};
 use slotmap::{SlotMap, new_key_type};
 use thiserror::Error;
 
-use crate::engine::{tensor::{EngineTensor, allowed_unit::AllowedUnit, factory::EngineTensorFactory, iter::EngineTensorIterator}, Engine, EngineError};
+use crate::engine::{tensor::{EngineTensor, allowed_unit::AllowedUnit, factory::EngineTensorFactory, iter::EngineTensorUnitIterator}, Engine, EngineError};
 
 use self::edge::Edge;
 
@@ -101,8 +101,8 @@ impl<T: AllowedUnit> CompGraph<T> {
         self.nodes.insert(Node::create_node(edge))
     }
 
-    pub fn iter(&self, tensor: &CompGraphTensor) -> EngineTensorIterator<T> {
-        EngineTensorIterator::new(self.get_node(tensor.node_key()).unwrap().tensor().unwrap())
+    pub fn iter(&self, tensor: &CompGraphTensor) -> EngineTensorUnitIterator<T> {
+        EngineTensorUnitIterator::new(self.get_node(tensor.node_key()).unwrap().tensor().unwrap())
     }
 
     //First return is open nodes, second is node_to_children
@@ -470,7 +470,7 @@ mod test {
 
         let tensor = new_node_keys.last().unwrap();
 
-        let expected = Array::from_iter( &mut expected_original.iter().map(|x| x * 2.0f32.pow(power)), expected_original.shape().clone());
+        let expected = Array::from_iter( &mut expected_original.iter_unit().map(|x| x * 2.0f32.pow(power)), expected_original.shape().clone());
 
         graph.non_populating_eval(&tensor).unwrap();
 

src/engine/basic.rs

@@ -1 +1,248 @@
-pub struct Basic {}
+use std::ops::IndexMut;
+
+use itertools::Itertools;
+
+use crate::{engine::{Engine, EngineError, EngineTensorFactory, EngineTensor, util::{err_if_incorrect_dimensions, err_if_dimension_mismatch}}, helper::{Position, Interval, Slice, Shape, Stride}};
+
+pub struct Basic {}
+
+macro_rules! conv_fn {
+    ($unit:ty) => {
+        fn conv2d<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, kernel: &dyn EngineTensor<Unit = $unit>, stride: usize) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+            //a: (batches, in_channels, y, x)
+            //kernel: (out_channels, in_channels, k_y, k_x)
+
+            err_if_incorrect_dimensions(a.shape(), 4)?;
+            err_if_incorrect_dimensions(kernel.shape(), 4)?;
+            err_if_dimension_mismatch(a.shape().dim(1), kernel.shape().dim(1))?;
+
+            let y = a.shape().dim(a.shape().dims() - 2);
+            let x = a.shape().dim(a.shape().dims() - 1);
+            let k_y = kernel.shape().dim(kernel.shape().dims() - 2) + 2 * (stride - 1);
+            let k_x = kernel.shape().dim(kernel.shape().dims() - 1) + 2 * (stride - 1);
+
+            let batches = a.shape().dim(0);
+            let out_channels = kernel.shape().dim(0);
+            let in_channels = kernel.shape().dim(1);
+
+            //(batches, out_channels, in_channels, y, x)
+            let a_broadcast = a.broadcast_splice(1, &[out_channels]);
+            //(batches, out_channels, in_channels, k_y, k_x)
+            let kernel_broadcast = kernel.broadcast_splice(0, &[batches]);
+
+            let half_k_y = k_y / 2;
+            let half_k_x = k_x / 2;
+
+            let y_out = y - half_k_y * 2;
+            let x_out = x - half_k_x * 2;
+            let out_shape = Shape::from([batches, out_channels, y_out, x_out].as_slice());
+            let out_stride = Stride::from(&out_shape);
+
+            let mut reordered_sums: Box<[$unit]> = vec![<$unit>::default(); batches * out_channels * y_out * x_out].into_boxed_slice();
+
+            for curr_y in 0..y_out {
+                for curr_x in 0..x_out {
+                    let a_sliced = a_broadcast.slice(&Slice::from([Interval::all(), Interval::all(), Interval::all(), Interval::between_with_step(curr_y, curr_y + k_y, stride), Interval::between_with_step(curr_x, curr_x + k_x, stride)].as_slice()));
+
+                    let chunked_products = a_sliced.iter_unit().zip(kernel_broadcast.iter_unit()).map(|(a_elem, k_elem)| a_elem * k_elem).chunks(in_channels * k_y * k_x);
+                    let curr_sums = chunked_products.into_iter().map(|i| -> $unit { i.sum() });
+
+                    for (i, sum) in curr_sums.enumerate() {
+                        let batch = i / out_channels;
+                        let out_channel = i % out_channels;
+
+                        let index = Position::from([batch, out_channel, curr_y, curr_x].as_slice()).tensor_index(&out_stride)?;
+
+                        *reordered_sums.index_mut(index) = sum;
+                    }
+                }
+            }
+
+            //(batches, out_channels, y_out, x_out)
+            Ok(E::from_slice(&reordered_sums, out_shape))
+        }
+    };
+}
+
+macro_rules! basic_impl {
+    ($unit:ty) => {
+        impl Engine<$unit> for Basic {
+
+            //Pointwise Single
+            fn abs<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+
+                Ok(E::from_iter(a.iter_unit().map(|x| x.abs()), a.shape().clone()))
+            }
+
+            fn neg<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| -x), a.shape().clone()))
+            }
+
+            //Scalar
+            fn add_scalar<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x + s), a.shape().clone()))
+            }
+
+            fn sub_scalar_lh<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| s - x), a.shape().clone()))
+            }
+
+            fn sub_scalar_rh<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, s: $unit) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x - s), a.shape().clone()))
+            }
+
+            fn mul_scalar<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x * s), a.shape().clone()))
+            }
+
+            fn div_scalar_lh<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| s / x), a.shape().clone()))
+            }
+
+            fn div_scalar_rh<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, s: $unit) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x / s), a.shape().clone()))
+            }
+
+            //Pointwise Double
+            fn add<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x + y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            fn sub<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x - y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            fn mul<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x * y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            fn div<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x / y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            conv_fn!($unit);
+        }
+    };
+}
+
+macro_rules! basic_unsigned_impl {
+    ($unit:ty) => {
+        impl Engine<$unit> for Basic {
+
+            //Pointwise Single
+            fn abs<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+
+                Ok(E::from_iter(&mut a.iter_unit().map(|x| x), a.shape().clone()))
+            }
+
+            fn neg<E: EngineTensorFactory<Unit = $unit>>(_a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                Err(crate::engine::EngineError::OperationUnsupportedForType())
+            }
+
+            //Scalar
+            fn add_scalar<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x + s), a.shape().clone()))
+            }
+
+            fn sub_scalar_lh<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| s - x), a.shape().clone()))
+            }
+
+            fn sub_scalar_rh<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, s: $unit) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x - s), a.shape().clone()))
+            }
+
+            fn mul_scalar<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x * s), a.shape().clone()))
+            }
+
+            fn div_scalar_lh<E: EngineTensorFactory<Unit = $unit>>(s: $unit, a: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| s / x), a.shape().clone()))
+            }
+
+            fn div_scalar_rh<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, s: $unit) -> Result<Box<dyn EngineTensor<Unit = $unit>>, EngineError> {
+                Ok(E::from_iter(a.iter_unit().map(|x| x / s), a.shape().clone()))
+            }
+
+            //Pointwise Double
+            fn add<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x + y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            fn sub<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x - y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            fn mul<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x * y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            fn div<E: EngineTensorFactory<Unit = $unit>>(a: &dyn EngineTensor<Unit = $unit>, b: &dyn EngineTensor<Unit = $unit>) -> Result<Box<dyn EngineTensor<Unit = $unit>>, crate::engine::EngineError> {
+                if a.shape() == b.shape() {
+                    Ok(E::from_iter(&mut a.iter_unit().zip(b.iter_unit()).map(|(x, y)| x / y), a.shape().clone()))
+                } else {
+                    Err(EngineError::ShapeMismatch(a.shape().clone(), b.shape().clone()))
+                }
+            }
+
+            conv_fn!($unit);
+        }
+    };
+}
+
+basic_impl!(f32);
+basic_impl!(f64);
+basic_impl!(i8);
+basic_impl!(i16);
+basic_impl!(i32);
+basic_impl!(i64);
+basic_unsigned_impl!(u8);
+basic_unsigned_impl!(u16);
+basic_unsigned_impl!(u32);
+basic_unsigned_impl!(u64);
+
+#[cfg(test)]
+mod test {
+    use crate::{helper::shape, engine::tensor::Array};
+
+    use super::*;
+
+    #[test]
+    pub fn conv() {
+        let a = Array::from_iter((1..=65536).map(|x| (x as f32) / 65536.0).cycle().take(1 * 3 * 256 * 256), shape![1, 3, 256, 256]);
+        let kernel = Array::from_iter((1..=9).map(|x| x as f32).cycle().take(1 * 3 * 3 * 3), shape![1, 3, 3, 3]);
+
+        let res =  Basic::conv2d::<Array<f32>>(a.as_ref(), kernel.as_ref(), 2).unwrap();
+
+        println!("{:?}", res.shape());
+        //println!("{:?}", res.iter_unit().collect::<Vec<f32>>());
+    }
+}