Add enough of the graph API to do mobilenet

.github/workflows/ci-intel.yml

@@ -1,64 +1,64 @@
 name: Build onednnl safe bindings for Intel
 
 on:
-  push:
-    branches:
-      - main
-  pull_request:
-    branches:
-      - main
+    push:
+        branches:
+            - main
+    pull_request:
+        branches:
+            - main
 
 jobs:
-  build:
-    runs-on: ubuntu-24.04
-    env:
-      RUSTFLAGS: -D warnings
-    strategy:
-      matrix:
-        rust:
-          - stable
-          - nightly
+    build:
+        runs-on: ubuntu-24.04
+        env:
+            RUSTFLAGS: -D warnings
+        strategy:
+            matrix:
+                rust:
+                    - stable
+                    - nightly
 
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
+        steps:
+            - name: Checkout code
+              uses: actions/checkout@v4
 
-      - name: Install Rust toolchain
-        uses: dtolnay/rust-toolchain@master
-        with:
-          toolchain: ${{ matrix.rust }}
+            - name: Install Rust toolchain
+              uses: dtolnay/rust-toolchain@master
+              with:
+                  toolchain: ${{ matrix.rust }}
 
-      - name: Cache oneAPI installation
-        id: cache-oneapi
-        uses: actions/cache@v3
-        with:
-          path: /opt/intel/oneapi
-          key: oneapi-${{ env.CACHE_NUMBER }}
-          restore-keys: |
-            oneapi-
+            - name: Cache oneAPI installation
+              id: cache-oneapi
+              uses: actions/cache@v3
+              with:
+                  path: /opt/intel/oneapi
+                  key: oneapi-${{ env.CACHE_NUMBER }}
+                  restore-keys: |
+                      oneapi-
 
-      - name: Set up Intel oneAPI APT repository
-        if: steps.cache-oneapi.outputs.cache-hit != 'true'
-        run: |
-          wget -qO- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | sudo gpg --dearmor -o /usr/share/keyrings/oneapi-archive-keyring.gpg
-          echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
-          sudo apt update
+            - name: Set up Intel oneAPI APT repository
+              if: steps.cache-oneapi.outputs.cache-hit != 'true'
+              run: |
+                  wget -qO- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | sudo gpg --dearmor -o /usr/share/keyrings/oneapi-archive-keyring.gpg
+                  echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+                  sudo apt update
 
-      - name: Install Intel oneAPI DNNL version 2025.0.1
-        if: steps.cache-oneapi.outputs.cache-hit != 'true'
-        run: sudo apt install -y intel-oneapi-dnnl-devel=2025.0.1-*
+            - name: Install Intel oneAPI OneDNN package 2025.2
+              if: steps.cache-oneapi.outputs.cache-hit != 'true'
+              run: sudo apt install -y intel-oneapi-dnnl-devel-2025.2
 
-      - name: Build
-        run: |
-          source /opt/intel/oneapi/setvars.sh
-          cargo build
+            - name: Build
+              run: |
+                  source /opt/intel/oneapi/setvars.sh
+                  cargo build
 
-      - name: Run tests
-        run: |
-          source /opt/intel/oneapi/setvars.sh
-          cargo test
-      
-      - name: Run docs
-        run: |
-          source /opt/intel/oneapi/setvars.sh
-          cargo doc
+            - name: Run tests
+              run: |
+                  source /opt/intel/oneapi/setvars.sh
+                  cargo test
+
+            - name: Run docs
+              run: |
+                  source /opt/intel/oneapi/setvars.sh
+                  cargo doc

src/graph.rs

@@ -0,0 +1,7 @@
+pub mod compiled_partition;
+pub mod graph;
+pub mod op;
+pub mod ops_builders;
+pub mod partition;
+pub mod spec;
+pub mod tensor;

src/graph/compiled_partition.rs

@@ -0,0 +1,71 @@
+use onednnl_sys::{
+    dnnl_graph_compiled_partition_create, dnnl_graph_compiled_partition_destroy,
+    dnnl_graph_compiled_partition_execute, dnnl_graph_compiled_partition_t, dnnl_status_t,
+};
+
+use crate::{
+    error::DnnlError,
+    graph::{partition::OneDNNGraphPartition, tensor::tensor::Tensor},
+    stream::Stream,
+};
+
+pub struct CompiledPartition {
+    pub(crate) handle: dnnl_graph_compiled_partition_t,
+    pub(crate) partition: OneDNNGraphPartition,
+}
+
+impl CompiledPartition {
+    pub fn create(partition: OneDNNGraphPartition) -> Result<Self, DnnlError> {
+        let mut handle = std::ptr::null_mut();
+        let status = unsafe { dnnl_graph_compiled_partition_create(&mut handle, partition.handle) };
+        if status != dnnl_status_t::dnnl_success {
+            return Err(status.into());
+        }
+        Ok(CompiledPartition { handle, partition })
+    }
+
+    pub fn execute(
+        &self,
+        stream: &Stream,
+        inputs: &[Tensor],
+        outputs: &[&mut Tensor],
+    ) -> Result<(), DnnlError> {
+        // Collect the input tensor handles into a vector. This ensures the collection
+        // of pointers has a stable memory location that lives long enough for the C call.
+        let mut input_handles: Vec<_> = inputs.iter().map(|t| t.handle as *const _).collect();
+
+        // Do the same for the output tensor handles.
+        let mut output_handles: Vec<_> = outputs.iter().map(|t| t.handle as *const _).collect();
+
+        // The C API expects the number of inputs/outputs as an integer type.
+        let num_inputs = input_handles.len();
+        let num_outputs = output_handles.len();
+
+        let status = unsafe {
+            // Now, we pass pointers to our vectors' data, which are guaranteed to be
+            // valid for the duration of this call.
+            dnnl_graph_compiled_partition_execute(
+                self.handle,
+                stream.handle,
+                num_inputs,
+                input_handles.as_mut_ptr(),
+                num_outputs,
+                output_handles.as_mut_ptr(), // The C API uses these handles to find the output buffers
+            )
+        };
+
+        if status != dnnl_status_t::dnnl_success {
+            return Err(status.into());
+        }
+
+        Ok(())
+    }
+}
+
+impl Drop for CompiledPartition {
+    fn drop(&mut self) {
+        unsafe {
+            dnnl_graph_compiled_partition_destroy(self.handle);
+        }
+    }
+}

src/graph/graph.rs

@@ -0,0 +1,141 @@
+use {
+    super::op::OneDNNGraphOp,
+    crate::{error::DnnlError, graph::partition::OneDNNGraphPartition},
+    onednnl_sys::{
+        dnnl_graph_add_op, dnnl_graph_graph_create, dnnl_graph_graph_create_with_fpmath_mode,
+        dnnl_graph_graph_destroy, dnnl_graph_graph_filter, dnnl_graph_graph_finalize,
+        dnnl_graph_graph_get_fpmath_mode, dnnl_graph_graph_get_partition_num,
+        dnnl_graph_graph_get_partitions, dnnl_graph_graph_is_finalized, dnnl_graph_graph_t,
+        dnnl_status_t,
+    },
+};
+
+pub struct OneDNNGraph {
+    handle: dnnl_graph_graph_t,
+    ops: Vec<OneDNNGraphOp>,
+}
+
+impl OneDNNGraph {
+    pub fn new(engine_type: onednnl_sys::dnnl_engine_kind_t::Type) -> Result<Self, DnnlError> {
+        let mut handle: dnnl_graph_graph_t = std::ptr::null_mut();
+        let status = unsafe { dnnl_graph_graph_create(&mut handle, engine_type) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok(Self {
+                handle,
+                ops: Vec::new(),
+            })
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn new_with_fpmath_mode(
+        engine_type: onednnl_sys::dnnl_engine_kind_t::Type,
+        fp_mode: onednnl_sys::dnnl_fpmath_mode_t::Type,
+    ) -> Result<Self, DnnlError> {
+        let mut handle: dnnl_graph_graph_t = std::ptr::null_mut();
+        let status =
+            unsafe { dnnl_graph_graph_create_with_fpmath_mode(&mut handle, engine_type, fp_mode) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok(Self {
+                handle,
+                ops: Vec::new(),
+            })
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn filter(
+        &self,
+        policy: onednnl_sys::dnnl_graph_partition_policy_t::Type,
+    ) -> Result<(), DnnlError> {
+        let status = unsafe { dnnl_graph_graph_filter(self.handle, policy) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok(())
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn finalize(&self) -> Result<(), DnnlError> {
+        let status = unsafe { dnnl_graph_graph_finalize(self.handle) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok(())
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn ops(&self) -> &[OneDNNGraphOp] {
+        &self.ops
+    }
+
+    pub fn is_finalized(&self) -> Result<bool, DnnlError> {
+        let mut is_finalized = 0;
+        let status = unsafe { dnnl_graph_graph_is_finalized(self.handle, &mut is_finalized) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok(is_finalized != 0)
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn get_fpmath_mode(
+        &self,
+    ) -> Result<(onednnl_sys::dnnl_fpmath_mode_t::Type, i32), DnnlError> {
+        let mut mode = onednnl_sys::dnnl_fpmath_mode_t::dnnl_fpmath_mode_strict;
+        let mut apply_to_int = 0;
+
+        let status =
+            unsafe { dnnl_graph_graph_get_fpmath_mode(self.handle, &mut mode, &mut apply_to_int) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok((mode, apply_to_int))
+        } else {
+            Err(status.into())
+        }
+    }
+    pub fn get_partition_num(&self) -> Result<usize, DnnlError> {
+        let mut num = 0;
+
+        let status = unsafe { dnnl_graph_graph_get_partition_num(self.handle, &mut num) };
+        if status == dnnl_status_t::dnnl_success {
+            Ok(num)
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn get_partitions(&self) -> Result<Vec<OneDNNGraphPartition>, DnnlError> {
+        let num = self.get_partition_num()?;
+        let mut partitions = Vec::with_capacity(num);
+
+        let status =
+            unsafe { dnnl_graph_graph_get_partitions(self.handle, num, partitions.as_mut_ptr()) };
+        if status == dnnl_status_t::dnnl_success {
+            unsafe { partitions.set_len(partitions.capacity()) };
+            Ok(partitions
+                .into_iter()
+                .map(|p| OneDNNGraphPartition { handle: p })
+                .collect())
+        } else {
+            Err(status.into())
+        }
+    }
+
+    pub fn add_op(&mut self, op: OneDNNGraphOp) -> Result<(), DnnlError> {
+        let status = unsafe { dnnl_graph_add_op(self.handle, op.handle) };
+        self.ops.push(op);
+        if status == dnnl_status_t::dnnl_success {
+            Ok(())
+        } else {
+            Err(status.into())
+        }
+    }
+}
+
+impl Drop for OneDNNGraph {
+    fn drop(&mut self) {
+        unsafe { dnnl_graph_graph_destroy(self.handle) };
+    }
+}