Add set_primitive_cache_capacity

boydjohnson · Dec 31, 2024 · f78af1a · f78af1a
1 parent bcc457c
commit f78af1a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/benches/binary.rs b/benches/binary.rs
@@ -16,6 +16,7 @@ use {
             config::binary::{Binary, ForwardBinaryConfig},
             ExecArg, ForwardBinary, Primitive, PropForwardInference,
         },
+        set_primitive_cache_capacity,
         stream::Stream,
     },
     onednnl_sys::{dnnl_data_type_t::dnnl_f32, DNNL_ARG_DST, DNNL_ARG_SRC_0, DNNL_ARG_SRC_1},
@@ -27,6 +28,8 @@ use {
 fn binary_add(b: &mut Bencher) {
     let engine = Engine::new(Engine::CPU, 0).unwrap();
 
+    set_primitive_cache_capacity(2).unwrap();
+
     let stream = Arc::new(Stream::new(engine.clone()).unwrap());
 
     let src0_desc = MemoryDescriptor::new::<1, x>([3], DataType::F32).unwrap();

diff --git a/src/lib.rs b/src/lib.rs
@@ -4,4 +4,15 @@ pub mod memory;
 pub mod primitive;
 pub mod stream;
 
+use error::DnnlError;
 pub use onednnl_sys;
+
+pub fn set_primitive_cache_capacity(capacity: std::ffi::c_int) -> Result<(), DnnlError> {
+    let status = unsafe { onednnl_sys::dnnl_set_primitive_cache_capacity(capacity) };
+
+    if status == onednnl_sys::dnnl_status_t::dnnl_success {
+        Ok(())
+    } else {
+        Err(status.into())
+    }
+}