pxl_rust/runtime/

map.rs

//! Abstracts data-parallel computation based on a map operation for a given job.
//!
//! `Map` module enables efficient execution of the same device kernel across
//! multiple chunks of input data in parallel. It is designed to work within a job
//! context, leveraging hardware-level parallelism in XCENA device.
//!
//! It applies the same device kernel across each chunk of the input data,
//! leveraging hardware-level parallelism for performance.
//!
//! This module provides a Rust abstraction around the C++ `pxl::runtime::Map` class.
use crate::ffi::ArgInfo_t;
use crate::ffi::{self, PxlResult};
use crate::runtime::Stream;
use std::ffi::c_void;
use std::pin::Pin;

/// Abstracts data-parallel computation based on a map operation for a given job.
///
/// This struct wraps a FFI pointer to the C++ `pxl::runtime::Map` class.
/// It provides methods for executing parallel processing, synchronizing operations,
/// and setting up asynchronous callbacks and stream.
pub struct Map {
    /// Pointer to the underlying FFI `Map` object.
    inner: *mut ffi::Map,
    /// Pointer to the underlying FFI `Stream` object.
    stream_: *mut ffi::Stream,
}

/// Implementation of the `Map` struct.
impl Map {
    /// Creates a new `Map` wrapper from a raw FFI pointer.
    /// This is called in runtime wrapper APIs requiring an actual `Map` object.
    /// # Arguments
    /// * `ptr` - A raw pointer to the underlying FFI `Map` object.
    /// # Safety
    /// Caller must ensure that `ptr` is valid and remains valid for the lifetime of `Map`.
    pub(crate) fn new(ptr: *mut ffi::Map) -> Self {
        Map {
            inner: ptr,
            stream_: unsafe { ffi::createStream() },
        }
    }

    pub(crate) fn destroy(&mut self) {
        if !self.inner.is_null() {
            unsafe { ffi::destroy_map(self.inner) };
            unsafe { ffi::destroyStream(self.stream_) };
            self.inner = std::ptr::null_mut();
            self.stream_ = std::ptr::null_mut();
        }
    }

    /// gets the raw pointer of the underlying FFI `Map` object.
    pub fn get(&self) -> *mut ffi::Map {
        self.inner
    }

    /// Sets success callback function with user data.
    /// # Arguments
    /// * `callback` - Callback function as a c_void pointer.
    /// * `user_data` - User data as a c_void pointer.
    /// # Safety
    /// Caller must ensure that the callback, user data remain valid for the lifetime of `Map`.
    /// # Example
    /// ```
    /// use std::ffi::c_void;
    /// use std::sync::{
    ///     atomic::{AtomicBool, Ordering},
    ///     Arc,
    /// };    ///
    /// unsafe extern "C" fn completion_callback(ptr: *mut c_void) {
    ///     let flag = ptr as *const AtomicBool;
    ///     println!("✅ Completion callback called with ptr: {:?}", flag);
    ///     (*flag).store(true, Ordering::SeqCst);
    /// }
    /// let done_flag = Arc::new(AtomicBool::new(false));
    /// let done_ptr = Arc::into_raw(done_flag.clone()) as *mut c_void;
    /// map.set_completion_callback(completion_callback as *mut c_void, done_ptr);
    /// ```
    pub fn set_completion_callback(&mut self, callback: *mut c_void, user_data: *mut c_void) {
        unsafe {
            ffi::map_set_completion_callback(
                self.inner,
                callback as *mut ffi::c_void,
                user_data as *mut ffi::c_void,
            );
        }
    }

    /// Sets error callback function with user data.
    /// # Arguments
    /// * `callback` - Callback function as a `c_void` pointer.
    /// * `user_data` - User data as a `c_void` pointer.
    /// # Safety
    /// Caller must ensure that the callback and user data remain valid for the lifetime of `Map`.
    /// # Example
    /// ```
    /// use std::ffi::c_void;
    /// use std::sync::{
    ///     atomic::{AtomicBool, Ordering},
    ///     Arc,
    /// };
    /// unsafe extern "C" fn error_callback(ptr: *mut c_void) {
    ///     let flag = ptr as *const AtomicBool;
    ///     println!("❌ Error callback called with ptr: {:?}", flag);
    ///     (*flag).store(true, Ordering::SeqCst);
    /// }
    /// let error_flag = Arc::new(AtomicBool::new(false));
    /// let error_ptr = Arc::into_raw(error_flag.clone()) as *mut c_void;
    /// map.set_error_callback(error_callback as *mut c_void, error_ptr);
    /// ```
    pub fn set_error_callback(&mut self, callback: *mut c_void, user_data: *mut c_void) {
        unsafe {
            ffi::map_set_error_callback(
                self.inner,
                callback as *mut ffi::c_void,
                user_data as *mut ffi::c_void,
            );
        }
    }

    /// Sets message callback function with user data.
    /// # Arguments
    /// * `callback` - Callback function as a `c_void` pointer.
    /// * `user_data` - User data as a `c_void` pointer.
    /// # Safety
    /// Caller must ensure that the callback and user data remain valid for the lifetime of `Map`.
    /// # Example
    /// ```
    /// use std::ffi::c_void;
    /// use std::sync::Arc;
    /// unsafe extern "C" fn message_callback(ptr: *mut c_void) {
    ///     println!("📩 Message callback called with ptr: {:?}", ptr);
    /// }
    /// let message_data = Arc::new(42); // Example user data
    /// let message_ptr = Arc::into_raw(message_data.clone()) as *mut c_void;
    /// map.set_message_callback(message_callback as *mut c_void, message_ptr);
    /// ```
    pub fn set_message_callback(&mut self, callback: *mut c_void, user_data: *mut c_void) {
        unsafe {
            ffi::map_set_message_callback(
                self.inner,
                callback as *mut ffi::c_void,
                user_data as *mut ffi::c_void,
            );
        }
    }

    pub fn set_batch_size(&mut self, batch_size: &u32) {
        unsafe {
            let inner_pin = Pin::new_unchecked(&mut *self.inner);
            inner_pin.setBatchSize(batch_size);
        }
    }

    /// Sets the cluster bitmap for the map operation.
    /// # Arguments
    /// * `cluster_bitmap` - The cluster bitmap to be used for the map operation (default = 0xF).
    /// # Example
    /// ```
    /// map.set_cluster_bitmap(0xF); // Use all clusters
    /// ```
    pub fn set_cluster_bitmap(&mut self, cluster_bitmap: u32) {
        unsafe {
            let inner_pin = Pin::new_unchecked(&mut *self.inner);
            inner_pin.setClusterBitmap(&cluster_bitmap);
        }
    }

    /// Sets the stream to the default stream.
    /// # Example
    /// ```
    /// map.set_default_stream()?;
    /// ```
    pub fn set_default_stream(&mut self) {
        unsafe {
            ffi::map_set_default_stream(self.inner);
        }
    }

    /// Sets the `Stream` for `Map` object.
    /// # Arguments
    /// * `stream` - A reference to the `Stream` object to be set.
    /// # Example
    /// ```rust
    /// let stream = pxl::runtime::create_stream();
    /// map.set_stream(&stream);
    /// ```
    pub fn set_stream(&mut self, stream: &Stream) {
        self.stream_ = stream.get();
    }

    /// Returns the stream ID that the map operation is associated with.
    /// # Returns
    /// The stream ID as a `u32`.
    /// # Example
    /// ```
    /// let stream_id = map.stream_id();
    /// println!("Current stream ID: {}", stream_id);
    /// ```
    pub fn stream_id(&self) -> u32 {
        unsafe { (*self.inner).streamId() }
    }

    /// Executes with the given arguments.
    /// This is called in execute! macro.
    /// # Arguments
    /// * `args` - A vector of `ArgInfo_t` objects representing the arguments to be passed.
    /// # Returns
    /// Execution status as a PxlResult.
    /// # Safety
    /// Caller must ensure that the `args` vector is valid and contains valid pointers.
    /// # Example
    /// Please refer to usage of `execute!` macro.
    pub fn execute(&mut self, args: &Vec<ArgInfo_t>) -> PxlResult {
        unsafe {
            // println!("Executing with arguments:");
            // for arg in args.iter() {
            //     println!("{:?}", arg);
            // }
            let args_ptr = args.as_ptr() as *mut ArgInfo_t;
            let num_args = args.len();
            ffi::map_execute(self.inner, args_ptr, num_args)
        }
    }

    /// Synchronizes the `Map` operation after execute request.
    /// # Returns
    /// Synchronization status as a PxlResult.
    /// # Safety
    /// Caller must ensure that execution status returned from `execute` is valid.
    /// # Example
    /// ```rust
    /// let ret = map.synchronize();
    /// ```
    pub fn synchronize(&mut self) -> PxlResult {
        unsafe { Pin::new_unchecked(&mut *self.inner).synchronize() }
    }

    /// Gets the execution status of the map operation.
    /// # Returns
    /// `ExecuteStatus` indicating the current execution status.
    /// # Example
    /// ```
    /// use pxl::ExecuteStatus;
    /// let status = map.get_execute_status();
    /// match status {
    ///     ExecuteStatus::ExecuteDone => println!("Execution completed successfully"),
    ///     ExecuteStatus::Fail => println!("Execution failed"),
    ///     _ => println!("Execution in progress or other status"),
    /// }
    /// ```
    pub fn get_execute_status(&self) -> ffi::ExecuteStatus {
        unsafe { (*self.inner).getExecuteStatus() }
    }

    /// Gets the error report of the map operation.
    /// # Returns
    /// `Vec<KernelError_t>` containing the error report.
    /// If successful, the beginIndex and endIndex will be 0xFFFFFFFF.
    /// If unsuccessful, they indicate the range of failed task index.
    /// # Example
    /// ```
    /// let kernel_errors = map.get_kernel_error();
    /// for error in &kernel_errors {
    ///     if error.beginIndex != 0xFFFFFFFF && error.endIndex != 0xFFFFFFFF {
    ///         println!("Error in range: {} - {}", error.beginIndex, error.endIndex);
    ///     }
    /// }
    /// ```
    pub fn get_kernel_error(&mut self) -> Vec<ffi::KernelError_t> {
        unsafe {
            let mut error_count: usize = 0;
            let error_ptr = ffi::map_get_kernel_error(self.inner, &mut error_count as *mut usize);

            if error_ptr.is_null() || error_count == 0 {
                return Vec::new();
            }

            let errors = std::slice::from_raw_parts(error_ptr, error_count);
            errors.to_vec()
        }
    }
}