Merge pull request #21 from boydjohnson/feature/more-features

Have MemoryDescriptors owned by PrimitiveConfigs/PrimitiveDescriptor

Author: boydjohnson

benches/binary.rs

@@ -36,16 +36,17 @@ fn binary_add(b: &mut Bencher) {
 
     let binary_config = ForwardBinaryConfig {
         alg_kind: Binary::ADD,
-        src0_desc: &src0_desc,
-        src1_desc: &src1_desc,
-        dst_desc: &dst_desc,
-        attr: &PrimitiveAttributes::new().unwrap(),
+        src0_desc: src0_desc.clone_desc().unwrap(),
+        src1_desc: src1_desc.clone_desc().unwrap(),
+        dst_desc: dst_desc.clone_desc().unwrap(),
+        attr: PrimitiveAttributes::new().unwrap(),
     };
 
-    let primitive =
-        Primitive::new::<_, PropForwardInference, ForwardBinary<_>>(binary_config, engine.clone());
+    let primitive = Primitive::<_, PropForwardInference, ForwardBinaryConfig>::new::<
+        ForwardBinary<_>,
+    >(binary_config, engine.clone());
     assert!(primitive.is_ok());
-    let primitive = primitive.unwrap();
+    let mut primitive = primitive.unwrap();
 
     let s0_buffer = AlignedBuffer::new(&[4.0f32, 5.0, 6.0]).unwrap().into();
 
@@ -87,7 +88,7 @@ fn binary_add(b: &mut Bencher) {
 
         assert!(stream.wait().is_ok());
 
-        assert_eq!(result, Ok(()));
+        assert!(result.is_ok());
 
         assert_eq!(dst_memory.to_vec(), Ok(vec![5.0, 7.0, 9.0]));
     });

src/primitive.rs

@@ -108,13 +108,13 @@ impl PropType<Backward> for PropBackwardData {
     const KIND: dnnl_prop_kind_t::Type = dnnl_prop_kind_t::dnnl_backward_data;
 }
 
-pub struct Primitive {
+pub struct Primitive<'a, D: Direction, P: PropType<D>, C: PrimitiveConfig<'a, D, P>> {
     pub handle: dnnl_primitive_t,
-    pub desc: PrimitiveDescriptor,
+    pub desc: Option<PrimitiveDescriptor<'a, D, P, C>>,
     pub engine: Arc<Engine>,
 }
 
-impl Primitive {
+impl<'a, D: Direction, P: PropType<D>, C: PrimitiveConfig<'a, D, P>> Primitive<'a, D, P, C> {
     /// Creates a new `Primitive`.
     ///
     /// # Example
@@ -141,45 +141,50 @@ impl Primitive {
     /// // Define a forward binary config
     /// let binary_config = ForwardBinaryConfig {
     ///     alg_kind: dnnl_alg_kind_t::dnnl_binary_add, // Example: addition operation
-    ///     src0_desc: &src0_desc,
-    ///     src1_desc: &src1_desc,
-    ///     dst_desc: &dst_desc,
-    ///     attr: &PrimitiveAttributes::new().unwrap(),
+    ///     src0_desc: src0_desc,
+    ///     src1_desc: src1_desc,
+    ///     dst_desc: dst_desc,
+    ///     attr: PrimitiveAttributes::new().unwrap(),
     /// };
     ///
-    /// let primitive =
-    ///     Primitive::new::<_, PropForwardInference, ForwardBinary<_>>(binary_config, engine);
+    /// let primitive = Primitive::<_, PropForwardInference, ForwardBinaryConfig>::new::<
+    ///     ForwardBinary<_>,
+    /// >(binary_config, engine);
     ///
     /// assert!(primitive.is_ok());
     /// ```
-    pub fn new<'a, D: Direction, P: PropType<D>, O: Operation<'a, D, P>>(
+    pub fn new<O: Operation<'a, D, P, OperationConfig = C>>(
         config: O::OperationConfig,
         engine: Arc<Engine>,
-    ) -> Result<Primitive, DnnlError> {
+    ) -> Result<Primitive<'a, D, P, C>, DnnlError> {
         let desc = config.create_primitive_desc(engine.clone())?;
         Self::from_descriptor(desc, engine)
     }
 
     pub fn from_descriptor(
-        desc: PrimitiveDescriptor,
+        desc: PrimitiveDescriptor<'a, D, P, C>,
         engine: Arc<Engine>,
-    ) -> Result<Primitive, DnnlError> {
+    ) -> Result<Primitive<'a, D, P, C>, DnnlError> {
         let mut handle = std::ptr::null_mut();
 
         let status = unsafe { dnnl_primitive_create(&mut handle, desc.handle) };
 
         if status == dnnl_status_t::dnnl_success {
-            Ok(Self {
+            Ok(Primitive::<'a, D, P, C> {
                 handle,
-                desc,
+                desc: Some(desc),
                 engine,
             })
         } else {
             Err(status.into())
         }
     }
 
-    pub fn execute<T>(&self, stream: &Stream, args: Vec<ExecArg<'_, T>>) -> Result<(), DnnlError> {
+    pub fn execute<T>(
+        &mut self,
+        stream: &Stream,
+        args: Vec<ExecArg<'_, T>>,
+    ) -> Result<Option<PrimitiveDescriptor<'a, D, P, C>>, DnnlError> {
         let c_args: Vec<dnnl_exec_arg_t> = args
             .iter()
             .map(|arg| dnnl_exec_arg_t {
@@ -198,14 +203,16 @@ impl Primitive {
         };
 
         if status == dnnl_status_t::dnnl_success {
-            Ok(())
+            Ok(self.desc.take())
         } else {
             Err(status.into())
         }
     }
 }
 
-impl Drop for Primitive {
+impl<'a, D: Direction, P: PropType<D>, C: PrimitiveConfig<'a, D, P>> Drop
+    for Primitive<'a, D, P, C>
+{
     fn drop(&mut self) {
         unsafe {
             dnnl_primitive_destroy(self.handle);

src/primitive/config.rs

@@ -4,6 +4,9 @@ use {
     std::sync::Arc,
 };
 
-pub trait PrimitiveConfig<'a, D: Direction, P: PropType<D>> {
-    fn create_primitive_desc(&self, engine: Arc<Engine>) -> Result<PrimitiveDescriptor, DnnlError>;
+pub trait PrimitiveConfig<'a, D: Direction, P: PropType<D>>: Sized {
+    fn create_primitive_desc(
+        self,
+        engine: Arc<Engine>,
+    ) -> Result<PrimitiveDescriptor<'a, D, P, Self>, DnnlError>;
 }

src/primitive/descriptor.rs

@@ -2,14 +2,26 @@ use {
     super::{config::PrimitiveConfig, Direction, Operation, PropType},
     crate::{engine::Engine, error::DnnlError},
     onednnl_sys::{dnnl_primitive_desc_destroy, dnnl_primitive_desc_t},
-    std::sync::Arc,
+    std::{marker::PhantomData, sync::Arc},
 };
 
-pub struct PrimitiveDescriptor {
+pub struct PrimitiveDescriptor<
+    'a,
+    D: Direction,
+    P: PropType<D>,
+    C: PrimitiveConfig<'a, D, P> + Sized,
+> {
     pub handle: dnnl_primitive_desc_t,
+    pub config: C,
+
+    pub(crate) _marker_a: PhantomData<&'a ()>,
+    pub(crate) _marker_d: PhantomData<D>,
+    pub(crate) _marker_p: PhantomData<P>,
 }
 
-impl PrimitiveDescriptor {
+impl<'a, D: Direction, P: PropType<D>, C: PrimitiveConfig<'a, D, P> + Sized>
+    PrimitiveDescriptor<'a, D, P, C>
+{
     /// Creates a new `PrimitiveDescriptor`.
     ///
     /// # Example
@@ -37,30 +49,30 @@ impl PrimitiveDescriptor {
     /// // Define a forward binary config
     /// let binary_config = ForwardBinaryConfig {
     ///     alg_kind: dnnl_alg_kind_t::dnnl_binary_add, // Example: addition operation
-    ///     src0_desc: &src0_desc,
-    ///     src1_desc: &src1_desc,
-    ///     dst_desc: &dst_desc,
-    ///     attr: &PrimitiveAttributes::new().unwrap(),
+    ///     src0_desc: src0_desc,
+    ///     src1_desc: src1_desc,
+    ///     dst_desc: dst_desc,
+    ///     attr: PrimitiveAttributes::new().unwrap(),
     /// };
     ///
     /// // Create a new PrimitiveDescriptor for the forward binary operation
-    /// let primitive_descriptor = PrimitiveDescriptor::new::<
-    ///     Forward,
-    ///     _,
+    /// let primitive_descriptor = PrimitiveDescriptor::<_, _, ForwardBinaryConfig>::new::<
     ///     ForwardBinary<PropForwardInference>,
     /// >(binary_config, engine);
     ///
     /// assert!(primitive_descriptor.is_ok());
     /// ```
-    pub fn new<'a, D: Direction, P: PropType<D>, O: Operation<'a, D, P>>(
+    pub fn new<O: Operation<'a, D, P, OperationConfig = C>>(
         config: O::OperationConfig,
         engine: Arc<Engine>,
-    ) -> Result<Self, DnnlError> {
+    ) -> Result<PrimitiveDescriptor<'a, D, P, C>, DnnlError> {
         config.create_primitive_desc(engine)
     }
 }
 
-impl Drop for PrimitiveDescriptor {
+impl<'a, D: Direction, P: PropType<D>, C: PrimitiveConfig<'a, D, P>> Drop
+    for PrimitiveDescriptor<'a, D, P, C>
+{
     fn drop(&mut self) {
         unsafe { dnnl_primitive_desc_destroy(self.handle) };
     }

src/primitives/au_gru.rs

@@ -5,32 +5,36 @@ use {
         memory::descriptor::MemoryDescriptor,
         primitive::{
             attributes::PrimitiveAttributes, config::PrimitiveConfig,
-            descriptor::PrimitiveDescriptor, Backward, Forward, Operation, OperationType, PropType,
+            descriptor::PrimitiveDescriptor, Backward, Forward, Operation, OperationType,
+            PropForwardTraining, PropType,
         },
     },
     onednnl_sys::{
         dnnl_augru_backward_primitive_desc_create, dnnl_augru_forward_primitive_desc_create,
-        dnnl_primitive_attr_t, dnnl_rnn_direction_t, dnnl_status_t,
+        dnnl_rnn_direction_t, dnnl_status_t,
     },
-    std::{ffi::c_uint, sync::Arc},
+    std::{ffi::c_uint, marker::PhantomData, sync::Arc},
 };
 
-pub struct ForwardAuGruConfig<'a> {
+pub struct ForwardAuGruConfig {
     direction: dnnl_rnn_direction_t::Type,
-    src_layer_desc: &'a MemoryDescriptor,
-    src_iter_desc: &'a MemoryDescriptor,
-    attention_desc: &'a MemoryDescriptor,
-    weights_layer_desc: &'a MemoryDescriptor,
-    weights_iter_desc: &'a MemoryDescriptor,
-    bias_desc: &'a MemoryDescriptor,
-    dst_layer_desc: &'a MemoryDescriptor,
-    dst_iter_desc: &'a MemoryDescriptor,
+    src_layer_desc: MemoryDescriptor,
+    src_iter_desc: MemoryDescriptor,
+    attention_desc: MemoryDescriptor,
+    weights_layer_desc: MemoryDescriptor,
+    weights_iter_desc: MemoryDescriptor,
+    bias_desc: MemoryDescriptor,
+    dst_layer_desc: MemoryDescriptor,
+    dst_iter_desc: MemoryDescriptor,
     flags: c_uint,
-    attr: &'a PrimitiveAttributes,
+    attr: PrimitiveAttributes,
 }
 
-impl<'a, P: PropType<Forward>> PrimitiveConfig<'a, Forward, P> for ForwardAuGruConfig<'a> {
-    fn create_primitive_desc(&self, engine: Arc<Engine>) -> Result<PrimitiveDescriptor, DnnlError> {
+impl<'a, P: PropType<Forward>> PrimitiveConfig<'a, Forward, P> for ForwardAuGruConfig {
+    fn create_primitive_desc(
+        self,
+        engine: Arc<Engine>,
+    ) -> Result<PrimitiveDescriptor<'a, Forward, P, ForwardAuGruConfig>, DnnlError> {
         let mut handle = std::ptr::null_mut();
         let status = unsafe {
             dnnl_augru_forward_primitive_desc_create(
@@ -52,7 +56,14 @@ impl<'a, P: PropType<Forward>> PrimitiveConfig<'a, Forward, P> for ForwardAuGruC
         };
 
         if status == dnnl_status_t::dnnl_success {
-            Ok(PrimitiveDescriptor { handle })
+            Ok(PrimitiveDescriptor::<'a, Forward, P, ForwardAuGruConfig> {
+                handle,
+                config: self,
+
+                _marker_a: PhantomData,
+                _marker_d: PhantomData,
+                _marker_p: PhantomData,
+            })
         } else {
             Err(status.into())
         }
@@ -61,29 +72,32 @@ impl<'a, P: PropType<Forward>> PrimitiveConfig<'a, Forward, P> for ForwardAuGruC
 
 pub struct BackwardAuGruConfig<'a> {
     direction: dnnl_rnn_direction_t::Type,
-    src_layer_desc: &'a MemoryDescriptor,
-    src_iter_desc: &'a MemoryDescriptor,
-    attention_desc: &'a MemoryDescriptor,
-    weights_layer_desc: &'a MemoryDescriptor,
-    weights_iter_desc: &'a MemoryDescriptor,
-    bias_desc: &'a MemoryDescriptor,
-    dst_layer_desc: &'a MemoryDescriptor,
-    dst_iter_desc: &'a MemoryDescriptor,
-    diff_src_layer_desc: &'a MemoryDescriptor,
-    diff_src_iter_desc: &'a MemoryDescriptor,
-    diff_attention_desc: &'a MemoryDescriptor,
-    diff_weights_layer_desc: &'a MemoryDescriptor,
-    diff_weights_iter_desc: &'a MemoryDescriptor,
-    diff_bias_desc: &'a MemoryDescriptor,
-    diff_dst_layer_desc: &'a MemoryDescriptor,
-    diff_dst_iter_desc: &'a MemoryDescriptor,
+    src_layer_desc: MemoryDescriptor,
+    src_iter_desc: MemoryDescriptor,
+    attention_desc: MemoryDescriptor,
+    weights_layer_desc: MemoryDescriptor,
+    weights_iter_desc: MemoryDescriptor,
+    bias_desc: MemoryDescriptor,
+    dst_layer_desc: MemoryDescriptor,
+    dst_iter_desc: MemoryDescriptor,
+    diff_src_layer_desc: MemoryDescriptor,
+    diff_src_iter_desc: MemoryDescriptor,
+    diff_attention_desc: MemoryDescriptor,
+    diff_weights_layer_desc: MemoryDescriptor,
+    diff_weights_iter_desc: MemoryDescriptor,
+    diff_bias_desc: MemoryDescriptor,
+    diff_dst_layer_desc: MemoryDescriptor,
+    diff_dst_iter_desc: MemoryDescriptor,
     flags: c_uint,
-    hint_fwd_pd: &'a PrimitiveDescriptor,
-    attr: dnnl_primitive_attr_t,
+    hint_fwd_pd: &'a PrimitiveDescriptor<'a, Forward, PropForwardTraining, ForwardAuGruConfig>,
+    attr: PrimitiveAttributes,
 }
 
 impl<'a, P: PropType<Backward>> PrimitiveConfig<'a, Backward, P> for BackwardAuGruConfig<'a> {
-    fn create_primitive_desc(&self, engine: Arc<Engine>) -> Result<PrimitiveDescriptor, DnnlError> {
+    fn create_primitive_desc(
+        self,
+        engine: Arc<Engine>,
+    ) -> Result<PrimitiveDescriptor<'a, Backward, P, BackwardAuGruConfig<'a>>, DnnlError> {
         let mut handle = std::ptr::null_mut();
         let status = unsafe {
             dnnl_augru_backward_primitive_desc_create(
@@ -109,12 +123,20 @@ impl<'a, P: PropType<Backward>> PrimitiveConfig<'a, Backward, P> for BackwardAuG
                 self.diff_dst_iter_desc.handle,
                 self.flags,
                 self.hint_fwd_pd.handle,
-                self.attr,
+                self.attr.handle,
             )
         };
 
         if status == dnnl_status_t::dnnl_success {
-            Ok(PrimitiveDescriptor { handle })
+            Ok(
+                PrimitiveDescriptor::<'a, Backward, P, BackwardAuGruConfig<'a>> {
+                    handle,
+                    config: self,
+                    _marker_a: PhantomData,
+                    _marker_d: PhantomData,
+                    _marker_p: PhantomData,
+                },
+            )
         } else {
             Err(status.into())
         }
@@ -125,9 +147,9 @@ pub struct ForwardAuGru<P: PropType<Forward>> {
     pub prop_type: P,
 }
 
-impl<'a, P: PropType<Forward>> Operation<'a, Forward, P> for ForwardAuGru<P> {
+impl<P: PropType<Forward>> Operation<'_, Forward, P> for ForwardAuGru<P> {
     const TYPE: OperationType = OperationType::Augru;
-    type OperationConfig = ForwardAuGruConfig<'a>;
+    type OperationConfig = ForwardAuGruConfig;
 }
 
 pub struct BackwardAuGru<P: PropType<Backward>> {