Better slicing

moose/src/boolean/mod.rs

@@ -4,7 +4,7 @@ use crate::computation::*;
 use crate::error::{Error, Result};
 use crate::execution::Session;
 use crate::floatingpoint::FloatTensor;
-use crate::host::HostPlacement;
+use crate::host::{HostPlacement, SliceInfo};
 use crate::integer::AbstractUint64Tensor;
 use crate::kernels::*;
 use crate::replicated::ReplicatedPlacement;
@@ -281,6 +281,44 @@ impl IndexAxisOp {
     }
 }
 
+impl SliceOp {
+    pub(crate) fn bool_rep_kernel<S: Session, HostT, RepT>(
+        sess: &S,
+        plc: &ReplicatedPlacement,
+        info: SliceInfo,
+        x: BoolTensor<HostT, RepT>,
+    ) -> Result<BoolTensor<HostT, RepT>>
+    where
+        ReplicatedPlacement: PlacementSlice<S, RepT, RepT>,
+        ReplicatedPlacement: PlacementShare<S, HostT, RepT>,
+    {
+        let x = match x {
+            BoolTensor::Host(v) => plc.share(sess, &v),
+            BoolTensor::Replicated(v) => v,
+        };
+        let z = plc.slice(sess, info, &x);
+        Ok(BoolTensor::Replicated(z))
+    }
+
+    pub(crate) fn bool_host_kernel<S: Session, HostT, RepT>(
+        sess: &S,
+        plc: &HostPlacement,
+        info: SliceInfo,
+        x: BoolTensor<HostT, RepT>,
+    ) -> Result<BoolTensor<HostT, RepT>>
+    where
+        HostPlacement: PlacementSlice<S, HostT, HostT>,
+        HostPlacement: PlacementReveal<S, RepT, HostT>,
+    {
+        let x = match x {
+            BoolTensor::Replicated(v) => plc.reveal(sess, &v),
+            BoolTensor::Host(v) => v,
+        };
+        let z = plc.slice(sess, info, &x);
+        Ok(BoolTensor::Host(z))
+    }
+}
+
 impl SaveOp {
     pub(crate) fn bool_kernel<S: Session, HostT, RepT>(
         sess: &S,

moose/src/fixedpoint/ops.rs

@@ -1108,6 +1108,65 @@ impl IndexAxisOp {
     }
 }
 
+impl SliceOp {
+    pub(crate) fn fixed_host_kernel<S: Session, HostFixedT, MirFixedT, RepFixedT>(
+        sess: &S,
+        plc: &HostPlacement,
+        info: SliceInfo,
+        x: FixedTensor<HostFixedT, MirFixedT, RepFixedT>,
+    ) -> Result<FixedTensor<HostFixedT, MirFixedT, RepFixedT>>
+    where
+        HostPlacement: PlacementReveal<S, RepFixedT, HostFixedT>,
+        HostPlacement: PlacementDemirror<S, MirFixedT, HostFixedT>,
+        HostPlacement: PlacementSlice<S, HostFixedT, HostFixedT>,
+    {
+        let x = match x {
+            FixedTensor::Replicated(v) => plc.reveal(sess, &v),
+            FixedTensor::Mirrored3(v) => plc.demirror(sess, &v),
+            FixedTensor::Host(v) => v,
+        };
+        let z = plc.slice(sess, info, &x);
+        Ok(FixedTensor::Host(z))
+    }
+
+    pub(crate) fn fixed_rep_kernel<S: Session, HostFixedT, MirFixedT, RepFixedT>(
+        sess: &S,
+        plc: &ReplicatedPlacement,
+        info: SliceInfo,
+        x: FixedTensor<HostFixedT, MirFixedT, RepFixedT>,
+    ) -> Result<FixedTensor<HostFixedT, MirFixedT, RepFixedT>>
+    where
+        ReplicatedPlacement: PlacementShare<S, HostFixedT, RepFixedT>,
+        ReplicatedPlacement: PlacementShare<S, MirFixedT, RepFixedT>,
+        ReplicatedPlacement: PlacementSlice<S, RepFixedT, RepFixedT>,
+    {
+        let x = match x {
+            FixedTensor::Host(v) => plc.share(sess, &v),
+            FixedTensor::Mirrored3(v) => plc.share(sess, &v),
+            FixedTensor::Replicated(v) => v,
+        };
+        let z = plc.slice(sess, info, &x);
+        Ok(FixedTensor::Replicated(z))
+    }
+
+    pub(crate) fn repfixed_kernel<S: Session, RepRingT>(
+        sess: &S,
+        plc: &ReplicatedPlacement,
+        info: SliceInfo,
+        x: RepFixedTensor<RepRingT>,
+    ) -> Result<RepFixedTensor<RepRingT>>
+    where
+        ReplicatedPlacement: PlacementSlice<S, RepRingT, RepRingT>,
+    {
+        let y = plc.slice(sess, info, &x.tensor);
+        Ok(RepFixedTensor {
+            tensor: y,
+            fractional_precision: x.fractional_precision,
+            integral_precision: x.integral_precision,
+        })
+    }
+}
+
 impl ShapeOp {
     pub(crate) fn host_fixed_kernel<S: Session, HostFixedT, MirFixedT, RepFixedT, HostShapeT>(
         sess: &S,

moose/src/floatingpoint/ops.rs

@@ -4,7 +4,7 @@ use crate::computation::*;
 use crate::error::Error;
 use crate::error::Result;
 use crate::execution::Session;
-use crate::host::HostPlacement;
+use crate::host::{HostPlacement, SliceInfo};
 use crate::kernels::*;
 use crate::mirrored::{Mir3Tensor, Mirrored3Placement};
 use crate::types::*;
@@ -631,3 +631,26 @@ impl MuxOp {
         Ok(FloatTensor::Host(z))
     }
 }
+
+impl SliceOp {
+    pub(crate) fn float_host_kernel<S: Session, HostFloatT, MirroredT>(
+        sess: &S,
+        plc: &HostPlacement,
+        slice: SliceInfo,
+        x: FloatTensor<HostFloatT, MirroredT>,
+    ) -> Result<FloatTensor<HostFloatT, MirroredT>>
+    where
+        HostPlacement: PlacementSlice<S, HostFloatT, HostFloatT>,
+    {
+        let x = match x {
+            FloatTensor::Host(v) => v,
+            FloatTensor::Mirrored3(_v) => {
+                return Err(Error::UnimplementedOperator(
+                    "SliceOp @ Mirrored3Placement".to_string(),
+                ))
+            }
+        };
+        let z = plc.slice(sess, slice, &x);
+        Ok(FloatTensor::Host(z))
+    }
+}

moose/src/host/bitarray.rs

@@ -1,4 +1,5 @@
 use crate::host::RawShape;
+use crate::host::SliceInfo;
 use anyhow::anyhow;
 use bitvec::prelude::*;
 use ndarray::{prelude::*, RemoveAxis};
@@ -125,6 +126,10 @@ impl BitArrayRepr {
             dim: Arc::new(IxDyn(&[len])),
         }
     }
+
+    pub(crate) fn slice(&self, _info: SliceInfo) -> anyhow::Result<BitArrayRepr> {
+        unimplemented!()
+    }
 }
 
 impl std::ops::BitXor for &BitArrayRepr {

moose/src/host/ops.rs

@@ -475,31 +475,71 @@ impl AtLeast2DOp {
 }
 
 impl SliceOp {
-    pub(crate) fn host_kernel<S: RuntimeSession, T>(
-        _sess: &S,
+    pub(crate) fn host_fixed_kernel<S: Session, HostRingT>(
+        sess: &S,
+        plc: &HostPlacement,
+        info: SliceInfo,
+        x: HostFixedTensor<HostRingT>,
+    ) -> Result<HostFixedTensor<HostRingT>>
+    where
+        HostPlacement: PlacementSlice<S, HostRingT, HostRingT>,
+    {
+        let tensor = plc.slice(sess, info, &x.tensor);
+        Ok(HostFixedTensor::<HostRingT> {
+            tensor,
+            fractional_precision: x.fractional_precision,
+            integral_precision: x.integral_precision,
+        })
+    }
+
+    pub(crate) fn host_bit_kernel<S: RuntimeSession>(
+        sess: &S,
         plc: &HostPlacement,
-        slice_info: SliceInfo,
+        info: SliceInfo,
+        x: HostBitTensor,
+    ) -> Result<HostBitTensor>
+    where
+        HostPlacement: PlacementPlace<S, HostBitTensor>,
+    {
+        let x = plc.place(sess, x);
+        x.slice(info)
+    }
+
+    pub(crate) fn host_float_kernel<S: RuntimeSession, T: Clone>(
+        sess: &S,
+        plc: &HostPlacement,
+        info: SliceInfo,
+        x: HostTensor<T>,
+    ) -> Result<HostTensor<T>>
+    where
+        HostPlacement: PlacementPlace<S, HostTensor<T>>,
+    {
+        let x = plc.place(sess, x);
+        x.slice(info)
+    }
+
+    pub(crate) fn host_ring_kernel<S: RuntimeSession, T>(
+        sess: &S,
+        plc: &HostPlacement,
+        info: SliceInfo,
         x: HostRingTensor<T>,
     ) -> Result<HostRingTensor<T>>
     where
         T: Clone,
+        HostPlacement: PlacementPlace<S, HostRingTensor<T>>,
     {
-        let slice_info =
-            ndarray::SliceInfo::<Vec<ndarray::SliceInfoElem>, IxDyn, IxDyn>::from(slice_info);
-        let sliced = x.0.slice(slice_info).to_owned();
-        Ok(HostRingTensor(sliced.to_shared(), plc.clone()))
+        let x = plc.place(sess, x);
+        x.slice(info)
     }
 
     pub(crate) fn shape_kernel<S: RuntimeSession>(
         _sess: &S,
         plc: &HostPlacement,
-        slice_info: SliceInfo,
+        info: SliceInfo,
         x: HostShape,
     ) -> Result<HostShape> {
-        let slice = x.0.slice(
-            slice_info.0[0].start as usize,
-            slice_info.0[0].end.unwrap() as usize,
-        );
+        let slice =
+            x.0.slice(info.0[0].start as usize, info.0[0].end.unwrap() as usize);
         Ok(HostShape(slice, plc.clone()))
     }
 }
@@ -561,6 +601,21 @@ impl<T: LinalgScalar> HostTensor<T> {
     }
 }
 
+impl<T: Clone> HostTensor<T> {
+    fn slice(&self, info: SliceInfo) -> Result<HostTensor<T>> {
+        if info.0.len() != self.0.ndim() {
+            return Err(Error::InvalidArgument(format!(
+                "The input dimension of `info` must match the array to be sliced. Used slice info dim {}, tensor had dim {}",
+                info.0.len(),
+                self.0.ndim()
+            )));
+        }
+        let info = ndarray::SliceInfo::<Vec<ndarray::SliceInfoElem>, IxDyn, IxDyn>::from(info);
+        let result = self.0.slice(info);
+        Ok(HostTensor(result.to_owned().into_shared(), self.1.clone()))
+    }
+}
+
 impl<T: Clone> HostRingTensor<T> {
     fn index_axis(self, axis: usize, index: usize) -> Result<HostRingTensor<T>> {
         if axis >= self.0.ndim() {
@@ -583,6 +638,24 @@ impl<T: Clone> HostRingTensor<T> {
     }
 }
 
+impl<T: Clone> HostRingTensor<T> {
+    fn slice(&self, info: SliceInfo) -> Result<HostRingTensor<T>> {
+        if info.0.len() != self.0.ndim() {
+            return Err(Error::InvalidArgument(format!(
+                "The input dimension of `info` must match the array to be sliced. Used slice info dim {}, tensor had dim {}",
+                info.0.len(),
+                self.0.ndim()
+            )));
+        }
+        let info = ndarray::SliceInfo::<Vec<ndarray::SliceInfoElem>, IxDyn, IxDyn>::from(info);
+        let result = self.0.slice(info);
+        Ok(HostRingTensor(
+            result.to_owned().into_shared(),
+            self.1.clone(),
+        ))
+    }
+}
+
 impl HostBitTensor {
     fn index_axis(self, axis: usize, index: usize) -> Result<HostBitTensor> {
         if axis >= self.0.ndim() {
@@ -602,6 +675,21 @@ impl HostBitTensor {
         let result = self.0.index_axis(axis, index);
         Ok(HostBitTensor(result, self.1))
     }
+
+    fn slice(&self, info: SliceInfo) -> Result<HostBitTensor> {
+        if info.0.len() != self.0.ndim() {
+            return Err(Error::InvalidArgument(format!(
+                "The input dimension of `info` must match the array to be sliced. Used slice info dim {}, tensor had dim {}",
+                info.0.len(),
+                self.0.ndim()
+            )));
+        }
+        let result = self
+            .0
+            .slice(info)
+            .map_err(|e| Error::KernelError(e.to_string()))?;
+        Ok(HostBitTensor(result, self.1.clone()))
+    }
 }
 
 impl IndexAxisOp {

moose/src/kernels/indexing.rs

@@ -9,7 +9,6 @@ modelled_kernel! {
     PlacementIndexAxis::index_axis, IndexAxisOp{axis: usize, index: usize},
     [
         (HostPlacement, (BooleanTensor) -> BooleanTensor => [concrete] Self::bool_host_kernel),
-        (HostPlacement, (Tensor) -> Tensor => [concrete] Self::logical_host_kernel),
         (HostPlacement, (Float32Tensor) -> Float32Tensor => [concrete] Self::float_host_kernel),
         (HostPlacement, (Float64Tensor) -> Float64Tensor => [concrete] Self::float_host_kernel),
         (HostPlacement, (Fixed64Tensor) -> Fixed64Tensor => [concrete] Self::fixed_host_kernel),
@@ -21,6 +20,7 @@ modelled_kernel! {
         (HostPlacement, (HostFloat64Tensor) -> HostFloat64Tensor => [runtime] Self::host_float_kernel),
         (HostPlacement, (HostRing64Tensor) -> HostRing64Tensor => [runtime] Self::host_ring_kernel),
         (HostPlacement, (HostRing128Tensor) -> HostRing128Tensor => [runtime] Self::host_ring_kernel),
+        (HostPlacement, (Tensor) -> Tensor => [concrete] Self::logical_host_kernel),
         (ReplicatedPlacement, (BooleanTensor) -> BooleanTensor => [concrete]  Self::bool_rep_kernel),
         (ReplicatedPlacement, (Tensor) -> Tensor => [concrete] Self::logical_rep_kernel),
         (ReplicatedPlacement, (Fixed64Tensor) -> Fixed64Tensor => [concrete] Self::fixed_rep_kernel),
@@ -57,11 +57,31 @@ pub trait PlacementSlice<S: Session, T, O> {
 modelled_kernel! {
     PlacementSlice::slice, SliceOp{slice: SliceInfo},
     [
-        (HostPlacement, (Shape) -> Shape => [concrete]  Self::logical_host_shape),
         (HostPlacement, (HostShape) -> HostShape => [runtime] Self::shape_kernel),
-        (HostPlacement, (HostRing64Tensor) -> HostRing64Tensor => [runtime] Self::host_kernel),
-        (HostPlacement, (HostRing128Tensor) -> HostRing128Tensor => [runtime] Self::host_kernel),
+        (HostPlacement, (Shape) -> Shape => [concrete]  Self::logical_host_shape),
+        (HostPlacement, (BooleanTensor) -> BooleanTensor => [concrete] Self::bool_host_kernel),
+        (HostPlacement, (Fixed64Tensor) -> Fixed64Tensor => [concrete] Self::fixed_host_kernel),
+        (HostPlacement, (Fixed128Tensor) -> Fixed128Tensor => [concrete] Self::fixed_host_kernel),
+        (HostPlacement, (Float32Tensor) -> Float32Tensor => [concrete] Self::float_host_kernel),
+        (HostPlacement, (Float64Tensor) -> Float64Tensor => [concrete] Self::float_host_kernel),
+        (HostPlacement, (HostBitTensor) -> HostBitTensor => [runtime] Self::host_bit_kernel),
+        (HostPlacement, (HostFixed64Tensor) -> HostFixed64Tensor => [concrete] Self::host_fixed_kernel),
+        (HostPlacement, (HostFixed128Tensor) -> HostFixed128Tensor => [concrete] Self::host_fixed_kernel),
+        (HostPlacement, (HostFloat32Tensor) -> HostFloat32Tensor => [runtime] Self::host_float_kernel),
+        (HostPlacement, (HostFloat64Tensor) -> HostFloat64Tensor => [runtime] Self::host_float_kernel),
+        (HostPlacement, (HostRing64Tensor) -> HostRing64Tensor => [runtime] Self::host_ring_kernel),
+        (HostPlacement, (HostRing128Tensor) -> HostRing128Tensor => [runtime] Self::host_ring_kernel),
+        (HostPlacement, (Tensor) -> Tensor => [concrete] Self::logical_host_kernel),
         (ReplicatedPlacement, (Shape) -> Shape => [concrete]  Self::logical_rep_shape),
-        (ReplicatedPlacement, (ReplicatedShape) -> ReplicatedShape => [concrete] Self::rep_kernel),
+        (ReplicatedPlacement, (BooleanTensor) -> BooleanTensor => [concrete] Self::bool_rep_kernel),
+        (ReplicatedPlacement, (ReplicatedShape) -> ReplicatedShape => [concrete] Self::rep_shape_kernel),
+        (ReplicatedPlacement, (ReplicatedBitTensor) -> ReplicatedBitTensor => [concrete] Self::rep_ring_kernel),
+        (ReplicatedPlacement, (ReplicatedRing64Tensor) -> ReplicatedRing64Tensor => [concrete] Self::rep_ring_kernel),
+        (ReplicatedPlacement, (ReplicatedRing128Tensor) -> ReplicatedRing128Tensor => [concrete] Self::rep_ring_kernel),
+        (ReplicatedPlacement, (ReplicatedFixed64Tensor) -> ReplicatedFixed64Tensor => [concrete] Self::repfixed_kernel),
+        (ReplicatedPlacement, (ReplicatedFixed128Tensor) -> ReplicatedFixed128Tensor => [concrete] Self::repfixed_kernel),
+        (ReplicatedPlacement, (Fixed64Tensor) -> Fixed64Tensor => [concrete] Self::fixed_rep_kernel),
+        (ReplicatedPlacement, (Fixed128Tensor) -> Fixed128Tensor => [concrete] Self::fixed_rep_kernel),
+        (ReplicatedPlacement, (Tensor) -> Tensor => [concrete] Self::logical_rep_kernel),
     ]
 }