From e7468ef40825fae42f1e75e66bb1854cf618deb2 Mon Sep 17 00:00:00 2001 From: Morten Dahl Date: Tue, 8 Feb 2022 14:40:39 +0100 Subject: [PATCH 1/2] add_n --- moose/src/kernels/arithmetic.rs | 23 ++--------------------- moose/src/lib.rs | 7 +++---- 2 files changed, 5 insertions(+), 25 deletions(-) diff --git a/moose/src/kernels/arithmetic.rs b/moose/src/kernels/arithmetic.rs index 98c061e47..afb13d182 100644 --- a/moose/src/kernels/arithmetic.rs +++ b/moose/src/kernels/arithmetic.rs @@ -275,27 +275,8 @@ pub trait PlacementAddN { fn add_n(&self, sess: &S, x: &[T]) -> O; } -modelled!(PlacementAddN::add_n, HostPlacement, vec[Tensor] -> Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[Float32Tensor] -> Float32Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[Float64Tensor] -> Float64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[HostFloat32Tensor] -> HostFloat32Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[HostFloat64Tensor] -> HostFloat64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[Fixed64Tensor] -> Fixed64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[Fixed128Tensor] -> Fixed128Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[HostFixed64Tensor] -> HostFixed64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[HostFixed128Tensor] -> HostFixed128Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[HostRing64Tensor] -> HostRing64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, HostPlacement, vec[HostRing128Tensor] -> HostRing128Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[ReplicatedRing64Tensor] -> ReplicatedRing64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[ReplicatedRing128Tensor] -> ReplicatedRing128Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[ReplicatedFixed64Tensor] -> ReplicatedFixed64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[ReplicatedFixed128Tensor] -> ReplicatedFixed128Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[Fixed64Tensor] -> Fixed64Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[Fixed128Tensor] -> Fixed128Tensor, AddNOp); -modelled!(PlacementAddN::add_n, ReplicatedPlacement, vec[Tensor] -> Tensor, AddNOp); - -kernel! { - AddNOp, +modelled_kernel! { + PlacementAddN::add_n, AddNOp, [ (HostPlacement, vec[Tensor] -> Tensor => [concrete] Self::host_logical_kernel), (HostPlacement, vec[Float32Tensor] -> Float32Tensor => [concrete] Self::float_kernel), diff --git a/moose/src/lib.rs b/moose/src/lib.rs index 3f6c7aff6..124a150ea 100644 --- a/moose/src/lib.rs +++ b/moose/src/lib.rs @@ -150,7 +150,6 @@ macro_rules! derive_runtime_kernel { })) } }; - (variadic, attributes[$($attr:ident$(: $prim_ty:ident)?),+] $k:expr, $self:ident) => { { $( @@ -4548,7 +4547,7 @@ macro_rules! modelled_kernel { } }; - (__variadic runtime, $trait:ident, $trait_fn:ident, $op:ident, $plc:ty, $([$($attr_id:ident: $attr_ty:ty),+])? ($t0:ty) -> $u:ty => $($kp:tt)+) => { + (__variadic runtime, $trait:ident, $trait_fn:ident, $op:ident, $plc:ty, $([$($attr_id:ident: $attr_ty:ty),+])? vec[$ts:ty] -> $u:ty => $($kp:tt)+) => { impl crate::kernels::VariadicKernel< crate::execution::SymbolicSession, $plc, @@ -4602,11 +4601,11 @@ macro_rules! modelled_kernel { $($($attr_id:$attr_ty),*,)? xs: &[<$ts as crate::computation::SymbolicType>::Type] ) -> <$u as crate::computation::SymbolicType>::Type { - use crate::computation::{KnownType, UnarySignature}; + use crate::computation::{KnownType, VariadicSignature}; use crate::execution::{Session, SymbolicSession}; use std::convert::TryInto; - let sig = UnarySignature { + let sig = VariadicSignature { args: <$ts as KnownType>::TY, ret: <$u as KnownType>::TY, }; From 87a75b52a2afd0ecc7638cc236f07ae9d1d096d7 Mon Sep 17 00:00:00 2001 From: Morten Dahl Date: Tue, 8 Feb 2022 14:51:48 +0100 Subject: [PATCH 2/2] concatenate --- moose/src/kernels/shapes.rs | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/moose/src/kernels/shapes.rs b/moose/src/kernels/shapes.rs index 1fe4bad57..03cac5eb1 100644 --- a/moose/src/kernels/shapes.rs +++ b/moose/src/kernels/shapes.rs @@ -100,27 +100,9 @@ pub trait PlacementConcatenate { fn concatenate(&self, sess: &S, axis: u32, xs: &[TS]) -> O; } -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[Tensor] -> Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[Float32Tensor] -> Float32Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[Float64Tensor] -> Float64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostFloat32Tensor] -> HostFloat32Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostFloat64Tensor] -> HostFloat64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostInt8Tensor] -> HostInt8Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostInt16Tensor] -> HostInt16Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostInt32Tensor] -> HostInt32Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostInt64Tensor] -> HostInt64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostRing64Tensor] -> HostRing64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, HostPlacement, attributes[axis: u32] vec[HostRing128Tensor] -> HostRing128Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[Tensor] -> Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[Fixed64Tensor] -> Fixed64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[Fixed128Tensor] -> Fixed128Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[ReplicatedFixed64Tensor] -> ReplicatedFixed64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[ReplicatedFixed128Tensor] -> ReplicatedFixed128Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[ReplicatedRing64Tensor] -> ReplicatedRing64Tensor, ConcatOp); -modelled!(PlacementConcatenate::concatenate, ReplicatedPlacement, attributes[axis: u32] vec[ReplicatedRing128Tensor] -> ReplicatedRing128Tensor, ConcatOp); - -kernel! { - ConcatOp, [ +modelled_kernel! { + PlacementConcatenate::concatenate, ConcatOp{axis: u32}, + [ (HostPlacement, vec[Tensor] -> Tensor => [concrete] attributes[axis] Self::logical_host_kernel), (HostPlacement, vec[Float32Tensor] -> Float32Tensor => [concrete] attributes[axis] Self::float_host_kernel), (HostPlacement, vec[Float64Tensor] -> Float64Tensor => [concrete] attributes[axis] Self::float_host_kernel),