API Reference - cuTENSORMp

General

---

cutensorMpHandle_t

typedef struct cuTensorMpHandle *cutensorMpHandle_t

---

cutensorMpCreate()

cutensorStatus_t cutensorMpCreate(

cutensorMpHandle_t *handle,

ncclComm_t comm,

int local_device_id,

cudaStream_t stream,

)

Initializes the cutensorMp library and creates a handle for distributed tensor operations.

This function creates a cutensorMp handle that serves as the context for all distributed tensor operations. The handle is associated with a specific MPI communicator, local CUDA device, and CUDA stream. This allows cutensorMp to coordinate tensor operations across multiple processes and GPUs.

The communicator defines the group of processes that will participate in distributed tensor operations. The local device ID specifies which CUDA device on the current process will be used for computations. The CUDA stream enables asynchronous execution and synchronization with other CUDA operations.

The user is responsible for calling cutensorMpDestroy to free the resources associated with the handle.

Remark

non-blocking, no reentrant, and thread-safe

Parameters:

• handle – [out] Pointer to cutensorMpHandle_t that will hold the created handle

• comm – [in] NCCL communicator that defines the group of processes for distributed operations

• local_device_id – [in] CUDA device ID to use on the current process (must be valid and accessible)

• stream – [in] CUDA stream for asynchronous operations; the default stream (0) is accepted

Return values:

CUTENSOR_STATUS_SUCCESS – on success and an error code otherwise

---

cutensorMpDestroy()

cutensorStatus_t cutensorMpDestroy(cutensorMpHandle_t handle)

Frees all resources associated with the provided cutensorMp handle.

This function deallocates all memory and resources associated with a cutensorMp handle that was previously created by cutensorMpCreate. After calling this function, the handle becomes invalid and should not be used in subsequent cutensorMp operations.

Remark

blocking, no reentrant, and thread-safe

Parameters:

handle – [inout] The cutensorMpHandle_t object that will be deallocated

Return values:

CUTENSOR_STATUS_SUCCESS – on success and an error code otherwise

---

Tensor Descriptors

---

cutensorMpTensorDescriptor_t

typedef struct cuTensorMpTensorDescriptor *cutensorMpTensorDescriptor_t

---

cutensorMpCreateTensorDescriptor()

cutensorStatus_t cutensorMpCreateTensorDescriptor(

cutensorMpHandle_t const handle,

cutensorMpTensorDescriptor_t *desc,

uint32_t const numModes,

int64_t const extent[],

int64_t const elementStride[],

int64_t const blockSize[],

int64_t const blockStride[],

int64_t const nranksPerMode[],

uint32_t const nranks,

int32_t const ranks[],

cudaDataType_t const dataType,

)

Creates a distributed tensor descriptor for multi-process tensor operations.

This function creates a tensor descriptor that defines the structure and distribution of a multi-dimensional tensor across multiple processes. Unlike regular cuTENSOR tensor descriptors, this descriptor includes information about how the tensor is partitioned and distributed across different processes in the MPI communicator.

The tensor is described by its modes (dimensions), extents (sizes along each mode), and strides for elements and blocks. Rank ownership is specified through block sizes, nranks-per-mode, and the rank mapping. Element and block strides describe the local physical layout and do not change which rank owns an element.

Let P = product(nranksPerMode[0:numModes]). With this API’s default logical-rank strides, when nranks is an integer multiple of P, the descriptor has a uniform replication factor of nranks / P. A factor of 1 means every participating rank has a distinct ownership coordinate; a factor greater than 1 means each coordinate is stored by that many ranks. Otherwise replica multiplicity may be non-uniform and no single integer replication factor describes the descriptor.

Rank ownership follows the order of ranks: ranks[l] has logical rank l. If ranks is null, communicator ranks in [0, nranks) are used. Logical rank coordinates use mode 0 as the fastest-varying dimension. A communicator rank not listed in ranks owns no local elements.

Each mode is mapped independently. nranksPerMode[i] == 1 replicates the mode. A null blockSize array or blockSize[i] == 0 creates near-even contiguous slices; lower rank coordinates receive any remainder. A positive blockSize[i] assigns fixed-size global blocks round-robin, and the final global block may be shorter.

For example, extent 10 with block size 4 on two rank coordinates assigns global elements [0, 4) and [8, 10) to coordinate 0, and [4, 8) to coordinate 1.

When blockStride is null, no extra padding is inserted between consecutive locally owned blocks: the realized block stride is the local block size multiplied by the realized element stride. A short final block is not padded. When elementStride is also null, each rank’s full local shard is compact, with mode 0 varying fastest. Use cutensorMpTensorDescriptorGetAttribute to query exact rank-local extents, strides, and storage size.

The user is responsible for calling cutensorMpDestroyTensorDescriptor to free the resources associated with the descriptor once it is no longer needed.

Remark

non-blocking, no reentrant, and thread-safe

Parameters:

• handle – [in] Opaque handle holding cutensorMp’s library context

• desc – [out] Pointer to the address where the allocated tensor descriptor object will be stored

• numModes – [in] Number of modes (dimensions) in the tensor (must be greater than zero)

• extent – [in] Extent (size) of each mode (size: numModes, all values must be greater than zero)

• elementStride – [in] Stride between consecutive elements in each mode (size: numModes), in elements. Passing null selects compact rank-local element strides.

• blockSize – [in] Size of each block along each mode for distribution (size: numModes). Passing null uses one contiguous partition per rank along each distributed mode. Passing a non-null array uses each positive entry as that mode’s block-cyclic block size; an entry of 0 has the same effect as the null default for that mode, so contiguous and block-cyclic modes can be mixed in one descriptor.

• blockStride – [in] Optional per-mode stride, in elements, between consecutive blocks stored in the local buffer (size: numModes). Let B_i be the effective block size selected by blockSize and s_i the resolved per-element stride. For local coordinate l in mode i, the physical offset contribution is (l / B_i) * blockStride[i] + (l % B_i) * s_i. This parameter controls local physical storage only; rank ownership is determined by blockSize, nranksPerMode, and the rank mapping. Passing null selects ordinary linear local addressing without inter-block padding; when supplied, every entry must be positive. Operations that consume explicit block strides require a non-overlapping physical layout, and support is operation-specific.

• nranksPerMode – [in] Number of processes along each mode (size: numModes)

• nranks – [in] Total number of ranks (processes) participating in the tensor distribution

• ranks – [in] Array of rank IDs for each participating process (size: nranks), passing null will use the range [0, nranks)

• dataType – [in] Data type of the tensor elements

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully

• CUTENSOR_STATUS_NOT_INITIALIZED – if the handle is not initialized

• CUTENSOR_STATUS_INVALID_VALUE – if some input data is invalid (this typically indicates a user error)

• CUTENSOR_STATUS_NOT_SUPPORTED – if the requested descriptor configuration is not supported

---

cutensorMpDestroyTensorDescriptor()

cutensorStatus_t cutensorMpDestroyTensorDescriptor(

cutensorMpTensorDescriptor_t desc,

)

Frees all resources related to the provided distributed tensor descriptor.

This function deallocates all memory and resources associated with a cutensorMp tensor descriptor that was previously created by cutensorMpCreateTensorDescriptor. After calling this function, the descriptor becomes invalid and should not be used in subsequent cutensorMp operations.

Remark

blocking, no reentrant, and thread-safe

Parameters:

desc – [inout] The cutensorMpTensorDescriptor_t object that will be deallocated

Return values:

CUTENSOR_STATUS_SUCCESS – on success and an error code otherwise

Generic Operation Functions

The following functions are generic and work with all the different operations.

---

cutensorMpDestroyOperationDescriptor()

cutensorStatus_t cutensorMpDestroyOperationDescriptor(

cutensorMpOperationDescriptor_t desc,

)

Frees all resources related to the provided distributed contraction descriptor.

This function deallocates all memory and resources associated with a cutensorMp operation descriptor that was previously created by cutensorMpCreateContraction. After calling this function, the descriptor becomes invalid and should not be used in subsequent cutensorMp operations.

Remark

blocking, no reentrant, and thread-safe

Parameters:

desc – [inout] The cutensorMpOperationDescriptor_t object that will be deallocated

Return values:

CUTENSOR_STATUS_SUCCESS – on success and an error code otherwise

---

Plan Preferences

---

cutensorMpAlgo_t

enum cutensorMpAlgo_t

Values:

enumerator CUTENSORMP_ALGO_DEFAULT

Default algorithm.

---

cutensorMpPlanPreference_t

typedef struct cuTensorMpPlanPreference *cutensorMpPlanPreference_t

---

cutensorMpCreatePlanPreference()

cutensorStatus_t cutensorMpCreatePlanPreference(

cutensorMpHandle_t const handle,

cutensorMpPlanPreference_t *pref,

cutensorMpAlgo_t const cutensormp_algo,

uint64_t const cutensormp_workspace_size_device,

uint64_t const cutensormp_workspace_size_host,

)

Creates a plan preference object for controlling distributed tensor operation planning.

This function creates a preference object that allows users to control various aspects of the execution plan for distributed tensor operations. The preferences include algorithm selection, workspace size limits, and JIT compilation options that affect both the underlying cuTENSOR operations and the distributed communication patterns.

The plan preference provides fine-grained control over:

• Local cuTENSOR algorithm selection and JIT mode

• Distributed algorithm strategy (non-packing, packing with permutation, or packing with P2P)

• Workspace size limits for both device and host memory

• cuTENSOR workspace preferences

The user is responsible for calling cutensorMpDestroyPlanPreference to free the resources associated with the preference object.

Remark

non-blocking, no reentrant, and thread-safe

Parameters:

• handle – [in] Opaque handle holding cutensorMp’s library context

• pref – [out] Pointer to the plan preference object that will be created

• cutensormp_algo – [in] Algorithm selection for distributed communication patterns

• cutensormp_workspace_size_device – [in] Maximum device workspace size for cutensorMp operations (bytes)

• cutensormp_workspace_size_host – [in] Maximum host workspace size for cutensorMp operations (bytes)

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully

• CUTENSOR_STATUS_NOT_INITIALIZED – if the handle is not initialized

• CUTENSOR_STATUS_INVALID_VALUE – if some input data is invalid (this typically indicates a user error)

---

cutensorMpDestroyPlanPreference()

cutensorStatus_t cutensorMpDestroyPlanPreference(

cutensorMpPlanPreference_t pref,

)

Frees all resources related to the provided plan preference object.

This function deallocates all memory and resources associated with a cutensorMp plan preference object that was previously created by cutensorMpCreatePlanPreference. After calling this function, the preference object becomes invalid and should not be used in subsequent cutensorMp operations.

Remark

blocking, no reentrant, and thread-safe

Parameters:

pref – [inout] The cutensorMpPlanPreference_t object that will be deallocated

Return values:

CUTENSOR_STATUS_SUCCESS – on success and an error code otherwise

---

Plans

---

cutensorMpPlan_t

typedef struct cuTensorMpPlan *cutensorMpPlan_t

---

cutensorMpPlanAttribute_t

enum cutensorMpPlanAttribute_t

Values:

enumerator CUTENSORMP_PLAN_REQUIRED_WORKSPACE_DEVICE

uint64_t: exact required workspace in bytes that is needed to execute the plan

enumerator CUTENSORMP_PLAN_REQUIRED_WORKSPACE_HOST

uint64_t: exact required workspace in bytes that is needed to execute the plan

---

cutensorMpCreatePlan()

cutensorStatus_t cutensorMpCreatePlan(

cutensorMpHandle_t const handle,

cutensorMpPlan_t *plan,

cutensorMpOperationDescriptor_t const desc,

cutensorMpPlanPreference_t const pref,

)

Creates an execution plan for a distributed tensor operation.

This function creates an optimized execution plan for the distributed tensor operation encoded by the operation descriptor. The plan selects the most appropriate algorithms and communication strategies based on the tensor distributions, available resources, and user preferences.

The planning process analyzes the distributed tensor layout, communication requirements, and computational resources to determine an efficient execution strategy. This may involve data redistribution, local contractions, and result aggregation phases that minimize communication overhead while maximizing computational efficiency.

Planning is a collective call on the handle communicator after a valid handle is provided. All ranks must pass matching operation descriptors and algorithm selection fields; workspace budgets are treated as local resource limits and may differ by rank.

The user is responsible for calling cutensorMpDestroyPlan to free the resources associated with the plan once it is no longer needed.

Remark

calls asynchronous functions, no reentrant, and thread-safe

Parameters:

• handle – [in] Opaque handle holding cutensorMp’s library context

• plan – [out] Pointer to the execution plan object that will be created

• desc – [in] Operation descriptor encoding the distributed operation (created by cutensorMpCreateContraction or cutensorMpCreateRedistribution)

• pref – [in] Plan preference object specifying algorithm and workspace preferences (may be NULL for default preferences)

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully

• CUTENSOR_STATUS_NOT_INITIALIZED – if the handle is not initialized

• CUTENSOR_STATUS_INVALID_VALUE – if some input data is invalid (this typically indicates a user error)

• CUTENSOR_STATUS_NOT_SUPPORTED – if no viable execution plan could be found

---

cutensorMpDestroyPlan()

cutensorStatus_t cutensorMpDestroyPlan(cutensorMpPlan_t plan)

Frees all resources related to the provided distributed contraction plan.

This function deallocates all memory and resources associated with a cutensorMp execution plan that was previously created by cutensorMpCreatePlan. After calling this function, the plan becomes invalid and should not be used in subsequent cutensorMp operations.

Remark

blocking, no reentrant, and thread-safe

Parameters:

plan – [inout] The cutensorMpPlan_t object that will be deallocated

Return values:

CUTENSOR_STATUS_SUCCESS – on success and an error code otherwise

---

cutensorMpSetStream()

cutensorStatus_t cutensorMpSetStream(

cutensorMpHandle_t const handle,

cutensorMpPlan_t const plan,

cudaStream_t const stream,

)

Overrides the CUDA stream that the given plan executes on.

Every plan starts out bound to the stream that was passed to cutensorMpCreate. This function retargets a single plan onto a different stream, so that plans created from the same handle can be issued on separate streams without creating one handle per stream. All subsequent executions of plan (e.g., cutensorMpContract, cutensorMpRedistribute, cutensorMpUnfold, cutensorMpFold) are enqueued on stream, and the caller must synchronize stream — not the handle’s stream — to observe their results.

The work that plan creation and any earlier execution already enqueued on the plan’s previous stream is ordered before stream, so no additional synchronization is required at the switch.

The stream is a rank-local property: this function performs no communication and ranks may pass different streams. It must not be called while an execution of plan is still in flight.

Remark

non-blocking, no reentrant, and thread-safe

Parameters:

• handle – [in] Opaque handle holding cutensorMp’s library context

• plan – [in] Execution plan to retarget; must have been created with handle

• stream – [in] The CUDA stream that subsequent executions of plan will use; the default stream (0) is accepted

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully

• CUTENSOR_STATUS_NOT_INITIALIZED – if the handle is not initialized

• CUTENSOR_STATUS_INVALID_VALUE – if handle or plan is NULL, or if the plan was not created with the provided handle

---

cutensorMpPlanGetAttribute()

cutensorStatus_t cutensorMpPlanGetAttribute(

cutensorMpHandle_t const handle,

cutensorMpPlan_t const plan,

cutensorMpPlanAttribute_t const attribute,

void *buf,

size_t const sizeInBytes,

)

Retrieves information about an already-created plan (see cutensorPlanAttribute_t)

Parameters:

• plan – [in] Denotes an already-created plan (e.g., via cutensorMpCreatePlan)

• attr – [in] Requested attribute.

• buf – [out] On successful exit: Holds the information of the requested attribute.

• sizeInBytes – [in] size of buf in bytes.

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully.

• CUTENSOR_STATUS_INVALID_VALUE – if some input data is invalid (this typically indicates an user error).

---

Contraction Operations

---

cutensorMpOperationDescriptor_t

typedef struct cuTensorMpOperationDescriptor *cutensorMpOperationDescriptor_t

---

cutensorMpCreateContraction()

cutensorStatus_t cutensorMpCreateContraction(

cutensorMpHandle_t const handle,

cutensorMpOperationDescriptor_t *desc,

cutensorMpTensorDescriptor_t const descA,

int32_t const modesA[],

cutensorOperator_t opA,

cutensorMpTensorDescriptor_t const descB,

int32_t const modesB[],

cutensorOperator_t opB,

cutensorMpTensorDescriptor_t const descC,

int32_t const modesC[],

cutensorOperator_t opC,

cutensorMpTensorDescriptor_t const descD,

int32_t const modesD[],

cutensorComputeDescriptor_t const descCompute,

)

Creates an operation descriptor that encodes a distributed tensor contraction.

This function creates an operation descriptor for distributed tensor contractions of the form \( D = \alpha \mathcal{A} \mathcal{B} + \beta \mathcal{C} \), where the tensors A, B, C, and D are distributed across multiple processes as specified by their respective tensor descriptors.

The distributed contraction leverages both intra-process cuTENSOR operations and inter-process communication to efficiently compute tensor contractions that exceed the memory capacity or computational resources of a single GPU. The operation automatically handles data redistribution, local contractions, and result aggregation across the participating processes.

Replicated input tensors (descA, descB, descC with replication factor > 1, see cutensorMpCreateTensorDescriptor) are accepted and may unlock additional planner optimizations that reduce communication. Replicated output tensors (descD with replication factor > 1) are accepted for planner-supported ownership combinations. cutensorMpCreateContraction or cutensorMpCreatePlan returns CUTENSOR_STATUS_NOT_SUPPORTED when cuTensorMp cannot prove a correct producer, reduction, or redistribution path for every output replica.

The user is responsible for calling cutensorMpDestroyOperationDescriptor to free the resources associated with the descriptor once it is no longer needed.

Remark

non-blocking, no reentrant, and thread-safe

Parameters:

• handle – [in] Opaque handle holding cutensorMp’s library context

• desc – [out] Pointer to the operation descriptor that will be created and filled with information encoding the distributed contraction operation

• descA – [in] Distributed tensor descriptor for input tensor A

• modesA – [in] Modes of the input tensor A

• opA – [in] Unary operator that will be applied to each element of A before it is further processed. The original data of this tensor remains unchanged.

• descB – [in] Distributed tensor descriptor for input tensor B

• modesB – [in] Modes of the input tensor B

• opB – [in] Unary operator that will be applied to each element of B before it is further processed. The original data of this tensor remains unchanged.

• descC – [in] Distributed tensor descriptor for input tensor C

• modesC – [in] Modes of the input tensor C

• opC – [in] Unary operator that will be applied to each element of C before it is further processed. The original data of this tensor remains unchanged.

• descD – [in] Distributed tensor descriptor for output tensor D (currently must be identical to descC)

• modesD – [in] Modes of the output tensor D

• descCompute – [in] Compute descriptor that determines the precision for the operation

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully

• CUTENSOR_STATUS_NOT_INITIALIZED – if the handle is not initialized

• CUTENSOR_STATUS_INVALID_VALUE – if some input data is invalid

• CUTENSOR_STATUS_NOT_SUPPORTED – if the combination of tensor configurations is not supported

---

cutensorMpContract()

cutensorStatus_t cutensorMpContract(

cutensorMpHandle_t const handle,

cutensorMpPlan_t const plan,

void const *alpha,

void const *A,

void const *B,

void const *beta,

void const *C,

void *D,

void *device_workspace,

void *host_workspace,

)

Performs a distributed tensor contraction across multiple processes.

This function executes the distributed tensor contraction \( D = \alpha \mathcal{A} \mathcal{B} + \beta \mathcal{C} \) according to the execution plan created by cutensorMpCreatePlan. The operation coordinates computation and communication across multiple processes and GPUs to efficiently perform tensor contractions that exceed the capacity of a single device.

The execution involves several phases:

1. Data redistribution to align tensor blocks for efficient computation

2. Local tensor contractions using cuTENSOR on each participating device

3. Communication and aggregation of partial results across processes

4. Final result assembly in the distributed output tensor

All participating processes in the MPI communicator must call this function with consistent parameters. The input and output tensors must be distributed according to their respective tensor descriptors, with each process providing its local portion of the data.

Remark

calls asynchronous functions, no reentrant, and thread-safe

Parameters:

• handle – [in] Opaque handle holding cutensorMp’s library context

• plan – [in] Execution plan for the distributed contraction (created by cutensorMpCreatePlan)

• alpha – [in] Scaling factor for the A*B product. Pointer to host memory with data type determined by the compute descriptor. The data type follows that of cuTENSOR (i.e., the data type of the scalar is deptermined by the data type of C:CUDA_R_16F and CUDA_R_16BF use CUDA_R_32Fscalars, all data types of the scalar are identical to the type of C)

• A – [in] Pointer to the local portion of distributed tensor A in GPU memory; may be NULL when this rank’s local A extent has zero volume

• B – [in] Pointer to the local portion of distributed tensor B in GPU memory; may be NULL when this rank’s local B extent has zero volume

• beta – [in] Scaling factor for tensor C. Pointer to host memory with data type determined by the compute descriptor. The data type follows that of cuTENSOR (i.e., the data type of the scalar is deptermined by the data type of C:CUDA_R_16F and CUDA_R_16BF use CUDA_R_32Fscalars, all data types of the scalar are identical to the type of C)

• C – [in] Pointer to the local portion of distributed tensor C in GPU memory; may be NULL when this rank’s local C extent has zero volume

• D – [out] Pointer to the local portion of distributed tensor D in GPU memory (may be identical to C); may be NULL when this rank’s local D extent has zero volume

• device_workspace – [in] Pointer to device workspace memory (size determined by cutensorMpPlanGetAttribute with CUTENSORMP_PLAN_REQUIRED_WORKSPACE_DEVICE); may be NULL when the required device workspace size is zero on this rank

• host_workspace – [in] Pointer to host workspace memory (size determined by cutensorMpPlanGetAttribute with CUTENSORMP_PLAN_REQUIRED_WORKSPACE_HOST)

Return values:

• CUTENSOR_STATUS_SUCCESS – The operation completed successfully

• CUTENSOR_STATUS_NOT_INITIALIZED – if the handle is not initialized

• CUTENSOR_STATUS_INVALID_VALUE – if some input data is invalid (this typically indicates a user error)

• CUTENSOR_STATUS_NOT_SUPPORTED – if the operation is not supported with the given configuration

• CUTENSOR_STATUS_INSUFFICIENT_WORKSPACE – if the provided workspace is insufficient

• CUTENSOR_STATUS_ARCH_MISMATCH – if the plan was created for a different device architecture

• CUTENSOR_STATUS_CUDA_ERROR – if a CUDA error occurred during execution

---