singe-cutensor-sys 0.1.0-alpha.7

/* automatically generated by rust-bindgen 0.72.1 */

pub const CUTENSORMP_MAJOR: u32 = 0;
pub const CUTENSORMP_MINOR: u32 = 1;
pub const CUTENSORMP_PATCH: u32 = 0;
pub const CUTENSORMP_VERSION: u32 = 100;
#[repr(u32)]
#[derive(
    Debug,
    Copy,
    Clone,
    Hash,
    PartialOrd,
    Ord,
    PartialEq,
    Eq,
    TryFromPrimitive,
    IntoPrimitive,
)]
pub enum cutensorMpAlgo_t {
    /// Default algorithm.
    CUTENSORMP_ALGO_DEFAULT = 0,
}
#[repr(u32)]
#[derive(
    Debug,
    Copy,
    Clone,
    Hash,
    PartialOrd,
    Ord,
    PartialEq,
    Eq,
    TryFromPrimitive,
    IntoPrimitive,
)]
pub enum cutensorMpPlanAttribute_t {
    /// uint64_t: exact required workspace in bytes that is needed to execute the plan.
    CUTENSORMP_PLAN_REQUIRED_WORKSPACE_DEVICE = 0,
    /// uint64_t: exact required workspace in bytes that is needed to execute the plan.
    CUTENSORMP_PLAN_REQUIRED_WORKSPACE_HOST = 1,
}
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpHandle {
    _unused: [u8; 0],
}
pub type cutensorMpHandle_t = *mut cuTensorMpHandle;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpTensorDescriptor {
    _unused: [u8; 0],
}
pub type cutensorMpTensorDescriptor_t = *mut cuTensorMpTensorDescriptor;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpPlanPreference {
    _unused: [u8; 0],
}
pub type cutensorMpPlanPreference_t = *mut cuTensorMpPlanPreference;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpPlan {
    _unused: [u8; 0],
}
pub type cutensorMpPlan_t = *mut cuTensorMpPlan;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpOperationDescriptor {
    _unused: [u8; 0],
}
pub type cutensorMpOperationDescriptor_t = *mut cuTensorMpOperationDescriptor;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpContractionDescriptor {
    _unused: [u8; 0],
}
pub type cutensorMpContractionDescriptor_t = *mut cuTensorMpContractionDescriptor;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct cuTensorMpContractionPlan {
    _unused: [u8; 0],
}
pub type cutensorMpContractionPlan_t = *mut cuTensorMpContractionPlan;
unsafe extern "C" {
    /// Initializes the cutensorMp library and creates a handle for distributed tensor operations.
    ///
    /// This function creates a cutensorMp handle that serves as the context for all distributed tensor operations. The handle is associated with a specific MPI communicator, local CUDA device, and CUDA stream. This allows cutensorMp to coordinate tensor operations across multiple processes and GPUs.
    ///
    /// The communicator defines the group of processes that will participate in distributed tensor operations. The local device ID specifies which CUDA device on the current process will be used for computations. The CUDA stream enables asynchronous execution and synchronization with other CUDA operations.
    ///
    /// The user is responsible for calling cutensorMpDestroy to free the resources associated with the handle.
    ///
    /// Remark
    ///
    /// non-blocking, no reentrant, and thread-safe.
    ///
    /// # Parameters
    ///
    /// - `handle`: Pointer to cutensorMpHandle_t that will hold the created handle.
    /// - `comm`: NCCL communicator that defines the group of processes for distributed operations.
    /// - `local_device_id`: CUDA device ID to use on the current process (must be valid and accessible).
    /// - `stream`: CUDA stream for asynchronous operations.
    pub fn cutensorMpCreate(
        handle: *mut cutensorMpHandle_t,
        comm: ncclComm_t,
        local_device_id: ::core::ffi::c_int,
        stream: cudaStream_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Frees all resources associated with the provided cutensorMp handle.
    ///
    /// This function deallocates all memory and resources associated with a cutensorMp handle that was previously created by cutensorMpCreate. After calling this function, the handle becomes invalid and should not be used in subsequent cutensorMp operations.
    ///
    /// Remark
    ///
    /// blocking, no reentrant, and thread-safe.
    pub fn cutensorMpDestroy(handle: cutensorMpHandle_t) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Creates a distributed tensor descriptor for multi-process tensor operations.
    ///
    /// This function creates a tensor descriptor that defines the structure and distribution of a multi-dimensional tensor across multiple processes. Unlike regular cuTENSOR tensor descriptors, this descriptor includes information about how the tensor is partitioned and distributed across different processes in the MPI communicator.
    ///
    /// The tensor is described by its modes (dimensions), extents (sizes along each mode), and strides for elements and blocks. The distribution is specified through block sizes, block strides, and nranks-per-mode, which determine how the tensor data is partitioned across the participating processes.
    ///
    /// The user is responsible for calling cutensorMpDestroyTensorDescriptor to free the resources associated with the descriptor once it is no longer needed.
    ///
    /// Remark
    ///
    /// non-blocking, no reentrant, and thread-safe.
    ///
    /// # Parameters
    ///
    /// - `handle`: Opaque handle holding cutensorMp’s library context.
    /// - `desc`: Pointer to the address where the allocated tensor descriptor object will be stored.
    /// - `numModes`: Number of modes (dimensions) in the tensor (must be greater than zero).
    /// - `extent`: Extent (size) of each mode (size: numModes, all values must be greater than zero).
    /// - `elementStride`: Stride between consecutive elements in each mode (size: numModes).
    /// - `blockSize`: Size of each block along each mode for distribution (size: numModes), passing null will using extent\[i\]/nranksPerMode\[i\].
    /// - `blockStride`: Stride between consecutive blocks in each mode (size: numModes).
    /// - `nranksPerMode`: Number of processes along each mode (size: numModes).
    /// - `nranks`: Total number of ranks (processes) participating in the tensor distribution.
    /// - `ranks`: Array of rank IDs for each participating process (size: nranks), passing null will use the range [0, nranks).
    /// - `dataType`: Data type of the tensor elements.
    ///
    /// # Return value
    ///
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INVALID_VALUE`]: if some input data is invalid (this typically indicates a user error).
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_INITIALIZED`]: if the handle is not initialized.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_SUPPORTED`]: if the requested descriptor configuration is not supported.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_SUCCESS`]: The operation completed successfully.
    pub fn cutensorMpCreateTensorDescriptor(
        handle: cutensorMpHandle_t,
        desc: *mut cutensorMpTensorDescriptor_t,
        numModes: u32,
        extent: *const i64,
        elementStride: *const i64,
        blockSize: *const i64,
        blockStride: *const i64,
        nranksPerMode: *const i64,
        nranks: u32,
        ranks: *const i32,
        dataType: cudaDataType_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Frees all resources related to the provided distributed tensor descriptor.
    ///
    /// This function deallocates all memory and resources associated with a cutensorMp tensor descriptor that was previously created by cutensorMpCreateTensorDescriptor. After calling this function, the descriptor becomes invalid and should not be used in subsequent cutensorMp operations.
    ///
    /// Remark
    ///
    /// blocking, no reentrant, and thread-safe.
    pub fn cutensorMpDestroyTensorDescriptor(
        desc: cutensorMpTensorDescriptor_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Creates an operation descriptor that encodes a distributed tensor contraction.
    ///
    /// This function creates an operation descriptor for distributed tensor contractions of the form $D = \alpha \mathcal{A} \mathcal{B} + \beta \mathcal{C}$, where the tensors A, B, C, and D are distributed across multiple processes as specified by their respective tensor descriptors.
    ///
    /// The distributed contraction leverages both intra-process cuTENSOR operations and inter-process communication to efficiently compute tensor contractions that exceed the memory capacity or computational resources of a single GPU. The operation automatically handles data redistribution, local contractions, and result aggregation across the participating processes.
    ///
    /// The user is responsible for calling cutensorMpDestroyOperationDescriptor to free the resources associated with the descriptor once it is no longer needed.
    ///
    /// Remark
    ///
    /// non-blocking, no reentrant, and thread-safe.
    ///
    /// # Parameters
    ///
    /// - `handle`: Opaque handle holding cutensorMp’s library context.
    /// - `desc`: Pointer to the operation descriptor that will be created and filled with information encoding the distributed contraction operation.
    /// - `descA`: Distributed tensor descriptor for input tensor A.
    /// - `modesA`: Modes of the input tensor A.
    /// - `opA`: Unary operator that will be applied to each element of A before it is further processed. The original data of this tensor remains unchanged.
    /// - `descB`: Distributed tensor descriptor for input tensor B.
    /// - `modesB`: Modes of the input tensor B.
    /// - `opB`: Unary operator that will be applied to each element of B before it is further processed. The original data of this tensor remains unchanged.
    /// - `descC`: Distributed tensor descriptor for input tensor C.
    /// - `modesC`: Modes of the input tensor C.
    /// - `opC`: Unary operator that will be applied to each element of C before it is further processed. The original data of this tensor remains unchanged.
    /// - `descD`: Distributed tensor descriptor for output tensor D (currently must be identical to descC).
    /// - `modesD`: Modes of the output tensor D.
    /// - `descCompute`: Compute descriptor that determines the precision for the operation.
    ///
    /// # Return value
    ///
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INVALID_VALUE`]: if some input data is invalid (this typically indicates a user error).
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_INITIALIZED`]: if the handle is not initialized.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_SUPPORTED`]: if the combination of tensor configurations is not supported.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_SUCCESS`]: The operation completed successfully.
    pub fn cutensorMpCreateContraction(
        handle: cutensorMpHandle_t,
        desc: *mut cutensorMpOperationDescriptor_t,
        descA: cutensorMpTensorDescriptor_t,
        modesA: *const i32,
        opA: cutensorOperator_t,
        descB: cutensorMpTensorDescriptor_t,
        modesB: *const i32,
        opB: cutensorOperator_t,
        descC: cutensorMpTensorDescriptor_t,
        modesC: *const i32,
        opC: cutensorOperator_t,
        descD: cutensorMpTensorDescriptor_t,
        modesD: *const i32,
        descCompute: cutensorComputeDescriptor_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Frees all resources related to the provided distributed contraction descriptor.
    ///
    /// This function deallocates all memory and resources associated with a cutensorMp operation descriptor that was previously created by cutensorMpCreateContraction. After calling this function, the descriptor becomes invalid and should not be used in subsequent cutensorMp operations.
    ///
    /// Remark
    ///
    /// blocking, no reentrant, and thread-safe.
    pub fn cutensorMpDestroyOperationDescriptor(
        desc: cutensorMpOperationDescriptor_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Creates a plan preference object for controlling distributed tensor operation planning.
    ///
    /// This function creates a preference object that allows users to control various aspects of the execution plan for distributed tensor operations. The preferences include algorithm selection, workspace size limits, and JIT compilation options that affect both the underlying cuTENSOR operations and the distributed communication patterns.
    ///
    /// The plan preference provides fine-grained control over:
    ///
    /// * Local cuTENSOR algorithm selection and JIT mode
    /// * Distributed algorithm strategy (non-packing, packing with permutation, or packing with P2P)
    /// * Workspace size limits for both device and host memory
    /// * cuTENSOR workspace preferences
    ///
    /// The user is responsible for calling cutensorMpDestroyPlanPreference to free the resources associated with the preference object.
    ///
    /// Remark
    ///
    /// non-blocking, no reentrant, and thread-safe.
    ///
    /// # Parameters
    ///
    /// - `handle`: Opaque handle holding cutensorMp’s library context.
    /// - `pref`: Pointer to the plan preference object that will be created.
    /// - `cutensormp_algo`: Algorithm selection for distributed communication patterns.
    /// - `cutensormp_workspace_size_device`: Maximum device workspace size for cutensorMp operations (bytes), minimum 2GB is required.
    /// - `cutensormp_workspace_size_host`: Maximum host workspace size for cutensorMp operations (bytes).
    ///
    /// # Return value
    ///
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INVALID_VALUE`]: if some input data is invalid (this typically indicates a user error).
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_INITIALIZED`]: if the handle is not initialized.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_SUCCESS`]: The operation completed successfully.
    pub fn cutensorMpCreatePlanPreference(
        handle: cutensorMpHandle_t,
        pref: *mut cutensorMpPlanPreference_t,
        cutensormp_algo: cutensorMpAlgo_t,
        cutensormp_workspace_size_device: u64,
        cutensormp_workspace_size_host: u64,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Frees all resources related to the provided plan preference object.
    ///
    /// This function deallocates all memory and resources associated with a cutensorMp plan preference object that was previously created by cutensorMpCreatePlanPreference. After calling this function, the preference object becomes invalid and should not be used in subsequent cutensorMp operations.
    ///
    /// Remark
    ///
    /// blocking, no reentrant, and thread-safe.
    pub fn cutensorMpDestroyPlanPreference(
        pref: cutensorMpPlanPreference_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Creates an execution plan for distributed tensor contractions.
    ///
    /// This function creates an optimized execution plan for the distributed tensor contraction encoded by the operation descriptor. The plan selects the most appropriate algorithms and communication strategies based on the tensor distributions, available resources, and user preferences.
    ///
    /// The planning process analyzes the distributed tensor layout, communication requirements, and computational resources to determine an efficient execution strategy. This may involve data redistribution, local contractions, and result aggregation phases that minimize communication overhead while maximizing computational efficiency.
    ///
    /// The user is responsible for calling cutensorMpDestroyPlan to free the resources associated with the plan once it is no longer needed.
    ///
    /// Remark
    ///
    /// calls asynchronous functions, no reentrant, and thread-safe.
    ///
    /// # Parameters
    ///
    /// - `handle`: Opaque handle holding cutensorMp’s library context.
    /// - `plan`: Pointer to the execution plan object that will be created.
    /// - `desc`: Operation descriptor encoding the distributed contraction (created by cutensorMpCreateContraction).
    /// - `pref`: Plan preference object specifying algorithm and workspace preferences (may be NULL for default preferences).
    ///
    /// # Return value
    ///
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INVALID_VALUE`]: if some input data is invalid (this typically indicates a user error).
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_INITIALIZED`]: if the handle is not initialized.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_SUPPORTED`]: if no viable execution plan could be found.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_SUCCESS`]: The operation completed successfully.
    pub fn cutensorMpCreatePlan(
        handle: cutensorMpHandle_t,
        plan: *mut cutensorMpPlan_t,
        desc: cutensorMpOperationDescriptor_t,
        pref: cutensorMpPlanPreference_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Frees all resources related to the provided distributed contraction plan.
    ///
    /// This function deallocates all memory and resources associated with a cutensorMp execution plan that was previously created by cutensorMpCreatePlan. After calling this function, the plan becomes invalid and should not be used in subsequent cutensorMp operations.
    ///
    /// Remark
    ///
    /// blocking, no reentrant, and thread-safe.
    pub fn cutensorMpDestroyPlan(plan: cutensorMpPlan_t) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Retrieves information about an already-created plan (see [`cutensorPlanAttribute_t`]).
    ///
    /// # Parameters
    ///
    /// - `plan`: Denotes an already-created plan (e.g., via cutensorMpCreatePlan).
    /// - `buf`: On successful exit: Holds the information of the requested attribute.
    /// - `sizeInBytes`: size of `buf` in bytes.
    ///
    /// # Return value
    ///
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INVALID_VALUE`]: if some input data is invalid (this typically indicates an user error).
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_SUCCESS`]: The operation completed successfully.
    pub fn cutensorMpPlanGetAttribute(
        handle: cutensorMpHandle_t,
        plan: cutensorMpPlan_t,
        attribute: cutensorMpPlanAttribute_t,
        buf: *mut ::core::ffi::c_void,
        sizeInBytes: size_t,
    ) -> cutensorStatus_t;
}
unsafe extern "C" {
    /// Performs a distributed tensor contraction across multiple processes.
    ///
    /// This function executes the distributed tensor contraction $D = \alpha \mathcal{A} \mathcal{B} + \beta \mathcal{C}$ according to the execution plan created by cutensorMpCreatePlan. The operation coordinates computation and communication across multiple processes and GPUs to efficiently perform tensor contractions that exceed the capacity of a single device.
    ///
    /// The execution involves several phases:
    ///
    /// 1. Data redistribution to align tensor blocks for efficient computation
    /// 2. Local tensor contractions using cuTENSOR on each participating device
    /// 3. Communication and aggregation of partial results across processes
    /// 4. Final result assembly in the distributed output tensor
    ///
    /// All participating processes in the MPI communicator must call this function with consistent parameters. The input and output tensors must be distributed according to their respective tensor descriptors, with each process providing its local portion of the data.
    ///
    /// Remark
    ///
    /// calls asynchronous functions, no reentrant, and thread-safe.
    ///
    /// # Parameters
    ///
    /// - `handle`: Opaque handle holding cutensorMp’s library context.
    /// - `plan`: Execution plan for the distributed contraction (created by cutensorMpCreatePlan).
    /// - `alpha`: Scaling factor for the A\*B product. Pointer to host memory with data type determined by the compute descriptor. The data type follows that of cuTENSOR (i.e., the data type of the scalar is deptermined by the data type of `C`:[`cudaDataType_t::CUDA_R_16F`] and [`cudaDataType_t::CUDA_R_16BF`] use [`cudaDataType_t::CUDA_R_32F`]scalars, all data types of the scalar are identical to the type of `C`).
    /// - `A`: Pointer to the local portion of distributed tensor A in GPU memory.
    /// - `B`: Pointer to the local portion of distributed tensor B in GPU memory.
    /// - `beta`: Scaling factor for tensor C. Pointer to host memory with data type determined by the compute descriptor. The data type follows that of cuTENSOR (i.e., the data type of the scalar is deptermined by the data type of `C`:[`cudaDataType_t::CUDA_R_16F`] and [`cudaDataType_t::CUDA_R_16BF`] use [`cudaDataType_t::CUDA_R_32F`]scalars, all data types of the scalar are identical to the type of `C`).
    /// - `C`: Pointer to the local portion of distributed tensor C in GPU memory.
    /// - `D`: Pointer to the local portion of distributed tensor D in GPU memory (may be identical to C).
    /// - `device_workspace`: Pointer to device workspace memory (size determined by cutensorMpPlanGetAttribute with [`cutensorMpPlanAttribute_t::CUTENSORMP_PLAN_REQUIRED_WORKSPACE_DEVICE`]).
    /// - `host_workspace`: Pointer to host workspace memory (size determined by cutensorMpPlanGetAttribute with [`cutensorMpPlanAttribute_t::CUTENSORMP_PLAN_REQUIRED_WORKSPACE_HOST`]).
    ///
    /// # Return value
    ///
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_ARCH_MISMATCH`]: if the plan was created for a different device architecture.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_CUDA_ERROR`]: if a CUDA error occurred during execution.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INSUFFICIENT_WORKSPACE`]: if the provided workspace is insufficient.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_INVALID_VALUE`]: if some input data is invalid (this typically indicates a user error).
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_INITIALIZED`]: if the handle is not initialized.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_NOT_SUPPORTED`]: if the operation is not supported with the given configuration.
    /// - [`cutensorStatus_t::CUTENSOR_STATUS_SUCCESS`]: The operation completed successfully.
    pub fn cutensorMpContract(
        handle: cutensorMpHandle_t,
        plan: cutensorMpPlan_t,
        alpha: *const ::core::ffi::c_void,
        A: *const ::core::ffi::c_void,
        B: *const ::core::ffi::c_void,
        beta: *const ::core::ffi::c_void,
        C: *const ::core::ffi::c_void,
        D: *mut ::core::ffi::c_void,
        device_workspace: *mut ::core::ffi::c_void,
        host_workspace: *mut ::core::ffi::c_void,
    ) -> cutensorStatus_t;
}