/*
 * Copyright 2026 NVIDIA Corporation.  All rights reserved.
 *
 * NOTICE TO LICENSEE:
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 *
 * These Licensed Deliverables contained herein is PROPRIETARY and
 * CONFIDENTIAL to NVIDIA and is being provided under the terms and
 * conditions of a form of NVIDIA software license agreement by and
 * between NVIDIA and Licensee ("License Agreement") or electronically
 * accepted by Licensee.  Notwithstanding any terms or conditions to
 * the contrary in the License Agreement, reproduction or disclosure
 * of the Licensed Deliverables to any third party without the express
 * written consent of NVIDIA is prohibited.
 *
 * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
 * LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
 * SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE.  IT IS
 * PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
 * NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
 * DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
 * NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
 * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
 * LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
 * SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
 * DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
 * WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
 * ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
 * OF THESE LICENSED DELIVERABLES.
 *
 * U.S. Government End Users.  These Licensed Deliverables are a
 * "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
 * 1995), consisting of "commercial computer software" and "commercial
 * computer software documentation" as such terms are used in 48
 * C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
 * only as a commercial end item.  Consistent with 48 C.F.R.12.212 and
 * 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
 * U.S. Government End Users acquire the Licensed Deliverables with
 * only those rights set forth herein.
 *
 * Any use of the Licensed Deliverables in individual and commercial
 * software must include, in the user documentation and internal
 * comments to the code, the above Disclaimer and U.S. Government End
 * Users Notice.
 */

#pragma once

#include <cublas_v2.h>
#include <inttypes.h>
#include <nccl.h>
#include <stdio.h>

#define CUBLASMP_VER_MAJOR 0
#define CUBLASMP_VER_MINOR 8
#define CUBLASMP_VER_PATCH 1
#define CUBLASMP_VERSION (CUBLASMP_VER_MAJOR * 1000 + CUBLASMP_VER_MINOR * 100 + CUBLASMP_VER_PATCH)

#ifdef __cplusplus
extern "C"
{
#endif

/**
 * \brief Status type returned by all cuBLASMp library functions.
 *
 * All cuBLASMp library functions return their status.
 */
typedef enum
{
    /** The operation completed successfully. */
    CUBLASMP_STATUS_SUCCESS = 0,
    /** The cuBLASMp library was not initialized. */
    CUBLASMP_STATUS_NOT_INITIALIZED = 1,
    /** Resource allocation failed inside the cuBLASMp library. */
    CUBLASMP_STATUS_ALLOCATION_FAILED = 2,
    /** An unsupported value or parameter was passed to the function. */
    CUBLASMP_STATUS_INVALID_VALUE = 3,
    /** The function requires a feature absent from the device architecture. */
    CUBLASMP_STATUS_ARCHITECTURE_MISMATCH = 4,
    /** The GPU program failed to execute. */
    CUBLASMP_STATUS_EXECUTION_FAILED = 5,
    /** An internal cuBLASMp operation failed. */
    CUBLASMP_STATUS_INTERNAL_ERROR = 6,
    /** The functionality requested is not supported. */
    CUBLASMP_STATUS_NOT_SUPPORTED = 7,
} cublasMpStatus_t;

/**
 * \brief Describes the ordering of the process grid.
 */
typedef enum
{
    /** The grid of processes will be accessed in column-major ordering. */
    CUBLASMP_GRID_LAYOUT_COL_MAJOR = 0,
    /** The grid of processes will be accessed in row-major ordering. */
    CUBLASMP_GRID_LAYOUT_ROW_MAJOR = 1
} cublasMpGridLayout_t;

/**
 * \brief Attributes of \ref cublasMpMatmulDescriptor_t that can be set and queried.
 *
 * These attributes are set using \ref cublasMpMatmulDescriptorSetAttribute and queried
 * using \ref cublasMpMatmulDescriptorGetAttribute.
 */
typedef enum
{
    /** Operation to perform on matrix A. Type: cublasOperation_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_TRANSA = 0,
    /** Operation to perform on matrix B. Type: cublasOperation_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_TRANSB = 1,
    /** Compute type of the matrix multiplication. Type: cublasComputeType_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_COMPUTE_TYPE = 2,
    /** Hint for the algorithm type to use. If not supported, cuBLASMp falls back to the
     *  default algorithm. Type: cublasMpMatmulAlgoType_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_ALGO_TYPE = 3,
    /** Number of SMs to use for communication. Type: int32_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_COMMUNICATION_SM_COUNT = 4,
    /** Epilogue operation after matrix multiplication. Type: cublasMpMatmulEpilogue_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE = 5,
    /** Device pointer to bias or bias gradient vector. Type: void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_BIAS_POINTER = 6,
    /** Stride (in elements) between bias values in batched operations. Must be 0
     *  (batched matmul is not currently supported). Type: int64_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_BIAS_BATCH_STRIDE = 7,
    /** Data type of the bias vector. Type: cudaDataType_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_BIAS_DATA_TYPE = 8,
    /** Device pointer for epilogue auxiliary buffer (e.g., ReLU bit-mask in forward pass,
     *  GELU input matrix). Requires setting EPILOGUE_AUX_LD. Type: void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_POINTER = 9,
    /** Leading dimension for epilogue auxiliary buffer. For ReLU bit-mask: in elements (bits),
     *  must be divisible by 128. For GELU input: in elements, must be divisible by 8.
     *  Type: int64_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_LD = 10,
    /** Stride (in bytes) between auxiliary buffers in batched operations. Must be 0
     *  (batched matmul is not currently supported). Type: int64_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_BATCH_STRIDE = 11,
    /** Data type stored in the epilogue auxiliary pointer. If unset, defaults to the output
     *  matrix element data type (with exceptions: ReLU uses a bit-mask). Type: cudaDataType_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_DATA_TYPE = 12,
    /** Device pointer to the scaling factor to convert results from compute type to storage
     *  data range in the auxiliary matrix. If not specified, assumed to be 1. Type: void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_SCALE_POINTER = 13,
    /** Device pointer to receive the max absolute value in the auxiliary buffer on completion.
     *  If not specified, the max absolute value is not computed. Type: void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_AMAX_POINTER = 14,
    /** Scaling mode for the auxiliary matrix. Default: CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32.
     *  Type: cublasMpMatmulMatrixScale_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_SCALE_MODE = 15,
    /** Device pointer to scale factor that converts data in matrix A to the compute data type
     *  range. Supported only when compute type is CUBLAS_COMPUTE_32F and the output matrix
     *  type is not complex. If not specified, assumed to be 1. Type: const void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_POINTER = 16,
    /** Scaling mode for matrix A. Default: CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32.
     *  Type: cublasMpMatmulMatrixScale_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_MODE = 17,
    /** Device pointer to scale factor for matrix B (equivalent to A_SCALE_POINTER for B).
     *  Type: const void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_POINTER = 18,
    /** Scaling mode for matrix B. Default: CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32.
     *  Type: cublasMpMatmulMatrixScale_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_MODE = 19,
    /** Device pointer to scale factor for matrix C (equivalent to A_SCALE_POINTER for C).
     *  Type: const void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_C_SCALE_POINTER = 20,
    /** Scaling mode for matrix C. Default: CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32.
     *  Type: cublasMpMatmulMatrixScale_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_C_SCALE_MODE = 21,
    /** Device pointer to scale factor for matrix D (equivalent to A_SCALE_POINTER for D).
     *  Type: const void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_D_SCALE_POINTER = 22,
    /** Scaling mode for matrix D. Default: CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32.
     *  Type: cublasMpMatmulMatrixScale_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_D_SCALE_MODE = 23,
    /** Device pointer to receive the max absolute value in the output matrix D on completion.
     *  If not specified, the max absolute value is not computed. Type: void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_AMAX_D_POINTER = 24,
    /** Device pointer to scale factors for converting output matrix D data. Default: NULL.
     *  Type: void*. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_D_OUT_SCALE_POINTER = 25,
    /** Scaling mode for the output matrix D scaling. Default: 0.
     *  Type: cublasMpMatmulMatrixScale_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_D_OUT_SCALE_MODE = 26,
    /** Data type used for inter-GPU communication in distributed matmul operations.
     *  When specified, the communication happens in this type, which can reduce communication
     *  bandwidth at the cost of precision. Supported by all GEMM+ReduceScatter variants.
     *  Must satisfy: d_type <= communication_type <= scale_type (by element size).
     *  Default: scale type derived from compute type. Type: cudaDataType_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_COMMUNICATION_TYPE = 27,
    /** Controls how epilogue output buffers (bias gradients) are handled after matrix
     *  multiplication. When set to FULL, performs AllReduce on the bias gradient output buffer.
     *  Applicable to bias gradient epilogues with GEMM+ReduceScatter algorithm.
     *  Default: CUBLASMP_RESULT_SCHEME_DEFAULT. Type: cublasMpResultScheme_t. */
    CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_BIAS_RESULT_SCHEME = 28
} cublasMpMatmulDescriptorAttribute_t;

/**
 * \brief Matrix-matrix multiplication algorithm types.
 *
 * This is treated as a hint and it is not guaranteed that cuBLASMp will use the
 * requested implementation.
 */
typedef enum
{
    /** Default algorithm. */
    CUBLASMP_MATMUL_ALGO_TYPE_DEFAULT = 0,
    /** Use split matmul with P2P communication. */
    CUBLASMP_MATMUL_ALGO_TYPE_SPLIT_P2P = 1,
    /** Use split matmul with multicast communication. Only valid for GEMM+ReduceScatter
     *  and GEMM+AllReduce on Compute Capability 9.0 (Hopper) and newer GPUs connected
     *  with an NVSwitch. */
    CUBLASMP_MATMUL_ALGO_TYPE_SPLIT_MULTICAST = 2,
    /** Use atomic matmul with P2P communication. */
    CUBLASMP_MATMUL_ALGO_TYPE_ATOMIC_P2P = 3,
    /** Use atomic matmul with multicast communication. Only valid for GEMM+ReduceScatter
     *  on Compute Capability 9.0 (Hopper) and newer GPUs connected with an NVSwitch.
     *  [DEPRECATED] */
    CUBLASMP_MATMUL_ALGO_TYPE_ATOMIC_MULTICAST = 4,
    /** Perform computation and communication sequentially without overlap. This algorithm
     *  may perform better on non-NVLINK multi-node systems. */
    CUBLASMP_MATMUL_ALGO_TYPE_NO_OVERLAP = 5
} cublasMpMatmulAlgoType_t;

/**
 * \brief Epilogue operation to perform after matrix multiplication.
 *
 * Epilogues enable fusion of additional operations such as activation functions,
 * bias addition, and their derivatives for training, providing better performance
 * by avoiding separate kernel launches.
 */
typedef enum
{
    /** No special postprocessing. */
    CUBLASMP_MATMUL_EPILOGUE_DEFAULT = 0,
    /** Performs AllReduce communication operation after matrix multiplication. */
    CUBLASMP_MATMUL_EPILOGUE_ALLREDUCE = 1,
    /** Apply ReLU point-wise transform: x := max(x, 0). */
    CUBLASMP_MATMUL_EPILOGUE_RELU = 2,
    /** Apply ReLU transform and produce auxiliary output (ReLU bit-mask). */
    CUBLASMP_MATMUL_EPILOGUE_RELU_AUX = (CUBLASMP_MATMUL_EPILOGUE_RELU | 128),
    /** Apply broadcast bias from the bias vector. Bias vector length must match
     *  matrix D rows. */
    CUBLASMP_MATMUL_EPILOGUE_BIAS = 4,
    /** Apply bias and then ReLU transform. */
    CUBLASMP_MATMUL_EPILOGUE_RELU_BIAS = (CUBLASMP_MATMUL_EPILOGUE_RELU | CUBLASMP_MATMUL_EPILOGUE_BIAS),
    /** Apply bias and then ReLU transform with auxiliary output. */
    CUBLASMP_MATMUL_EPILOGUE_RELU_AUX_BIAS = (CUBLASMP_MATMUL_EPILOGUE_RELU_AUX | CUBLASMP_MATMUL_EPILOGUE_BIAS),
    /** ReLU gradient. Currently not supported. */
    CUBLASMP_MATMUL_EPILOGUE_DRELU = (8 | 128),
    /** ReLU gradient with bias gradient. Currently not supported. */
    CUBLASMP_MATMUL_EPILOGUE_DRELU_BGRAD = (CUBLASMP_MATMUL_EPILOGUE_DRELU | 16),
    /** Apply GELU point-wise transform: x := GELU(x). */
    CUBLASMP_MATMUL_EPILOGUE_GELU = 32,
    /** Apply GELU transform and output GELU input as a separate matrix (useful for training). */
    CUBLASMP_MATMUL_EPILOGUE_GELU_AUX = (CUBLASMP_MATMUL_EPILOGUE_GELU | 128),
    /** Apply bias and then GELU transform. */
    CUBLASMP_MATMUL_EPILOGUE_GELU_BIAS = (CUBLASMP_MATMUL_EPILOGUE_GELU | CUBLASMP_MATMUL_EPILOGUE_BIAS),
    /** Apply bias and then GELU transform with auxiliary output (useful for training). */
    CUBLASMP_MATMUL_EPILOGUE_GELU_AUX_BIAS = (CUBLASMP_MATMUL_EPILOGUE_GELU_AUX | CUBLASMP_MATMUL_EPILOGUE_BIAS),
    /** Apply GELU gradient to matmul output. Requires auxiliary input. */
    CUBLASMP_MATMUL_EPILOGUE_DGELU = (64 | 128),
    /** Apply GELU and bias gradients independently to matmul output. Requires auxiliary input. */
    CUBLASMP_MATMUL_EPILOGUE_DGELU_BGRAD = (CUBLASMP_MATMUL_EPILOGUE_DGELU | 16),
    /** Apply bias gradient over the GEMM's "k" dimension for input matrix A. Bias size
     *  corresponds to the number of rows of matrix D. */
    CUBLASMP_MATMUL_EPILOGUE_BGRADA = 256,
    /** Apply bias gradient over the GEMM's "k" dimension for input matrix B. Bias size
     *  corresponds to the number of columns of matrix D. */
    CUBLASMP_MATMUL_EPILOGUE_BGRADB = 512
} cublasMpMatmulEpilogue_t;

/**
 * \brief Scaling mode that defines how matrix scaling factor pointers are interpreted.
 */
typedef enum
{
    /** Scaling factors are single-precision scalars applied to whole tensors.
     *  Supported with Compute Capability 8.9 and above. */
    CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32 = 0,
    /** Scaling factors contain a dedicated 8-bit CUDA_R_8F_UE4M3 value for each 16-element
     *  block in the innermost dimension. Supported with Compute Capability 10.0 and above. */
    CUBLASMP_MATMUL_MATRIX_SCALE_VEC16_UE4M3 = 1,
    /** Scaling factors contain a dedicated 8-bit CUDA_R_8F_UE8M0 value for each 32-element
     *  block in the innermost dimension. Supported with Compute Capability 10.0 and above. */
    CUBLASMP_MATMUL_MATRIX_SCALE_VEC32_UE8M0 = 2,
    /** Scaling factors are vectors of CUDA_R_32F values. Only applicable to matrices A and B,
     *  where the vectors have M and N elements respectively. Supported with Compute
     *  Capability 9.0. */
    CUBLASMP_MATMUL_MATRIX_SCALE_OUTER_VEC_FP32 = 3,
    /** Scaling factors contain a dedicated CUDA_R_32F scaling factor for each 128-element
     *  block in the innermost dimension. Supported with Compute Capability 9.0. */
    CUBLASMP_MATMUL_MATRIX_SCALE_VEC128_FP32 = 4,
    /** Scaling factors contain a dedicated CUDA_R_32F scaling factor for each 128x128-element
     *  block. Supported with Compute Capability 9.0. */
    CUBLASMP_MATMUL_MATRIX_SCALE_BLK128x128_FP32 = 5
} cublasMpMatmulMatrixScale_t;

/**
 * \brief Emulation strategy for floating point emulation algorithms.
 *
 * Used in \ref cublasMpSetEmulationStrategy to choose how to leverage floating
 * point emulation algorithms. The emulation strategy must be set after handle
 * creation using \ref cublasMpCreate and before performing computations.
 */
typedef enum
{
    /** Default emulation strategy (equivalent to PERFORMANT). */
    CUBLASMP_EMULATION_STRATEGY_DEFAULT = 0,
    /** Utilizes emulation whenever it provides a performance benefit. */
    CUBLASMP_EMULATION_STRATEGY_PERFORMANT = 1,
    /** Utilizes emulation whenever possible. */
    CUBLASMP_EMULATION_STRATEGY_EAGER = 2
} cublasMpEmulationStrategy_t;

/**
 * \brief Controls how epilogue output buffers (such as bias gradients) are handled
 * after matrix multiplication in distributed algorithms.
 */
typedef enum
{
    /** Use default behavior (equivalent to PARTIAL). */
    CUBLASMP_RESULT_SCHEME_DEFAULT = 0,
    /** Leave epilogue output as partial sums. Each rank retains its local partial result.
     *  The user is responsible for performing any required reduction externally. */
    CUBLASMP_RESULT_SCHEME_PARTIAL = 1,
    /** Perform AllReduce on the epilogue output buffer to produce the full (reduced) result
     *  on all ranks. Useful for GEMM+ReduceScatter where each rank computes partial sums
     *  of the bias gradient. */
    CUBLASMP_RESULT_SCHEME_FULL = 2,
} cublasMpResultScheme_t;

/**
 * \brief Opaque structure holding the cuBLASMp library context.
 *
 * The cublasMpHandle_t holds the cuBLASMp library context (device properties,
 * system information, etc.). The handle must be initialized and destroyed using
 * \ref cublasMpCreate and \ref cublasMpDestroy respectively.
 */
struct cublasMpHandle;
typedef struct cublasMpHandle* cublasMpHandle_t;

/**
 * \brief Opaque structure holding information about the process grid.
 *
 * The cublasMpGrid_t holds information about the grid dimensions and stores the
 * communicator associated with the grid of processes. It must be initialized and
 * destroyed using \ref cublasMpGridCreate and \ref cublasMpGridDestroy respectively.
 */
struct cublasMpGrid;
typedef struct cublasMpGrid* cublasMpGrid_t;

/**
 * \brief Opaque structure describing the shape and characteristics of a distributed matrix.
 *
 * The cublasMpMatrixDescriptor_t captures the shape and characteristics of a distributed
 * matrix. It must be initialized and destroyed using \ref cublasMpMatrixDescriptorCreate
 * and \ref cublasMpMatrixDescriptorDestroy respectively.
 */
struct cublasMpMatrixDescriptor;
typedef struct cublasMpMatrixDescriptor* cublasMpMatrixDescriptor_t;

/**
 * \brief Opaque structure describing properties of a distributed matrix-matrix multiplication.
 *
 * The cublasMpMatmulDescriptor_t captures the properties of a distributed matrix-matrix
 * multiplication performed using \ref cublasMpMatmul. It must be initialized and destroyed
 * using \ref cublasMpMatmulDescriptorCreate and \ref cublasMpMatmulDescriptorDestroy
 * respectively.
 */
struct cublasMpMatmulDescriptor;
typedef struct cublasMpMatmulDescriptor* cublasMpMatmulDescriptor_t;

/**
 * \brief Initializes the cuBLASMp library handle.
 *
 * This function initializes the cuBLASMp library handle (cublasMpHandle_t) which holds
 * the cuBLASMp library context. It allocates light hardware resources on the host, and
 * must be called prior to making any other cuBLASMp library calls.
 *
 * Calling any cuBLASMp function which uses cublasMpHandle_t without a previous call
 * of cublasMpCreate() will return an error.
 *
 * The cuBLASMp library context is tied to the current CUDA device and the given CUDA stream.
 * Sharing a device with multiple processes may result in undefined behavior.
 *
 * \param[out] handle  cuBLASMp library handle.
 * \param[in]  stream  Stream that will be assigned to the handle.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpCreate(cublasMpHandle_t* handle, cudaStream_t stream);

/**
 * \brief Destroys the cuBLASMp library handle.
 *
 * This function destroys the cuBLASMp library handle (cublasMpHandle_t) which holds
 * the cuBLASMp library context. The cuBLASMp library context is tied to the CUDA
 * device that was set when calling cublasMpCreate(). Only one handle per process
 * and per GPU is supported.
 *
 * \param[in,out] handle  cuBLASMp library handle to destroy.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpDestroy(cublasMpHandle_t handle);

/**
 * \brief Sets the CUDA stream for cuBLASMp computations.
 *
 * This function sets the CUDA stream to be used in the computations.
 *
 * \param[in] handle  cuBLASMp library handle.
 * \param[in] stream  CUDA stream to set.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetStream(cublasMpHandle_t handle, cudaStream_t stream);

/**
 * \brief Gets the current CUDA stream used by the cuBLASMp library.
 *
 * This function returns the current CUDA stream that is being used in the computations.
 *
 * \param[in]  handle  cuBLASMp library handle.
 * \param[out] stream  Pointer to receive the current CUDA stream.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetStream(cublasMpHandle_t handle, cudaStream_t* stream);

/**
 * \brief Sets the floating point emulation strategy.
 *
 * This function allows you to select how the library should make use of floating
 * point emulation. See \ref cublasMpEmulationStrategy_t for available strategies.
 * The emulation strategy must be set after handle creation and before performing
 * computations.
 *
 * \param[in] handle             cuBLASMp library handle.
 * \param[in] emulationStrategy  Emulation strategy to use.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetEmulationStrategy(cublasMpHandle_t handle, cublasMpEmulationStrategy_t emulationStrategy);

/**
 * \brief Gets the current floating point emulation strategy.
 *
 * This function retrieves the current emulation strategy from the cuBLASMp
 * library handle.
 *
 * \param[in]  handle             cuBLASMp library handle.
 * \param[out] emulationStrategy  Pointer to receive the current emulation strategy.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetEmulationStrategy(cublasMpHandle_t handle, cublasMpEmulationStrategy_t* emulationStrategy);

/**
 * \brief Returns the version number of the cuBLASMp library.
 *
 * The version value is computed as
 * CUBLASMP_VER_MAJOR * 1000 + CUBLASMP_VER_MINOR * 100 + CUBLASMP_VER_PATCH.
 *
 * \param[out] version  cuBLASMp library version.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetVersion(int* version);

/**
 * \brief Creates and initializes a process grid.
 *
 * This function initializes the cublasMpGrid_t opaque data structure. It maps the
 * given resources (communicator, grid dimensions and grid layout) to a grid object.
 *
 * All the processes defined to be in this grid must enter this function.
 *
 * \param[in]     nprow   Number of row processes in the grid.
 * \param[in]     npcol   Number of column processes in the grid.
 * \param[in]     layout  Grid layout (cublasMpGridLayout_t).
 * \param[in]     comm    NCCL communicator associated with the grid.
 * \param[in,out] grid    Pointer to a grid object to initialize.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGridCreate(
    int64_t nprow,
    int64_t npcol,
    cublasMpGridLayout_t layout,
    ncclComm_t comm,
    cublasMpGrid_t* grid);

/**
 * \brief Destroys a process grid.
 *
 * This function destroys the given grid object. All the processes defined to
 * be in this grid must enter this function.
 *
 * \param[in,out] grid  Grid object to destroy.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGridDestroy(cublasMpGrid_t grid);

/**
 * \brief Registers a device memory buffer with NCCL symmetric memory for a process grid.
 *
 * Registered buffers can be accessed by all processes in the grid for efficient
 * communication-computation overlap in cuBLASMp operations. Buffers allocated with
 * ncclMemAlloc or other CUDA VMM allocation methods should be registered with the grid
 * associated with the output matrix (e.g., matrix D in matmul operations) before calling
 * cuBLASMp computation functions.
 *
 * All the processes defined to be in this grid must enter this function.
 *
 * \param[in] grid  Process grid to register the buffer with.
 * \param[in] ptr   Pointer to device memory buffer to register.
 * \param[in] size  Size of the buffer in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpBufferRegister(cublasMpGrid_t grid, void* ptr, size_t size);

/**
 * \brief Deregisters a previously registered device memory buffer from NCCL symmetric memory.
 *
 * This should be called before freeing the buffer with ncclMemFree. For buffers allocated
 * with \ref cublasMpMalloc, use \ref cublasMpFree instead, which combines deregistration
 * and freeing.
 *
 * All the processes defined to be in this grid must enter this function.
 *
 * \param[in] grid  Process grid the buffer was registered with.
 * \param[in] ptr   Pointer to device memory buffer to deregister.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpBufferDeregister(cublasMpGrid_t grid, void* ptr);

/**
 * \brief Creates and initializes a matrix descriptor.
 *
 * This function creates and initializes a new cublasMpMatrixDescriptor_t object that
 * describes a distributed matrix using 2D block-cyclic distribution.
 *
 * \param[in]  m     Number of rows in the global matrix.
 * \param[in]  n     Number of columns in the global matrix.
 * \param[in]  mb    Blocking factor used to distribute the rows of the global matrix.
 * \param[in]  nb    Blocking factor used to distribute the columns of the global matrix.
 * \param[in]  rsrc  Row rank of the process who owns the first row block of the global matrix.
 * \param[in]  csrc  Column rank of the process who owns the first column block of the global matrix.
 * \param[in]  lld   Leading dimension of the local matrix.
 * \param[in]  type  Data type of the matrix (e.g., CUDA_R_32F, CUDA_R_64F, CUDA_R_16F, CUDA_R_16BF,
 *                   CUDA_R_8F_E4M3, CUDA_R_8F_E5M2, CUDA_R_4F_E2M1, CUDA_R_8I, CUDA_R_32I,
 *                   CUDA_C_32F, CUDA_C_64F).
 * \param[in]  grid  Grid object associated with the matrix descriptor.
 * \param[out] desc  Matrix descriptor object initialized by this function.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatrixDescriptorCreate(
    int64_t m,
    int64_t n,
    int64_t mb,
    int64_t nb,
    int64_t rsrc,
    int64_t csrc,
    int64_t lld,
    cudaDataType_t type,
    cublasMpGrid_t grid,
    cublasMpMatrixDescriptor_t* desc);

/**
 * \brief Destroys a matrix descriptor.
 *
 * This function destroys a cublasMpMatrixDescriptor_t object.
 *
 * \param[in,out] desc  Matrix descriptor object to destroy.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatrixDescriptorDestroy(cublasMpMatrixDescriptor_t desc);

/**
 * \brief Initializes the values of an existing matrix descriptor.
 *
 * This function initializes the values of a cublasMpMatrixDescriptor_t object.
 * Unlike \ref cublasMpMatrixDescriptorCreate, this function does not allocate
 * additional memory.
 *
 * \param[in]     m     Number of rows in the global matrix.
 * \param[in]     n     Number of columns in the global matrix.
 * \param[in]     mb    Blocking factor used to distribute the rows of the global matrix.
 * \param[in]     nb    Blocking factor used to distribute the columns of the global matrix.
 * \param[in]     rsrc  Row rank of the process who owns the first row block of the global matrix.
 * \param[in]     csrc  Column rank of the process who owns the first column block of the global matrix.
 * \param[in]     lld   Leading dimension of the local matrix.
 * \param[in]     type  Data type of the matrix (e.g., CUDA_R_32F, CUDA_R_64F, CUDA_R_16F, CUDA_R_16BF,
 *                      CUDA_R_8F_E4M3, CUDA_R_8F_E5M2, CUDA_R_4F_E2M1, CUDA_R_8I, CUDA_R_32I,
 *                      CUDA_C_32F, CUDA_C_64F).
 * \param[in]     grid  Grid object associated with the matrix descriptor.
 * \param[in,out] desc  Matrix descriptor object to reinitialize.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatrixDescriptorInit(
    int64_t m,
    int64_t n,
    int64_t mb,
    int64_t nb,
    int64_t rsrc,
    int64_t csrc,
    int64_t lld,
    cudaDataType_t type,
    cublasMpGrid_t grid,
    cublasMpMatrixDescriptor_t desc);

/**
 * \brief Creates and initializes a matmul descriptor.
 *
 * This function initializes the values of a cublasMpMatmulDescriptor_t object
 * used in \ref cublasMpMatmul.
 *
 * \param[in,out] matmulDesc   Pointer to a cublasMpMatmulDescriptor object to initialize.
 * \param[in]     computeType  cuBLAS compute type used for computations.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmulDescriptorCreate(
    cublasMpMatmulDescriptor_t* matmulDesc,
    cublasComputeType_t computeType);

/**
 * \brief Destroys a matmul descriptor.
 *
 * This function destroys a cublasMpMatmulDescriptor_t object used in \ref cublasMpMatmul.
 *
 * \param[in,out] matmulDesc  Matmul descriptor object to destroy.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmulDescriptorDestroy(cublasMpMatmulDescriptor_t matmulDesc);

/**
 * \brief Initializes the values of an existing matmul descriptor.
 *
 * This function reinitializes a cublasMpMatmulDescriptor_t object without allocating
 * new memory.
 *
 * \param[in,out] matmulDesc   Matmul descriptor object to reinitialize.
 * \param[in]     computeType  cuBLAS compute type used for computations.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmulDescriptorInit(cublasMpMatmulDescriptor_t matmulDesc, cublasComputeType_t computeType);

/**
 * \brief Sets an attribute of a matmul descriptor.
 *
 * This function sets attributes of a cublasMpMatmulDescriptor_t object. The attributes
 * are of type \ref cublasMpMatmulDescriptorAttribute_t.
 *
 * \param[in] matmulDesc   Matmul descriptor object to set its attribute.
 * \param[in] attr         Matmul descriptor attribute to set.
 * \param[in] buf          Pointer to the attribute value to set.
 * \param[in] sizeInBytes  Attribute buffer size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmulDescriptorSetAttribute(
    cublasMpMatmulDescriptor_t matmulDesc,
    cublasMpMatmulDescriptorAttribute_t attr,
    const void* buf,
    size_t sizeInBytes);

/**
 * \brief Gets an attribute of a matmul descriptor.
 *
 * This function returns the attributes of a cublasMpMatmulDescriptor_t object.
 *
 * \param[in]  matmulDesc   Matmul descriptor object to query.
 * \param[in]  attr         Matmul descriptor attribute to query.
 * \param[out] buf          Buffer to receive the attribute value.
 * \param[in]  sizeInBytes  Attribute buffer size in bytes.
 * \param[out] sizeWritten  Size of the attribute written into buf in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmulDescriptorGetAttribute(
    cublasMpMatmulDescriptor_t matmulDesc,
    cublasMpMatmulDescriptorAttribute_t attr,
    void* buf,
    size_t sizeInBytes,
    size_t* sizeWritten);

/**
 * \brief Computes the workspace size required by \ref cublasMpTrsm.
 *
 * This function returns the required buffer sizes to perform cublasMpTrsm() on the
 * given input. This function requires square block size.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  side                       Indicates if matrix A is on the left or right of X.
 * \param[in]  uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]  trans                      Operation op(A) that is non- or (conj.) transpose.
 * \param[in]  diag                       Indicates if the elements on the main diagonal of matrix A are unity.
 * \param[in]  m                          Number of rows of matrix sub(B).
 * \param[in]  n                          Number of columns of matrix sub(B).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B). Must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). Must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  computeType                cuBLAS compute type for internal matrix-matrix multiplications.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTrsm_bufferSize(
    cublasMpHandle_t handle,
    cublasSideMode_t side,
    cublasFillMode_t uplo,
    cublasOperation_t trans,
    cublasDiagType_t diag,
    int64_t m,
    int64_t n,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    cublasComputeType_t computeType,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Solves a distributed triangular linear system with multiple right-hand-sides.
 *
 * This function solves:
 *   op(A) * X = alpha * B   if side == CUBLAS_SIDE_LEFT
 *   X * op(A) = alpha * B   if side == CUBLAS_SIDE_RIGHT
 *
 * where A is a triangular matrix stored in lower or upper mode with or without the
 * main diagonal, X and B are m x n matrices, and alpha is a scalar.
 * op(A) = A, A^T, or A^H depending on the trans parameter.
 * The solution X overwrites the right-hand-sides B on exit.
 *
 * This function requires square block size.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     side                       Indicates if matrix A is on the left or right of X.
 * \param[in]     uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]     trans                      Operation op(A) that is non- or (conj.) transpose.
 * \param[in]     diag                       Indicates if the diagonal elements of A are unity.
 * \param[in]     m                          Number of rows of matrix sub(B).
 * \param[in]     n                          Number of columns of matrix sub(B).
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]     descA                      Matrix descriptor for global matrix A.
 * \param[in,out] b                          Pointer to local portion of global matrix B (device).
 *                                           On exit, overwritten with the solution X.
 * \param[in]     ib                         Row index of the first row of sub(B). Must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). Must be a multiple of nbB.
 * \param[in]     descB                      Matrix descriptor for global matrix B.
 * \param[in]     computeType                cuBLAS compute type for internal matrix-matrix multiplications.
 * \param[out]    d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]     workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpTrsm_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpTrsm_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTrsm(
    cublasMpHandle_t handle,
    cublasSideMode_t side,
    cublasFillMode_t uplo,
    cublasOperation_t trans,
    cublasDiagType_t diag,
    int64_t m,
    int64_t n,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    cublasComputeType_t computeType,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost);

/**
 * \brief Computes the workspace size required by \ref cublasMpGemm.
 *
 * This function returns the required buffer sizes to perform cublasMpGemm() on the
 * given input.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  transA                     Operation op(A) that is non- or (conj.) transpose.
 * \param[in]  transB                     Operation op(B) that is non- or (conj.) transpose.
 * \param[in]  m                          Number of rows of op(A) and C.
 * \param[in]  n                          Number of columns of op(B) and C.
 * \param[in]  k                          Number of columns of op(A) and rows of op(B).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B). Must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). Must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  beta                       Scalar used for multiplication (host pointer).
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C.
 * \param[in]  computeType                cuBLAS compute type for internal matrix-matrix multiplications.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGemm_bufferSize(
    cublasMpHandle_t handle,
    cublasOperation_t transA,
    cublasOperation_t transB,
    int64_t m,
    int64_t n,
    int64_t k,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    cublasComputeType_t computeType,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Performs a distributed matrix-matrix multiplication (C = alpha * op(A) * op(B) + beta * C).
 *
 * This function performs the matrix-matrix multiplication:
 *   C = alpha * op(A) * op(B) + beta * C
 *
 * where alpha and beta are scalars, and A, B and C are matrices stored in column-major
 * format with dimensions op(A) m x k, op(B) k x n and C m x n.
 * op(A) = A, A^T, or A^H depending on transA.
 * op(B) = B, B^T, or B^H depending on transB.
 *
 * This routine will internally call cublasMpMatmul() with d == c.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     transA                     Operation op(A) that is non- or (conj.) transpose.
 * \param[in]     transB                     Operation op(B) that is non- or (conj.) transpose.
 * \param[in]     m                          Number of rows of op(A) and C.
 * \param[in]     n                          Number of columns of op(B) and C.
 * \param[in]     k                          Number of columns of op(A) and rows of op(B).
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]     descA                      Matrix descriptor for global matrix A.
 * \param[in]     b                          Pointer to local portion of global matrix B (device).
 * \param[in]     ib                         Row index of the first row of sub(B). Must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). Must be a multiple of nbB.
 * \param[in]     descB                      Matrix descriptor for global matrix B.
 * \param[in]     beta                       Scalar used for multiplication (host pointer).
 * \param[in,out] c                          Pointer to local portion of global matrix C (device).
 * \param[in]     ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]     descC                      Matrix descriptor for global matrix C.
 * \param[in]     computeType                cuBLAS compute type for internal matrix-matrix multiplications.
 * \param[out]    d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]     workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpGemm_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpGemm_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGemm(
    cublasMpHandle_t handle,
    cublasOperation_t transA,
    cublasOperation_t transB,
    int64_t m,
    int64_t n,
    int64_t k,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    cublasComputeType_t computeType,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost);

/**
 * \brief Computes the workspace size required by \ref cublasMpMatmul.
 *
 * This function returns the required buffer sizes to perform cublasMpMatmul() on the
 * given input.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  matmulDesc                 Descriptor of the operation, created with cublasMpMatmulDescriptorCreate().
 * \param[in]  m                          Number of rows of op(A) and C.
 * \param[in]  n                          Number of columns of op(B) and C.
 * \param[in]  k                          Number of columns of op(A) and rows of op(B).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B). Must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). Must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  beta                       Scalar used for multiplication (host pointer).
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C.
 * \param[out] d                          Pointer to local portion of global matrix D (device).
 * \param[in]  id                         Row index of the first row of sub(D). Must be a multiple of mbD.
 * \param[in]  jd                         Column index of the first column of sub(D). Must be a multiple of nbD.
 * \param[in]  descD                      Matrix descriptor for global matrix D.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmul_bufferSize(
    cublasMpHandle_t handle,
    cublasMpMatmulDescriptor_t matmulDesc,
    int64_t m,
    int64_t n,
    int64_t k,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    const void* beta,
    const void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    void* d,
    int64_t id,
    int64_t jd,
    cublasMpMatrixDescriptor_t descD,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Performs a distributed matrix-matrix multiplication (D = alpha * op(A) * op(B) + beta * C).
 *
 * This function performs the matrix-matrix multiplication:
 *   D = alpha * op(A) * op(B) + beta * C
 *
 * where alpha and beta are scalars, and A, B, C and D are matrices stored in column-major
 * format with dimensions op(A) m x k, op(B) k x n, C m x n and D m x n.
 * op(A) depends on the TRANSA attribute of the matmulDesc descriptor.
 * op(B) depends on the TRANSB attribute of the matmulDesc descriptor.
 *
 * This is the recommended API for tensor parallelism operations such as AllGather+GEMM,
 * GEMM+ReduceScatter, and GEMM+AllReduce with communication-computation overlap.
 *
 * FP8 matrix multiplication is only supported for the TN format on Compute Capability 9.0+ GPUs.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  matmulDesc                 Descriptor of the operation, created with cublasMpMatmulDescriptorCreate().
 * \param[in]  m                          Number of rows of op(A) and C.
 * \param[in]  n                          Number of columns of op(B) and C.
 * \param[in]  k                          Number of columns of op(A) and rows of op(B).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B). Must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). Must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  beta                       Scalar used for multiplication (host pointer).
 * \param[in]  c                          Pointer to local portion of global matrix C (device). Can be NULL if beta == 0.
 * \param[in]  ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C.
 * \param[out] d                          Pointer to local portion of global matrix D (device).
 * \param[in]  id                         Row index of the first row of sub(D). Must be a multiple of mbD.
 * \param[in]  jd                         Column index of the first column of sub(D). Must be a multiple of nbD.
 * \param[in]  descD                      Matrix descriptor for global matrix D.
 * \param[out] d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]  workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpMatmul_bufferSize().
 * \param[out] h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]  workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpMatmul_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMatmul(
    cublasMpHandle_t handle,
    cublasMpMatmulDescriptor_t matmulDesc,
    int64_t m,
    int64_t n,
    int64_t k,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    const void* beta,
    const void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    void* d,
    int64_t id,
    int64_t jd,
    cublasMpMatrixDescriptor_t descD,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost);

/**
 * \brief Computes the workspace size required by \ref cublasMpSyrk.
 *
 * This function returns the required buffer sizes to perform cublasMpSyrk() on the
 * given input. This function requires square block size.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  uplo                       Indicates if matrix C lower or upper part is stored.
 * \param[in]  trans                      Operation op(A) that is non- or transpose.
 * \param[in]  n                          Number of rows of op(A) and C.
 * \param[in]  k                          Number of columns of op(A).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  beta                       Scalar used for multiplication (host pointer).
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C.
 * \param[in]  computeType                cuBLAS compute type for internal matrix-matrix multiplications.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSyrk_bufferSize(
    cublasMpHandle_t handle,
    cublasFillMode_t uplo,
    cublasOperation_t trans,
    int64_t n,
    int64_t k,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    cublasComputeType_t computeType,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Performs a distributed symmetric rank-k update (C = alpha * op(A) * op(A)^T + beta * C).
 *
 * This function performs the symmetric rank-k update:
 *   C = alpha * op(A) * op(A)^T + beta * C
 *
 * where alpha and beta are scalars, C is a symmetric matrix stored in lower or upper
 * mode, and A is a matrix with dimensions op(A) n x k.
 * op(A) = A if trans == CUBLAS_OP_N, or A^T if trans == CUBLAS_OP_T.
 *
 * This function requires square block size.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     uplo                       Indicates if matrix C lower or upper part is stored.
 * \param[in]     trans                      Operation op(A) that is non- or transpose.
 * \param[in]     n                          Number of rows of op(A) and C.
 * \param[in]     k                          Number of columns of op(A).
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]     descA                      Matrix descriptor for global matrix A.
 * \param[in]     beta                       Scalar used for multiplication (host pointer).
 * \param[in,out] c                          Pointer to local portion of global matrix C (device).
 * \param[in]     ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]     descC                      Matrix descriptor for global matrix C.
 * \param[in]     computeType                cuBLAS compute type for internal matrix-matrix multiplications.
 * \param[out]    d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]     workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpSyrk_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpSyrk_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSyrk(
    cublasMpHandle_t handle,
    cublasFillMode_t uplo,
    cublasOperation_t trans,
    int64_t n,
    int64_t k,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    cublasComputeType_t computeType,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost);

/**
 * \brief Computes the number of rows or columns of a distributed matrix owned by a process.
 *
 * Computes the number of rows or columns of a distributed matrix owned by the process
 * indicated by the iproc argument.
 *
 * \param[in] n         Number of rows or columns in the global distributed matrix.
 * \param[in] nb        Row or column blocking size of the global matrix.
 * \param[in] iproc     Coordinate of the process whose local array row or column is to be determined.
 * \param[in] isrcproc  Coordinate of the process that owns the first row or column of the distributed matrix.
 * \param[in] nprocs    Total number of row or column processes over which the matrix is distributed.
 * \return Number of local rows or columns owned by the specified process.
 */
int64_t cublasMpNumroc(int64_t n, int64_t nb, uint32_t iproc, uint32_t isrcproc, uint32_t nprocs);

/**
 * \brief Returns a human-readable string description for a status code.
 *
 * This function returns a human-readable string description for a given
 * cublasMpStatus_t error code. Useful for error handling and debugging.
 * The returned string is a static string and remains valid for the lifetime
 * of the application. For unknown status codes, the function returns a string
 * indicating an unknown status.
 *
 * \param[in] status  The cublasMpStatus_t status code to convert to a string.
 * \return Pointer to a null-terminated string containing the status description.
 */
const char* cublasMpGetStatusString(cublasMpStatus_t status);

/**
 * \brief Computes the workspace size required by \ref cublasMpGemr2D.
 *
 * This function returns the required buffer sizes to perform cublasMpGemr2D() on the
 * given input.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  m                          Number of rows of sub(A) and sub(B).
 * \param[in]  n                          Number of columns of sub(A) and sub(B).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A).
 * \param[in]  ja                         Column index of the first column of sub(A).
 * \param[in]  descA                      Matrix descriptor for global matrix A. Grid must be set to
 *                                        null for processes not part of the grid of A.
 * \param[in]  b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B).
 * \param[in]  jb                         Column index of the first column of sub(B).
 * \param[in]  descB                      Matrix descriptor for global matrix B. Grid must be set to
 *                                        null for processes not part of the grid of B.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \param[in]  global_comm                Communicator containing at least the union of all processes
 *                                        in the communicators of A and B.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGemr2D_bufferSize(
    cublasMpHandle_t handle,
    int64_t m,
    int64_t n,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost,
    ncclComm_t global_comm);

/**
 * \brief Redistributes a general rectangular matrix from one distribution to another.
 *
 * This function redistributes general rectangular matrix A according to the distribution
 * properties of matrix B. All processes in the global communicator must call this function,
 * even if they do not own a piece of either matrix.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  m                          Number of rows of sub(A) and sub(B).
 * \param[in]  n                          Number of columns of sub(A) and sub(B).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A).
 * \param[in]  ja                         Column index of the first column of sub(A).
 * \param[in]  descA                      Matrix descriptor for global matrix A. Grid must be set to
 *                                        null for processes not part of the grid of A.
 * \param[out] b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B).
 * \param[in]  jb                         Column index of the first column of sub(B).
 * \param[in]  descB                      Matrix descriptor for global matrix B. Grid must be set to
 *                                        null for processes not part of the grid of B.
 * \param[out] d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]  workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpGemr2D_bufferSize().
 * \param[out] h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]  workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpGemr2D_bufferSize().
 * \param[in]  global_comm                Communicator containing at least the union of all processes
 *                                        in the communicators of A and B.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGemr2D(
    cublasMpHandle_t handle,
    int64_t m,
    int64_t n,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost,
    ncclComm_t global_comm);

/**
 * \brief Computes the workspace size required by \ref cublasMpTrmr2D.
 *
 * This function returns the required buffer sizes to perform cublasMpTrmr2D() on the
 * given input.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]  diag                       Indicates if diagonal elements of A are unity.
 * \param[in]  m                          Number of rows of sub(A) and sub(B).
 * \param[in]  n                          Number of columns of sub(A) and sub(B).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A).
 * \param[in]  ja                         Column index of the first column of sub(A).
 * \param[in]  descA                      Matrix descriptor for global matrix A. Grid must be set to
 *                                        null for processes not part of the grid of A.
 * \param[in]  b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B).
 * \param[in]  jb                         Column index of the first column of sub(B).
 * \param[in]  descB                      Matrix descriptor for global matrix B. Grid must be set to
 *                                        null for processes not part of the grid of B.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \param[in]  global_comm                Communicator containing at least the union of all processes
 *                                        in the communicators of A and B.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTrmr2D_bufferSize(
    cublasMpHandle_t handle,
    cublasFillMode_t uplo,
    cublasDiagType_t diag,
    int64_t m,
    int64_t n,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost,
    ncclComm_t global_comm);

/**
 * \brief Redistributes a trapezoidal matrix from one distribution to another.
 *
 * This function redistributes trapezoidal matrix A according to the distribution
 * properties of trapezoidal matrix B. All processes in the global communicator must
 * call this function, even if they do not own a piece of either matrix.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]  diag                       Indicates if diagonal elements of A are unity.
 * \param[in]  m                          Number of rows of sub(A) and sub(B).
 * \param[in]  n                          Number of columns of sub(A) and sub(B).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A).
 * \param[in]  ja                         Column index of the first column of sub(A).
 * \param[in]  descA                      Matrix descriptor for global matrix A. Grid must be set to
 *                                        null for processes not part of the grid of A.
 * \param[out] b                          Pointer to local portion of global matrix B (device).
 * \param[in]  ib                         Row index of the first row of sub(B).
 * \param[in]  jb                         Column index of the first column of sub(B).
 * \param[in]  descB                      Matrix descriptor for global matrix B. Grid must be set to
 *                                        null for processes not part of the grid of B.
 * \param[out] d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]  workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpTrmr2D_bufferSize().
 * \param[out] h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]  workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpTrmr2D_bufferSize().
 * \param[in]  global_comm                Communicator containing at least the union of all processes
 *                                        in the communicators of A and B.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTrmr2D(
    cublasMpHandle_t handle,
    cublasFillMode_t uplo,
    cublasDiagType_t diag,
    int64_t m,
    int64_t n,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost,
    ncclComm_t global_comm);

/**
 * \brief Computes the workspace size required by \ref cublasMpGeadd.
 *
 * This function returns the required buffer sizes to perform cublasMpGeadd() on the
 * given input. This function requires square block size.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  trans                      Operation op(A) that is non- or (conj.) transpose.
 * \param[in]  m                          Number of rows of sub(A) and sub(C).
 * \param[in]  n                          Number of columns of sub(A) and sub(C).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  beta                       Scalar used for multiplication (host pointer).
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGeadd_bufferSize(
    cublasMpHandle_t handle,
    cublasOperation_t trans,
    int64_t m,
    int64_t n,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Performs a distributed general matrix-matrix addition (C = alpha * op(A) + beta * C).
 *
 * This function performs the matrix-matrix addition:
 *   C = alpha * op(A) + beta * C
 *
 * where alpha and beta are scalars, and A and C are matrices stored in column-major
 * format with dimensions op(A) m x n and C m x n.
 * op(A) = A, A^T, or A^H depending on the trans parameter.
 *
 * This function requires square block size.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     trans                      Operation op(A) that is non- or (conj.) transpose.
 * \param[in]     m                          Number of rows of sub(A) and sub(C).
 * \param[in]     n                          Number of columns of sub(A) and sub(C).
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]     descA                      Matrix descriptor for global matrix A.
 * \param[in]     beta                       Scalar used for multiplication (host pointer).
 * \param[in,out] c                          Pointer to local portion of global matrix C (device).
 * \param[in]     ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]     descC                      Matrix descriptor for global matrix C.
 * \param[out]    d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]     workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpGeadd_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpGeadd_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGeadd(
    cublasMpHandle_t handle,
    cublasOperation_t trans,
    int64_t m,
    int64_t n,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost);

/**
 * \brief Computes the workspace size required by \ref cublasMpTradd.
 *
 * This function returns the required buffer sizes to perform cublasMpTradd() on the
 * given input. This function requires square block size.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  uplo                       Indicates if matrix C lower or upper part is stored.
 * \param[in]  trans                      Operation op(A) that is non- or (conj.) transpose.
 * \param[in]  m                          Number of rows of sub(A) and sub(C).
 * \param[in]  n                          Number of columns of sub(A) and sub(C).
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  beta                       Scalar used for multiplication (host pointer).
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C.
 * \param[out] workspaceSizeInBytesOnDevice  Required device workspace size in bytes.
 * \param[out] workspaceSizeInBytesOnHost    Required host workspace size in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTradd_bufferSize(
    cublasMpHandle_t handle,
    cublasFillMode_t uplo,
    cublasOperation_t trans,
    int64_t m,
    int64_t n,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Performs a distributed trapezoidal matrix-matrix addition (C = alpha * op(A) + beta * C).
 *
 * This function performs the trapezoidal matrix-matrix addition:
 *   C = alpha * op(A) + beta * C
 *
 * where alpha and beta are scalars, and A and C are matrices stored in column-major
 * format with dimensions op(A) m x n and C m x n.
 * op(A) = A, A^T, or A^H depending on the trans parameter.
 * Only the upper or lower trapezoidal part of C is updated, as indicated by uplo.
 *
 * This function requires square block size.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     uplo                       Indicates if matrix C lower or upper part is stored.
 * \param[in]     trans                      Operation op(A) that is non- or (conj.) transpose.
 * \param[in]     m                          Number of rows of sub(A) and sub(C).
 * \param[in]     n                          Number of columns of sub(A) and sub(C).
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). Must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). Must be a multiple of nbA.
 * \param[in]     descA                      Matrix descriptor for global matrix A.
 * \param[in]     beta                       Scalar used for multiplication (host pointer).
 * \param[in,out] c                          Pointer to local portion of global matrix C (device).
 * \param[in]     ic                         Row index of the first row of sub(C). Must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). Must be a multiple of nbC.
 * \param[in]     descC                      Matrix descriptor for global matrix C.
 * \param[out]    d_work                     Device workspace of size workspaceSizeInBytesOnDevice.
 * \param[in]     workspaceSizeInBytesOnDevice  Device workspace size as provided by cublasMpTradd_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpTradd_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTradd(
    cublasMpHandle_t handle,
    cublasFillMode_t uplo,
    cublasOperation_t trans,
    int64_t m,
    int64_t n,
    const void* alpha,
    const void* a,
    int64_t ia,
    int64_t ja,
    cublasMpMatrixDescriptor_t descA,
    const void* beta,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    void* d_work,
    size_t workspaceSizeInBytesOnDevice,
    void* h_work,
    size_t workspaceSizeInBytesOnHost);

/**
 * \brief Allocates and registers device memory with NCCL symmetric memory.
 *
 * This is a convenience function that allocates device memory using NCCL symmetric
 * memory and automatically registers it with the specified process grid. It combines
 * ncclMemAlloc and \ref cublasMpBufferRegister into a single call.
 *
 * Alternatively, users can use ncclMemAlloc, CUDA VMM APIs, or other CUDA allocation
 * methods followed by manual registration with \ref cublasMpBufferRegister.
 *
 * All the processes defined to be in this grid must enter this function.
 * Use \ref cublasMpFree to free buffers allocated with this function.
 *
 * \param[in]  grid  Process grid to allocate and register the buffer with.
 * \param[out] ptr   Pointer to receive the allocated device memory pointer.
 * \param[in]  size  Size of the buffer to allocate in bytes.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpMalloc(cublasMpGrid_t grid, void** ptr, size_t size);

/**
 * \brief Deregisters and frees device memory allocated with \ref cublasMpMalloc.
 *
 * This is a convenience function that combines \ref cublasMpBufferDeregister and
 * ncclMemFree into a single call. It should only be used for buffers allocated with
 * \ref cublasMpMalloc. For buffers allocated manually, use \ref cublasMpBufferDeregister
 * followed by the appropriate free function.
 *
 * All the processes defined to be in this grid must enter this function.
 *
 * \param[in] grid  Process grid the buffer was allocated with.
 * \param[in] ptr   Pointer to device memory buffer to free.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpFree(cublasMpGrid_t grid, void* ptr);

/**
 * \brief Function pointer type for cuBLASMp logging callbacks.
 *
 * The callback function is called by the library to report logging information.
 * It can be set using \ref cublasMpLoggerSetCallback.
 *
 * \param logLevel      The severity level of the log message.
 * \param functionName  Name of the function generating the log message.
 * \param message       The actual log message text.
 */
typedef void (*cublasMpLoggerCallback_t)(int logLevel, const char* functionName, const char* message);

/**
 * \brief Sets the logging callback function.
 *
 * \warning This is an experimental feature.
 *
 * \param[in] callback  Pointer to a callback function (cublasMpLoggerCallback_t).
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpLoggerSetCallback(cublasMpLoggerCallback_t callback);

/**
 * \brief Sets the logging output file.
 *
 * Once registered using this function call, the provided file handle must not be
 * closed unless the function is called again to switch to a different file handle.
 *
 * \warning This is an experimental feature.
 *
 * \param[in] file  Pointer to an open file with write permission.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpLoggerSetFile(FILE* file);

/**
 * \brief Opens a logging output file at the given path.
 *
 * \warning This is an experimental feature.
 *
 * \param[in] logFile  Path of the logging output file.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpLoggerOpenFile(const char* logFile);

/**
 * \brief Sets the logging level.
 *
 * Log levels: 0=Off, 1=Error, 2=Trace, 3=Hints, 4=Info, 5=API Trace.
 *
 * \warning This is an experimental feature.
 *
 * \param[in] level  Value of the logging level.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpLoggerSetLevel(int level);

/**
 * \brief Sets the value of the logging mask.
 *
 * Log mask values can be combined: 0=Off, 1=Error, 2=Trace, 4=Hints, 8=Info, 16=API Trace.
 *
 * \warning This is an experimental feature.
 *
 * \param[in] mask  Value of the logging mask.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpLoggerSetMask(int mask);

/**
 * \brief Disables logging for the entire run.
 *
 * \warning This is an experimental feature.
 *
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpLoggerForceDisable();

#ifdef __cplusplus
}
#endif
