/*
 * 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 9
#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 and
     *  GEMM+AllReduce 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),
    /** Apply ReLU gradient to matmul output. Store ReLU gradient in the output matrix.
     *  Requires an extra input (ReLU bit-mask from forward pass) via EPILOGUE_AUX_POINTER. */
    CUBLASMP_MATMUL_EPILOGUE_DRELU = (8 | 128),
    /** Apply independently ReLU and Bias gradients to matmul output. Store ReLU gradient in
     *  the output matrix and Bias gradient in the bias buffer. Requires an extra input
     *  (ReLU bit-mask from forward pass) via EPILOGUE_AUX_POINTER. */
    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. Most compute routines use descriptors associated with a non-null process grid.
 * Redistribution routines (\ref cublasMpGemr2D and \ref cublasMpTrmr2D) also allow
 * descriptors whose grid is NULL on ranks that do not participate in that matrix's grid.
 * Descriptors 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 handle remains associated with the CUDA
 * device that was current when cublasMpCreate() was called.
 *
 * \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.
 * The new stream must belong to the same CUDA context as the stream associated
 * with the handle at creation time.
 *
 * \param[in] handle  cuBLASMp library handle.
 * \param[in] stream  CUDA stream to set. The stream must belong to the same
 *                    CUDA context as the current handle stream.
 * \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.
 *
 * Sets 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 Sets the maximum number of mantissa bits for FP64 fixed-point emulation.
 *
 * Specifies the maximum number of mantissa bits used in fixed-point emulation
 * of FP64 matrix multiplications. In dynamic control mode, if the required
 * bit count exceeds this value, the library falls back to native FP64. In fixed
 * control mode, this value is used directly as the bit count.
 *
 * This setting is applied to both the underlying cuBLAS and cuBLASLt handles.
 * Requires CUDA Toolkit 13.0.2 or later; earlier toolkits return
 * CUBLASMP_STATUS_NOT_SUPPORTED.
 *
 * \param[in] handle    cuBLASMp library handle.
 * \param[in] bitCount  Maximum number of mantissa bits (default: 79 for dynamic,
 *                      55 for fixed mode).
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetFixedPointEmulationMaxMantissaBitCount(cublasMpHandle_t handle, int bitCount);

/**
 * \brief Gets the maximum number of mantissa bits for FP64 fixed-point emulation.
 *
 * Requires CUDA Toolkit 13.0.2 or later; earlier toolkits return
 * CUBLASMP_STATUS_NOT_SUPPORTED.
 *
 * \param[in]  handle    cuBLASMp library handle.
 * \param[out] bitCount  Pointer to receive the current maximum mantissa bit count.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetFixedPointEmulationMaxMantissaBitCount(cublasMpHandle_t handle, int* bitCount);

/**
 * \brief Sets the mantissa bit offset for FP64 fixed-point emulation.
 *
 * Adjusts the automatic dynamic precision framework to use fewer or more bits
 * than required for accuracy equal to native FP64. A positive offset increases
 * the bit count (higher accuracy, lower performance); a negative offset
 * decreases it (lower accuracy, higher performance).
 *
 * This setting is applied to both the underlying cuBLAS and cuBLASLt handles.
 * Requires CUDA Toolkit 13.0.2 or later; earlier toolkits return
 * CUBLASMP_STATUS_NOT_SUPPORTED.
 *
 * \param[in] handle     cuBLASMp library handle.
 * \param[in] bitOffset  Mantissa bit offset to apply to the automatic precision
 *                       framework (default: 0).
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetFixedPointEmulationMantissaBitOffset(cublasMpHandle_t handle, int bitOffset);

/**
 * \brief Gets the mantissa bit offset for FP64 fixed-point emulation.
 *
 * Requires CUDA Toolkit 13.0.2 or later; earlier toolkits return
 * CUBLASMP_STATUS_NOT_SUPPORTED.
 *
 * \param[in]  handle     cuBLASMp library handle.
 * \param[out] bitOffset  Pointer to receive the current mantissa bit offset.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetFixedPointEmulationMantissaBitOffset(cublasMpHandle_t handle, int* bitOffset);

/**
 * \brief Sets a device pointer to receive the actual mantissa bit count used during
 * FP64 fixed-point emulation.
 *
 * When set, the library writes the number of mantissa bits actually retained by the
 * emulation to the provided device pointer after each FP64 emulated computation.
 * This is useful for diagnostics and performance analysis.
 *
 * This setting is applied to both the underlying cuBLAS and cuBLASLt handles.
 * Requires CUDA Toolkit 13.0.2 or later; earlier toolkits return
 * CUBLASMP_STATUS_NOT_SUPPORTED.
 *
 * \param[in] handle  cuBLASMp library handle.
 * \param[in] ptr     Device pointer to an int where the bit count will be written
 *                    (default: NULL, meaning the bit count is not reported).
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetFixedPointEmulationMantissaBitCountPointer(cublasMpHandle_t handle, int* ptr);

/**
 * \brief Gets the current device pointer for the mantissa bit count output.
 *
 * Requires CUDA Toolkit 13.0.2 or later; earlier toolkits return
 * CUBLASMP_STATUS_NOT_SUPPORTED.
 *
 * \param[in]  handle  cuBLASMp library handle.
 * \param[out] ptr     Pointer to receive the current device pointer.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetFixedPointEmulationMantissaBitCountPointer(cublasMpHandle_t handle, int** ptr);

#if CUBLAS_VERSION >= 130002

/**
 * \brief Sets the FP64 fixed-point emulation mantissa control mode.
 *
 * Requires CUDA Toolkit 13.0.2 or later. This declaration is available only when
 * compiling with CUDA Toolkit 13.0.2 or later because the CUDA enum type is
 * introduced there.
 *
 * Configures how the number of mantissa bits is determined for fixed-point
 * emulation of FP64 matrix multiplications. With dynamic control, the library
 * computes the required bit count at runtime to maintain accuracy equal to or
 * better than native FP64. With fixed control, a user-provided bit count is
 * used.
 *
 * This setting is applied to both the underlying cuBLAS and cuBLASLt handles.
 * The FP64 emulation settings must be configured after handle creation and
 * before performing computations.
 *
 * \param[in] handle           cuBLASMp library handle.
 * \param[in] mantissaControl  Mantissa control mode (CUDA_EMULATION_MANTISSA_CONTROL_DYNAMIC
 *                             or CUDA_EMULATION_MANTISSA_CONTROL_FIXED).
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetFixedPointEmulationMantissaControl(
    cublasMpHandle_t handle,
    cudaEmulationMantissaControl_t mantissaControl);

/**
 * \brief Gets the current FP64 fixed-point emulation mantissa control mode.
 *
 * Requires CUDA Toolkit 13.0.2 or later. This declaration is available only when
 * compiling with CUDA Toolkit 13.0.2 or later because the CUDA enum type is
 * introduced there.
 *
 * \param[in]  handle           cuBLASMp library handle.
 * \param[out] mantissaControl  Pointer to receive the current mantissa control mode.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetFixedPointEmulationMantissaControl(
    cublasMpHandle_t handle,
    cudaEmulationMantissaControl_t* mantissaControl);

/**
 * \brief Sets the special values support bitmask for FP64 fixed-point emulation.
 *
 * Requires CUDA Toolkit 13.0.2 or later. This declaration is available only when
 * compiling with CUDA Toolkit 13.0.2 or later because the CUDA enum type is
 * introduced there.
 *
 * Configures which special floating-point values (e.g. NaN, Inf, denormals) are
 * supported during fixed-point emulation of FP64 matrix multiplications. The mask
 * is a bitmask of cudaEmulationSpecialValuesSupport_t flags.
 *
 * This setting is applied to both the underlying cuBLAS and cuBLASLt handles.
 *
 * \param[in] handle  cuBLASMp library handle.
 * \param[in] mask    Bitmask of special values to support (default: 0xFFFF).
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSetEmulationSpecialValuesSupport(
    cublasMpHandle_t handle,
    cudaEmulationSpecialValuesSupport_t mask);

/**
 * \brief Gets the current special values support bitmask for FP64 fixed-point emulation.
 *
 * Requires CUDA Toolkit 13.0.2 or later. This declaration is available only when
 * compiling with CUDA Toolkit 13.0.2 or later because the CUDA enum type is
 * introduced there.
 *
 * \param[in]  handle  cuBLASMp library handle.
 * \param[out] mask    Pointer to receive the current special values support bitmask.
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpGetEmulationSpecialValuesSupport(
    cublasMpHandle_t handle,
    cudaEmulationSpecialValuesSupport_t* mask);

#endif // CUBLAS_VERSION >= 130002

/**
 * \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.
 * The communicator must contain exactly nprow * npcol ranks, and all participating
 * ranks must pass the same nprow, npcol, and layout values.
 *
 * All ranks in the communicator must enter this function collectively.
 *
 * \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 ranks in the grid
 * communicator must enter this function collectively.
 *
 * \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.
 *
 * Registering the device workspace buffer (d_work) enables cublasMpMatmul to use the
 * communication-computation overlap algorithms for AllGather+GEMM, GEMM+ReduceScatter, and GEMM+AllReduce operations.
 * Without a registered workspace, the library falls back to the no_overlap algorithm when using
 * the DEFAULT algorithm type, or returns CUBLASMP_STATUS_NOT_SUPPORTED for explicit pipelined
 * algorithm requests. Input/output matrix buffers do not need to be registered; Matrix B buffer registration
 * is optional and may only improve performance for the AllGather+GEMM algorithm.
 *
 * The buffer must be allocated using ncclMemAlloc, CUDA Virtual Memory Management (VMM) APIs,
 * or other VMM-compatible methods.
 *
 * All ranks in the grid communicator must enter this function collectively.
 *
 * \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 ranks in the grid communicator must enter this function collectively.
 *
 * \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.
 * Parameters mb, nb, and lld are expressed in elements and must be positive.
 * When grid is non-NULL, rsrc and csrc are zero-based source coordinates and must
 * satisfy 0 <= rsrc < grid->nprow and 0 <= csrc < grid->npcol. The local leading
 * dimension lld must be large enough for the local allocation on the calling rank.
 * The grid object, when non-NULL, must outlive the descriptor. Redistribution
 * routines (\ref cublasMpGemr2D and \ref cublasMpTrmr2D) also allow grid == NULL on
 * ranks that do not participate in that matrix's process grid.
 *
 * \param[in]  m     Number of rows in the global matrix.
 * \param[in]  n     Number of columns in the global matrix.
 * \param[in]  mb    Positive blocking factor used to distribute the rows of the global matrix.
 * \param[in]  nb    Positive blocking factor used to distribute the columns of the global matrix.
 * \param[in]  rsrc  Row rank of the process that owns the first row block of the global matrix.
 *                   When grid is non-NULL, must satisfy 0 <= rsrc < grid->nprow.
 * \param[in]  csrc  Column rank of the process that owns the first column block of the global matrix.
 *                   When grid is non-NULL, must satisfy 0 <= csrc < grid->npcol.
 * \param[in]  lld   Positive leading dimension of the local matrix. It must be large
 *                   enough for the local allocation on the calling rank.
 * \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, or NULL on
 *                   redistribution-only ranks that do not participate in this matrix's grid.
 * \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. The same validity requirements apply: mb/nb/lld must be
 * positive and, when grid is non-NULL, rsrc/csrc must be valid source coordinates
 * for that grid. The grid object, when non-NULL, must outlive the descriptor.
 * Redistribution routines (\ref cublasMpGemr2D and \ref cublasMpTrmr2D) also allow
 * grid == NULL on ranks that do not participate in that matrix's process grid.
 *
 * \param[in]     m     Number of rows in the global matrix.
 * \param[in]     n     Number of columns in the global matrix.
 * \param[in]     mb    Positive blocking factor used to distribute the rows of the global matrix.
 * \param[in]     nb    Positive blocking factor used to distribute the columns of the global matrix.
 * \param[in]     rsrc  Row rank of the process that owns the first row block of the global matrix.
 *                      When grid is non-NULL, must satisfy 0 <= rsrc < grid->nprow.
 * \param[in]     csrc  Column rank of the process that owns the first column block of the global matrix.
 *                      When grid is non-NULL, must satisfy 0 <= csrc < grid->npcol.
 * \param[in]     lld   Positive leading dimension of the local matrix. It must be large
 *                      enough for the local allocation on the calling rank.
 * \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, or NULL on
 *                      redistribution-only ranks that do not participate in this matrix's grid.
 * \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.
 *
 * Contract:
 * - ia, ja, ib, and jb are 1-based and may be misaligned within a tile.
 * - m and n must be non-negative.
 * - descA and descB must have non-NULL grids, matching data types, and compatible
 *   process grids. Compatible grids are built from the same parent NCCL communicator
 *   and have the same nprow, npcol, layout, and per-rank coordinates; pointer identity
 *   of the grid objects is not required.
 * - The submatrices must fit inside descA and descB. The triangular dimension is m
 *   when side == CUBLAS_SIDE_LEFT and n when side == CUBLAS_SIDE_RIGHT.
 * - This query does not dereference alpha, a, or b.
 *
 * Workspace:
 * - workspaceSizeInBytesOnDevice and workspaceSizeInBytesOnHost must be non-NULL
 *   output pointers.
 *
 * \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).
 * \param[in]  ja                         Column index of the first column of sub(A).
 * \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).
 * \param[in]  jb                         Column index of the first column of sub(B).
 * \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.
 *
 * Contract:
 * - ia, ja, ib, and jb are 1-based and may be misaligned within a tile.
 * - m and n must be non-negative.
 * - descA and descB must have non-NULL grids, matching data types, and compatible
 *   process grids. Compatible grids are built from the same parent NCCL communicator
 *   and have the same nprow, npcol, layout, and per-rank coordinates; pointer identity
 *   of the grid objects is not required.
 * - The submatrices must fit inside descA and descB. The triangular dimension is m
 *   when side == CUBLAS_SIDE_LEFT and n when side == CUBLAS_SIDE_RIGHT.
 *
 * Workspace and null pointers:
 * - d_work and h_work are caller-provided temporary workspaces.
 * - The matrix pointers a and b may be NULL only on ranks that own no local data for
 *   the referenced matrices; ranks owning local data must pass valid device pointers.
 *
 * \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).
 * \param[in]     ja                         Column index of the first column of sub(A).
 * \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).
 * \param[in]     jb                         Column index of the first column of sub(B).
 * \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. cublasMpGemm() uses cublasMpMatmul() semantics with D == C.
 *
 * Contract:
 * - Submatrix indices are 1-based and must be aligned to the corresponding descriptor
 *   block dimensions.
 * - descA, descB, and descC must be non-NULL. Matrix C is both input and output;
 *   unlike cublasMpMatmul(), cublasMpGemm() always writes the result to C.
 * - Block-size compatibility follows op(A) and op(B):
 *   transA=N, transB=N: nbA == mbB, mbA == mbC, nbB == nbC.
 *   transA=N, transB=T/C: nbA == nbB, mbA == mbC, mbB == nbC.
 *   transA=T/C, transB=N: mbA == mbB, nbA == mbC, nbB == nbC.
 *   transA=T/C, transB=T/C: mbA == nbB, nbA == mbC, mbB == nbC.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpGemm() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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.
 *
 * Contract:
 * - Uses the same descriptor, offset, and block-size contract as
 *   \ref cublasMpGemm_bufferSize.
 * - Matrix C is both input and output; unlike cublasMpMatmul(), cublasMpGemm()
 *   always writes the result to C.
 *
 * Workspace:
 * - Call cublasMpGemm_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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.
 *
 * Contract:
 * - Submatrix indices are 1-based and must be aligned to the corresponding descriptor
 *   block dimensions.
 * - descA, descB, and descD must be non-NULL.
 * - Block-size compatibility follows the transA and transB attributes in matmulDesc:
 *   transA=N, transB=N: nbA == mbB, mbA == mbD, nbB == nbD.
 *   transA=N, transB=T/C: nbA == nbB, mbA == mbD, mbB == nbD.
 *   transA=T/C, transB=N: mbA == mbB, nbA == mbD, nbB == nbD.
 *   transA=T/C, transB=T/C: mbA == nbB, nbA == mbD, mbB == nbD.
 * - This query does not dereference alpha. If c is NULL, beta must be a valid host
 *   pointer to a scalar zero; the query dereferences beta to validate that
 *   condition. In that case descC may also be NULL, and descD defines the output
 *   layout used in place of C.
 * - If c is non-NULL, descC must be non-NULL and C/D output blocking must match:
 *   mbC == mbD and nbC == nbD.
 *
 * Workspace:
 * - The matmul descriptor attributes, dimensions, descriptors, offsets, and null-c
 *   choice used for this query must match the subsequent cublasMpMatmul() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  beta                       Scalar used for multiplication (host pointer). Must point
 *                                        to a scalar zero when c is NULL.
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 *                                        May be NULL for a matching call that uses beta == 0.
 * \param[in]  ic                         Row index of the first row of sub(C). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C. May be NULL when
 *                                        c is NULL for a matching call that uses beta == 0; in that
 *                                        case descD provides the output layout.
 * \param[out] d                          Pointer to local portion of global matrix D (device).
 * \param[in]  id                         Row index of the first row of sub(D). id - 1 must be a multiple of mbD.
 * \param[in]  jd                         Column index of the first column of sub(D). jd - 1 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.
 *
 * Contract:
 * - Submatrix indices are 1-based and must be aligned to the corresponding descriptor
 *   block dimensions.
 * - descA, descB, and descD must be non-NULL.
 * - Block-size compatibility follows the transA and transB attributes in matmulDesc:
 *   transA=N, transB=N: nbA == mbB, mbA == mbD, nbB == nbD.
 *   transA=N, transB=T/C: nbA == nbB, mbA == mbD, mbB == nbD.
 *   transA=T/C, transB=N: mbA == mbB, nbA == mbD, nbB == nbD.
 *   transA=T/C, transB=T/C: mbA == nbB, nbA == mbD, mbB == nbD.
 * - If c is NULL, the beta scalar value must be zero. In that case descC may also
 *   be NULL, and descD defines the output layout used in place of C.
 * - If c is non-NULL, descC must be non-NULL and C/D output blocking must match:
 *   mbC == mbD and nbC == nbD.
 *
 * Workspace:
 * - Call cublasMpMatmul_bufferSize() with matching matmul descriptor attributes,
 *   dimensions, descriptors, offsets, and null-c choice. Pass device and host
 *   workspaces of the returned sizes to this function.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 must be a multiple of nbC.
 * \param[in]  descC                      Matrix descriptor for global matrix C. May be NULL when
 *                                        c is NULL and beta == 0; in that case descD provides the
 *                                        output layout.
 * \param[out] d                          Pointer to local portion of global matrix D (device).
 * \param[in]  id                         Row index of the first row of sub(D). id - 1 must be a multiple of mbD.
 * \param[in]  jd                         Column index of the first column of sub(D). jd - 1 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.
 *
 * Contract:
 * - ia, ja, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA and descC must be non-NULL.
 * - C must have square block size (mbC == nbC).
 * - The n-dimension blocking of op(A) must match C: mbA == mbC when trans=N,
 *   nbA == mbC when trans=T.
 * - uplo selects the stored triangle of C; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpSyrk() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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.
 *
 * Contract:
 * - Uses the same descriptor, offset, uplo, and block-size contract as
 *   \ref cublasMpSyrk_bufferSize.
 * - Only the triangle of C selected by uplo is updated.
 *
 * Workspace:
 * - Call cublasMpSyrk_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 workspace size required by \ref cublasMpSyr2k.
 *
 * This function returns the required buffer sizes to perform cublasMpSyr2k() on the
 * given input.
 *
 * Contract:
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL, and A and B must have matching data types.
 * - C must have square block size (mbC == nbC).
 * - The n-dimension blocking of op(A) and op(B) must match C: mbA == mbB == mbC
 *   when trans=N, nbA == nbB == mbC when trans=T.
 * - The k-dimension blocking of A and B must match: nbA == nbB when trans=N,
 *   mbA == mbB when trans=T.
 * - uplo selects the stored triangle of C; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpSyr2k() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpSyr2k_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* 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 symmetric rank-2k update.
 *
 * This function performs the symmetric rank-2k update:
 *   C = alpha * op(A) * op(B)^T + conj(alpha) * op(B) * op(A)^T + beta * C
 *
 * where alpha and beta are scalars, C is a symmetric matrix stored in lower or upper
 * mode, and A and B are matrices with dimensions op(A), op(B) n x k.
 * For complex types, this computes the Hermitian rank-2k update (HER2K).
 *
 * Contract:
 * - Uses the same descriptor, offset, uplo, data-type, and block-size contract as
 *   \ref cublasMpSyr2k_bufferSize.
 * - Only the triangle of C selected by uplo is updated.
 *
 * Workspace:
 * - Call cublasMpSyr2k_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpSyr2k_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpSyr2k_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSyr2k(
    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* 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 cublasMpSyrkx.
 *
 * This function returns the required buffer sizes to perform cublasMpSyrkx() on the
 * given input.
 *
 * Contract:
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL, and A and B must have matching data types.
 * - C must have square block size (mbC == nbC).
 * - The n-dimension blocking of op(A) and op(B) must match C: mbA == mbB == mbC
 *   when trans=N, nbA == nbB == mbC when trans=T.
 * - The k-dimension blocking of A and B must match: nbA == nbB when trans=N,
 *   mbA == mbB when trans=T.
 * - uplo selects the stored triangle of C; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpSyrkx() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpSyrkx_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* 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 extended symmetric rank-k update.
 *
 * This function performs the operation:
 *   C = alpha * op(A) * op(B)^T + beta * C
 *
 * where alpha and beta are scalars, C is a symmetric matrix stored in lower or upper
 * mode, and A and B are matrices with dimensions op(A), op(B) n x k.
 * Only the stored triangle of C is updated.
 * For complex types, this computes the extended Hermitian rank-k update (HERKX).
 *
 * Contract:
 * - Uses the same descriptor, offset, uplo, data-type, and block-size contract as
 *   \ref cublasMpSyrkx_bufferSize.
 * - Only the triangle of C selected by uplo is updated.
 *
 * Workspace:
 * - Call cublasMpSyrkx_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpSyrkx_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpSyrkx_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSyrkx(
    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* 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 cublasMpSymm.
 *
 * This function returns the required buffer sizes to perform cublasMpSymm() on the
 * given input.
 *
 * Contract:
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL.
 * - A is m x m when side == CUBLAS_SIDE_LEFT and n x n when side == CUBLAS_SIDE_RIGHT.
 * - A must have square block size (mbA == nbA).
 * - Row blocking must match across all matrices: mbA == mbB == mbC.
 * - Column blocking of B and C must match: nbB == nbC.
 * - uplo selects the stored symmetric part of A; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpSymm() call.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  side                       Indicates if A appears on the left or right of B.
 * \param[in]  uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]  m                          Number of rows of C.
 * \param[in]  n                          Number of columns of C.
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global symmetric matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpSymm_bufferSize(
    cublasMpHandle_t handle,
    cublasSideMode_t side,
    cublasFillMode_t uplo,
    int64_t m,
    int64_t n,
    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 symmetric matrix-matrix multiply.
 *
 * This function performs:
 *   C = alpha * A * B + beta * C  (side == CUBLAS_SIDE_LEFT)
 *   C = alpha * B * A + beta * C  (side == CUBLAS_SIDE_RIGHT)
 *
 * where alpha and beta are scalars, A is a symmetric matrix stored in lower or upper
 * mode, and B and C are m x n matrices.
 * For complex types, this computes the Hermitian matrix multiply (HEMM).
 *
 * Contract:
 * - Uses the same descriptor, offset, side, uplo, and block-size contract as
 *   \ref cublasMpSymm_bufferSize.
 * - C is updated in place.
 *
 * Workspace:
 * - Call cublasMpSymm_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     side                       Indicates if A appears on the left or right of B.
 * \param[in]     uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]     m                          Number of rows of C.
 * \param[in]     n                          Number of columns of C.
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global symmetric matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpSymm_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpSymm_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpSymm(
    cublasMpHandle_t handle,
    cublasSideMode_t side,
    cublasFillMode_t uplo,
    int64_t m,
    int64_t n,
    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 cublasMpHerk.
 *
 * This function returns the required buffer sizes to perform cublasMpHerk() on the
 * given input. This function supports complex types only (CUDA_C_32F, CUDA_C_64F)
 * and returns CUBLASMP_STATUS_NOT_SUPPORTED for real types.
 *
 * Contract:
 * - trans must be CUBLAS_OP_N or CUBLAS_OP_C.
 * - ia, ja, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA and descC must be non-NULL.
 * - C must have square block size (mbC == nbC).
 * - The n-dimension blocking of op(A) must match C: mbA == mbC when trans=N, nbA == mbC when trans=C.
 * - uplo selects the stored Hermitian triangle of C; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpHerk() call.
 *
 * \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 conjugate-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, real-valued).
 * \param[in]  a                          Pointer to local portion of global matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 must be a multiple of nbA.
 * \param[in]  descA                      Matrix descriptor for global matrix A.
 * \param[in]  beta                       Scalar used for multiplication (host pointer, real-valued).
 * \param[in]  c                          Pointer to local portion of global matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHerk_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 Hermitian rank-k update (complex types only).
 *
 * This function performs the Hermitian rank-k update:
 *   C = alpha * op(A) * op(A)^H + beta * C
 *
 * where alpha is a real scalar, beta is a real scalar, C is a Hermitian 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^H if trans == CUBLAS_OP_C.
 *
 * Contract:
 * - Uses the same descriptor, offset, uplo, trans, and block-size contract as
 *   \ref cublasMpHerk_bufferSize.
 * - Only the triangle of C selected by uplo is updated.
 *
 * Workspace:
 * - Call cublasMpHerk_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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 conjugate-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, real-valued).
 * \param[in]     a                          Pointer to local portion of global matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 must be a multiple of nbA.
 * \param[in]     descA                      Matrix descriptor for global matrix A.
 * \param[in]     beta                       Scalar used for multiplication (host pointer, real-valued).
 * \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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHerk_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpHerk_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpHerk(
    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 workspace size required by \ref cublasMpHer2k.
 *
 * This function returns the required buffer sizes to perform cublasMpHer2k() on the
 * given input. This function supports complex types only (CUDA_C_32F, CUDA_C_64F)
 * and returns CUBLASMP_STATUS_NOT_SUPPORTED for real types.
 *
 * Contract:
 * - trans must be CUBLAS_OP_N or CUBLAS_OP_C.
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL, and A and B must have matching data types.
 * - C must have square block size (mbC == nbC).
 * - The n-dimension blocking of op(A) and op(B) must match C: mbA == mbB == mbC
 *   when trans=N, nbA == nbB == mbC when trans=C.
 * - The k-dimension blocking of A and B must match: nbA == nbB when trans=N,
 *   mbA == mbB when trans=C.
 * - uplo selects the stored Hermitian triangle of C; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpHer2k() call.
 *
 * \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 conjugate-transpose.
 * \param[in]  n                          Number of rows of op(A), op(B), and C.
 * \param[in]  k                          Number of columns of op(A) and op(B).
 * \param[in]  alpha                      Complex 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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  beta                       Real 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHer2k_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* 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 Hermitian rank-2k update (complex types only).
 *
 * This function performs the Hermitian rank-2k update:
 *   C = alpha * op(A) * op(B)^H + conj(alpha) * op(B) * op(A)^H + beta * C
 *
 * where alpha is a complex scalar, beta is a real scalar, C is a Hermitian matrix stored
 * in lower or upper mode, and A, B are matrices with dimensions op(A), op(B) n x k.
 *
 * Contract:
 * - Uses the same descriptor, offset, uplo, trans, data-type, and block-size contract as
 *   \ref cublasMpHer2k_bufferSize.
 * - Only the triangle of C selected by uplo is updated.
 *
 * Workspace:
 * - Call cublasMpHer2k_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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 conjugate-transpose.
 * \param[in]     n                          Number of rows of op(A), op(B), and C.
 * \param[in]     k                          Number of columns of op(A) and op(B).
 * \param[in]     alpha                      Complex 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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]     descB                      Matrix descriptor for global matrix B.
 * \param[in]     beta                       Real 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHer2k_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpHer2k_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpHer2k(
    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* 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 cublasMpHerkx.
 *
 * This function returns the required buffer sizes to perform cublasMpHerkx() on the
 * given input. This function supports complex types only (CUDA_C_32F, CUDA_C_64F)
 * and returns CUBLASMP_STATUS_NOT_SUPPORTED for real types.
 *
 * Contract:
 * - trans must be CUBLAS_OP_N or CUBLAS_OP_C.
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL, and A and B must have matching data types.
 * - C must have square block size (mbC == nbC).
 * - The n-dimension blocking of op(A) and op(B) must match C: mbA == mbB == mbC
 *   when trans=N, nbA == nbB == mbC when trans=C.
 * - The k-dimension blocking of A and B must match: nbA == nbB when trans=N,
 *   mbA == mbB when trans=C.
 * - uplo selects the stored Hermitian triangle of C; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpHerkx() call.
 *
 * \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 conjugate-transpose.
 * \param[in]  n                          Number of rows of op(A), op(B), and C.
 * \param[in]  k                          Number of columns of op(A) and op(B).
 * \param[in]  alpha                      Complex 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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[in]  beta                       Real 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHerkx_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* 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 extended Hermitian rank-k update (complex types only).
 *
 * This function performs the operation:
 *   C = alpha * op(A) * op(B)^H + beta * C
 *
 * where alpha is a complex scalar, beta is a real scalar, C is a Hermitian matrix stored
 * in lower or upper mode, and A, B are matrices with dimensions op(A), op(B) n x k.
 * Only the stored triangle of C is updated.
 *
 * Contract:
 * - Uses the same descriptor, offset, uplo, trans, data-type, and block-size contract as
 *   \ref cublasMpHerkx_bufferSize.
 * - Only the triangle of C selected by uplo is updated.
 *
 * Workspace:
 * - Call cublasMpHerkx_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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 conjugate-transpose.
 * \param[in]     n                          Number of rows of op(A), op(B), and C.
 * \param[in]     k                          Number of columns of op(A) and op(B).
 * \param[in]     alpha                      Complex 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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]     descB                      Matrix descriptor for global matrix B.
 * \param[in]     beta                       Real 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHerkx_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpHerkx_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpHerkx(
    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* 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 cublasMpHemm.
 *
 * This function returns the required buffer sizes to perform cublasMpHemm() on the
 * given input. This function supports complex types only (CUDA_C_32F, CUDA_C_64F)
 * and returns CUBLASMP_STATUS_NOT_SUPPORTED for real types.
 *
 * Contract:
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL.
 * - A is m x m when side == CUBLAS_SIDE_LEFT and n x n when side == CUBLAS_SIDE_RIGHT.
 * - A must have square block size (mbA == nbA).
 * - Row blocking must match across all matrices: mbA == mbB == mbC.
 * - Column blocking of B and C must match: nbB == nbC.
 * - uplo selects the stored Hermitian triangle of A; the other triangle is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpHemm() call.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  side                       Indicates if the Hermitian matrix A appears on the left or right of B.
 * \param[in]  uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]  m                          Number of rows of B and C.
 * \param[in]  n                          Number of columns of B and C.
 * \param[in]  alpha                      Complex scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global Hermitian matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHemm_bufferSize(
    cublasMpHandle_t handle,
    cublasSideMode_t side,
    cublasFillMode_t uplo,
    int64_t m,
    int64_t n,
    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 Hermitian matrix-matrix multiply (complex types only).
 *
 * This function performs:
 *   C = alpha * A * B + beta * C  (side == CUBLAS_SIDE_LEFT)
 *   C = alpha * B * A + beta * C  (side == CUBLAS_SIDE_RIGHT)
 *
 * where alpha and beta are complex scalars, A is a Hermitian matrix (A = A^H) stored in
 * lower or upper mode, and B, C are m x n matrices. A is m x m when side == LEFT and
 * n x n when side == RIGHT.
 *
 * Contract:
 * - Uses the same descriptor, offset, side, uplo, and block-size contract as
 *   \ref cublasMpHemm_bufferSize.
 * - C is updated in place.
 *
 * Workspace:
 * - Call cublasMpHemm_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     side                       Indicates if the Hermitian matrix A appears on the left or right of B.
 * \param[in]     uplo                       Indicates if matrix A lower or upper part is stored.
 * \param[in]     m                          Number of rows of C.
 * \param[in]     n                          Number of columns of C.
 * \param[in]     alpha                      Complex scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global Hermitian matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpHemm_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpHemm_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpHemm(
    cublasMpHandle_t handle,
    cublasSideMode_t side,
    cublasFillMode_t uplo,
    int64_t m,
    int64_t n,
    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 cublasMpTrmm.
 *
 * This function returns the required buffer sizes to perform cublasMpTrmm() on the
 * given input.
 *
 * Contract:
 * - ia, ja, ib, jb, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - descA, descB, and descC must be non-NULL.
 * - A is m x m when side == CUBLAS_SIDE_LEFT and n x n when side == CUBLAS_SIDE_RIGHT.
 * - A must have square block size (mbA == nbA).
 * - Row blocking must match across all matrices: mbA == mbB == mbC.
 * - Column blocking of B and C must match: nbB == nbC.
 * - uplo selects the stored triangular part of A. If diag == CUBLAS_DIAG_UNIT, the
 *   diagonal elements of A are assumed to be one and are not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpTrmm() call.
 *
 * \param[in]  handle                     cuBLASMp library handle.
 * \param[in]  side                       Indicates if A appears on the left or right of B.
 * \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 diagonal elements of A are unity.
 * \param[in]  m                          Number of rows of B.
 * \param[in]  n                          Number of columns of B.
 * \param[in]  alpha                      Scalar used for multiplication (host pointer).
 * \param[in]  a                          Pointer to local portion of global triangular matrix A (device).
 * \param[in]  ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]  jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]  descB                      Matrix descriptor for global matrix B.
 * \param[out] c                          Pointer to local portion of global output matrix C (device).
 * \param[in]  ic                         Row index of the first row of sub(C). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpTrmm_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,
    const void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    void* c,
    int64_t ic,
    int64_t jc,
    cublasMpMatrixDescriptor_t descC,
    cublasComputeType_t computeType,
    size_t* workspaceSizeInBytesOnDevice,
    size_t* workspaceSizeInBytesOnHost);

/**
 * \brief Performs a distributed triangular matrix-matrix multiply.
 *
 * This function performs:
 *   C = alpha * op(A) * B  (side == CUBLAS_SIDE_LEFT)
 *   C = alpha * B * op(A)  (side == CUBLAS_SIDE_RIGHT)
 *
 * where alpha is a scalar, A is a triangular matrix stored in lower or upper mode,
 * and B and C are m x n matrices.
 *
 * Contract:
 * - Uses the same descriptor, offset, side, uplo, diag, and block-size contract as
 *   \ref cublasMpTrmm_bufferSize.
 * - Matrix C receives the result.
 *
 * Workspace:
 * - Call cublasMpTrmm_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \param[in]     handle                     cuBLASMp library handle.
 * \param[in]     side                       Indicates if A appears on the left or right of B.
 * \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 diagonal elements of A are unity.
 * \param[in]     m                          Number of rows of B.
 * \param[in]     n                          Number of columns of B.
 * \param[in]     alpha                      Scalar used for multiplication (host pointer).
 * \param[in]     a                          Pointer to local portion of global triangular matrix A (device).
 * \param[in]     ia                         Row index of the first row of sub(A). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ib - 1 must be a multiple of mbB.
 * \param[in]     jb                         Column index of the first column of sub(B). jb - 1 must be a multiple of nbB.
 * \param[in]     descB                      Matrix descriptor for global matrix B.
 * \param[out]    c                          Pointer to local portion of global matrix C (device, output).
 * \param[in]     ic                         Row index of the first row of sub(C). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 cublasMpTrmm_bufferSize().
 * \param[out]    h_work                     Host workspace of size workspaceSizeInBytesOnHost.
 * \param[in]     workspaceSizeInBytesOnHost    Host workspace size as provided by cublasMpTrmm_bufferSize().
 * \return cublasMpStatus_t
 */
cublasMpStatus_t cublasMpTrmm(
    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,
    const void* b,
    int64_t ib,
    int64_t jb,
    cublasMpMatrixDescriptor_t descB,
    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. Process coordinates are zero-based.
 * Parameters nb and nprocs must be positive.
 *
 * \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.
 *                      Must be in [0, nprocs).
 * \param[in] isrcproc  Coordinate of the process that owns the first row or column of the distributed matrix.
 *                      Must be in [0, nprocs).
 * \param[in] nprocs    Total number of row or column processes over which the matrix is distributed.
 *                      Must be positive.
 * \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.
 *
 * Contract:
 * - All ranks in global_comm must call this function, including ranks that do not
 *   participate in either matrix grid.
 * - ia, ja, ib, and jb are 1-based submatrix offsets. The requested submatrices must
 *   fit inside descA and descB.
 * - descA and descB must be non-NULL and have matching data types. A descriptor grid
 *   may be NULL only on ranks outside that matrix's process grid.
 * - global_comm must contain at least the union of the A and B process grids.
 * - For FP4 data, m, the row block sizes, and local leading dimensions must be even;
 *   ia and ib must be odd 1-based offsets.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpGemr2D() call.
 *
 * \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 part of either matrix.
 *
 * Contract:
 * - Uses the same descriptor, offset, communicator, data-type, and FP4 contract as
 *   \ref cublasMpGemr2D_bufferSize.
 *
 * Workspace:
 * - Call cublasMpGemr2D_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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.
 *
 * Contract:
 * - All ranks in global_comm must call this function, including ranks that do not
 *   participate in either matrix grid.
 * - ia, ja, ib, and jb are 1-based submatrix offsets. The requested submatrices must
 *   fit inside descA and descB.
 * - descA and descB must be non-NULL and have matching data types. A descriptor grid
 *   may be NULL only on ranks outside that matrix's process grid.
 * - global_comm must contain at least the union of the A and B process grids.
 * - uplo selects the trapezoidal part to redistribute. If diag == CUBLAS_DIAG_UNIT,
 *   diagonal elements are treated as one.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpTrmr2D() call.
 *
 * \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 part of either matrix.
 *
 * Contract:
 * - Uses the same descriptor, offset, communicator, data-type, uplo, and diag contract
 *   as \ref cublasMpTrmr2D_bufferSize.
 *
 * Workspace:
 * - Call cublasMpTrmr2D_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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. A and C block sizes must be compatible with trans: for CUBLAS_OP_N,
 * mbA == mbC and nbA == nbC; otherwise, nbA == mbC and mbA == nbC.
 *
 * Contract:
 * - ia, ja, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - handle, alpha, beta, descA, and descC must be non-NULL.
 * - For in-place use where A and C describe the same memory and offsets, only
 *   trans == CUBLAS_OP_N is supported.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpGeadd() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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.
 *
 * A and C block sizes must be compatible with trans: for CUBLAS_OP_N,
 * mbA == mbC and nbA == nbC; otherwise, nbA == mbC and mbA == nbC.
 *
 * Contract:
 * - ia, ja, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - handle, alpha, beta, descA, and descC must be non-NULL.
 * - For in-place use where A and C describe the same memory and offsets, only
 *   trans == CUBLAS_OP_N is supported.
 *
 * Workspace:
 * - Call cublasMpGeadd_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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. A and C block sizes must be compatible with trans: for CUBLAS_OP_N,
 * mbA == mbC and nbA == nbC; otherwise, nbA == mbC and mbA == nbC.
 *
 * Contract:
 * - ia, ja, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - handle, alpha, beta, descA, and descC must be non-NULL.
 * - uplo must select the stored trapezoidal part of C; the other part is not referenced.
 *
 * Workspace:
 * - Use the returned workspace sizes unchanged for the matching cublasMpTradd() call.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]  ja                         Column index of the first column of sub(A). ja - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]  jc                         Column index of the first column of sub(C). jc - 1 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.
 *
 * A and C block sizes must be compatible with trans: for CUBLAS_OP_N,
 * mbA == mbC and nbA == nbC; otherwise, nbA == mbC and mbA == nbC.
 *
 * Contract:
 * - ia, ja, ic, and jc are 1-based and must be aligned to the corresponding
 *   descriptor block dimensions.
 * - handle, alpha, beta, descA, and descC must be non-NULL.
 * - uplo must select the stored trapezoidal part of C; the other part is not referenced.
 *
 * Workspace:
 * - Call cublasMpTradd_bufferSize() with the same arguments first, then pass device
 *   and host workspaces of the returned sizes.
 *
 * \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). ia - 1 must be a multiple of mbA.
 * \param[in]     ja                         Column index of the first column of sub(A). ja - 1 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). ic - 1 must be a multiple of mbC.
 * \param[in]     jc                         Column index of the first column of sub(C). jc - 1 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 combines ncclMemAlloc and \ref cublasMpBufferRegister
 * into a single call. See \ref cublasMpBufferRegister for details on how workspace registration
 * affects algorithm selection and optional matrix B registration.
 *
 * All ranks in the grid communicator must enter this function collectively.
 * 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 ranks in the grid communicator must enter this function collectively.
 *
 * \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
