Input / Output Format

Value Format

BLAS::a_value_type
BLAS::b_value_type
BLAS::c_value_type

For complex numbers of every precision, the first value in a complex number is the real part and the second is the imaginary part. For real numbers, BLAS::<a/b/c>_value_type is the same as P in Precision<PA, PB, PC> used to describe BLAS (or the default precision).

Input/Output Data Format

This section describes the input and output data format (layout) required for correct calculations.

The tensor API for general matrix multiplication (execute()) and triangular solve (execute() and execute()) are execution descriptor methods which expect matrices represented using (cublasdx::tensor). They both accept matrices represented by tensors with arbitrary layouts. Since the tensor object carries all the information about the dimensions, the memory location and layout of a matrix, no other implicit assumptions are needed.

The dimensions of the matrices must match the dimensions defined by Size operator. See also Get Memory Layout and Suggested Shared Memory Layout sections for using optimized layouts.

The pointer API methods for general matrix multiplication (#2 and #3 overloads of execute()) and triangular solve (#2 and #3 overloads of execute() and execute()) assume that values in input matrices matrix_a, matrix_b (and optionally matrix_c for GEMM) are stored as defined by the Arrangement operator added to the description (by default it’s row-major format for matrix_a, column-major for matrix_b, and column-major for matrix_c).

Shared Memory Usage

It’s important to note that large BLAS operations (as defined by Size operator) may require more than 48 KB of shared memory per CUDA block for the matrices. Therefore, as described in CUDA Programming Guide (Technical Specifications per Compute Capability table, Shared Memory section for each compute capability in Compute Capabilities chapter), kernels with such BLAS operations must use the dynamic shared memory rather than statically sized shared memory arrays. Additionally, these kernels require an explicit opt-in using cudaFuncSetAttribute() to set the cudaFuncAttributeMaxDynamicSharedMemorySize. See example code below.

#include <cublasdx.hpp>
using namespace cublasdx;

using GEMM = decltype(cublasdx::Size<128, 128, 64>()
              + cublasdx::Precision<__nv_fp8_e4m3, __nv_fp8_e5m2, float>()
              + cublasdx::Type<cublasdx::type::real>()
              + cublasdx::Arrangement<cublasdx::row_major, cublasdx::col_major>()
              + cublasdx::Function</* MM or TRSM */>()
              + cublasdx::SM<900>()
              + cublasdx::Block()); // Thread() will output single problem requirement for TRSM

void example() {
  (...)

  // Get required shared memory sizes, options:

  // Shared Memory API GEMM
  // 1 - Shared memory size required for matrices based on GEMM definition
  auto shared_memory_size = cublasdx::get_shared_storage_size<GEMM>();
  // 2 - Shared memory size when dynamic leading dimensions are used
  auto shared_memory_size = cublasdx::get_shared_storage_size<GEMM>(lda, ldb, ldc);
  // 3 - Shared memory size calculated based on custom matrix layouts for A, B, C matrices
  auto shared_memory_size = cublasdx::get_shared_storage_size<GEMM>(matrix_a_layout, matrix_b_layout, matrix_c_layout);

  // TRSM and register API GEMM
  // 1 - Shared memory size required for matrices A and B based on GEMM definition
  auto shared_memory_size = cublasdx::get_shared_storage_size_ab<GEMM>();
  // 2 - Shared memory size when dynamic leading dimensions are used
  auto shared_memory_size = cublasdx::get_shared_storage_size_ab<GEMM>(lda, ldb);
  // 3 - Shared memory size calculated based on custom matrix layouts for A, B matrices
  auto shared_memory_size = cublasdx::get_shared_storage_size_ab<GEMM>(matrix_a_layout, matrix_b_layout);

  // Increases the max dynamic shared memory size to match GEMM requirements
  cudaFuncSetAttribute(gemm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size);
  // Invokes kernel with GEMM::block_dim threads in CUDA block
  gemm_kernel<GEMM><<<1, GEMM::block_dim, shared_memory_size>>>(alpha, a, b, beta, c);

  (...)
}

For block-level TRSM, use a TRSM descriptor with Size<M, N>() and Function<cublasdx::function::TRSM>(). get_shared_storage_size_ab<TRSM>() returns the shared-memory requirement for the triangular A matrix and the input/output B matrix. Thread-level TRSM does not require shared memory.

TRSM Block Shared-Memory Flow

Block-level TRSM uses the same shared-memory utilities as GEMM, but only matrices A and B are staged. A is read-only, while B is input/output and is overwritten in-place with the solution X.

A typical block-level TRSM flow is:

1. Build batched global-memory tensors for A and B.

2. Use get_batch to select the batch tile handled by the CUDA block.

3. Slice dynamic shared memory into A and B tensors.

4. Copy A and B from global memory to shared memory, then wait for the copies to complete.

5. Call TRSM{}.execute(smem_a, smem_b).

6. Synchronize the block, then copy smem_b back to global memory.

template<class TRSM, class GlobalTensorA, class GlobalTensorB>
__global__ void trsm_block_kernel(GlobalTensorA global_a, GlobalTensorB global_b) {
    extern __shared__ __align__(16) cublasdx::byte smem[];

    using TA        = typename TRSM::a_value_type;
    using TB        = typename TRSM::b_value_type;
    using alignment = cublasdx::alignment_of<TRSM>;

    const unsigned first_batch = blockIdx.x * TRSM::batches_per_block;
    auto batch_a = cublasdx::get_batch(global_a, TRSM::get_layout_gmem_a(), first_batch);
    auto batch_b = cublasdx::get_batch(global_b, TRSM::get_layout_gmem_b(), first_batch);

    auto [smem_a, smem_b] =
        cublasdx::shared_memory::slice<TA, TB>(smem,
                                             alignment::a, TRSM::get_layout_smem_a(),
                                             alignment::b, TRSM::get_layout_smem_b());

    cublasdx::copy<TRSM, alignment::a>(batch_a, smem_a);
    cublasdx::copy<TRSM, alignment::b>(batch_b, smem_b);
    cublasdx::copy_wait();

    TRSM{}.execute(smem_a, smem_b);
    __syncthreads();

    cublasdx::copy<TRSM, alignment::b>(smem_b, batch_b);
}

Compute the dynamic shared-memory size from the same layouts used for slicing:

using TA = typename TRSM::a_value_type;
using TB = typename TRSM::b_value_type;

const unsigned smem_bytes =
    cublasdx::make_shared_storage_calculator()
        .add(cublasdx::alignment_of<TRSM>::a, sizeof(TA), TRSM::get_layout_smem_a())
        .add(cublasdx::alignment_of<TRSM>::b, sizeof(TB), TRSM::get_layout_smem_b())
        .get();

cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
kernel<<<num_blocks, TRSM::block_dim, smem_bytes>>>(global_a, global_b);

When BatchesPerBlock is greater than one, get_batch and the descriptor layouts select a 3D tensor containing the consecutive TRSM instances owned by the block. For direct global-memory execution without shared memory, use thread-level TRSM instead; see TRSM Thread Execute Methods.