cuBLASDx#
API reference: cuBLASDx C++ documentation.
Operators#
Operators are set with cublasdxSetOperatorInt64() or cublasdxSetOperatorInt64s():
CUBLASDX_OPERATOR_FUNCTION—CUBLASDX_FUNCTION_MM(matmul).CUBLASDX_OPERATOR_EXECUTION—COMMONDX_EXECUTION_BLOCKorCOMMONDX_EXECUTION_THREAD.CUBLASDX_OPERATOR_API—CUBLASDX_API_SMEM,CUBLASDX_API_SMEM_DYNAMIC_LD, orCUBLASDX_API_TENSORS.CUBLASDX_OPERATOR_PRECISION— one value, or 3 for (A, B, C) withSetOperatorInt64s.CUBLASDX_OPERATOR_TYPE—CUBLASDX_TYPE_REAL,CUBLASDX_TYPE_COMPLEX.CUBLASDX_OPERATOR_SM— one int64 or 2 for (sm, arch modifier) withSetOperatorInt64s.CUBLASDX_OPERATOR_BLOCK_DIM— 3 int64s (x, y, z).CUBLASDX_OPERATOR_SIZE— 3 int64s (M, N, K).CUBLASDX_OPERATOR_TRANSPOSE_MODE— 2 int64s (transpose A, B).CUBLASDX_OPERATOR_ARRANGEMENT— 3 int64s (A, B, C layout).CUBLASDX_OPERATOR_LEADING_DIMENSION— 3 int64s (LDA, LDB, LDC) (optional).CUBLASDX_OPERATOR_ALIGNMENT— 3 int64s (optional).CUBLASDX_OPERATOR_STATIC_BLOCK_DIM— 1 to enable static block dimensions (optional, experimental).CUBLASDX_OPERATOR_ENABLE_INPUT_STREAMING— 1 to enable streaming tiles (pipelining).CUBLASDX_OPERATOR_WITH_PIPELINE— 1 for pipelining.
Options#
Use cublasdxSetOptionStr() to set COMMONDX_OPTION_SYMBOL_NAME for the device function name.
SMEM API device signatures#
void symbol(value_type* alpha, value_type* a, value_type* b, value_type* beta, value_type* c)With dynamic leading dimensions:
void symbol(value_type* alpha, value_type* a, unsigned* lda, value_type* b, unsigned* ldb, value_type* beta, value_type* c, unsigned* ldc)
Traits: CUBLASDX_TRAIT_SYMBOL_NAME, CUBLASDX_TRAIT_BLOCK_DIM, CUBLASDX_TRAIT_SIZE, etc.
Tensor API#
Create tensors with
cublasdxCreateTensor()— e.g.CUBLASDX_TENSOR_SUGGESTED_SMEM_A/B,CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C,CUBLASDX_TENSOR_SUGGESTED_SMEM_C,CUBLASDX_TENSOR_SUGGESTED_RMEM_C.For GMEM use
cublasdxCreateTensorStrided()withLIBMATHDX_RUNTIMEfor runtime dimensions.Create device functions with
cublasdxCreateDeviceFunction().Finalize code with
cublasdxFinalizeDeviceFunctions()(and set code target SM).Query tensor traits:
cublasdxGetTensorTraitInt64(),cublasdxGetTensorTraitStrSize()/cublasdxGetTensorTraitStr()(e.g.CUBLASDX_TENSOR_TRAIT_STORAGE_BYTES,CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME).
Pipelines#
Create pipelines with
cublasdxCreateDevicePipeline()orcublasdxCreateTilePipeline().Finalize with
cublasdxFinalizePipelines()(orcublasdxFinalize()).Use device functions that take pipelines and accumulator; finalize with
cublasdxFinalizeDeviceFunctions().Epilogue callback: use
cublasdxSetDeviceFunctionOptionStr()withCUBLASDX_DEVICE_FUNCTION_OPTION_CALLBACK.
Examples#
cuBLASDx “pointer API” example#
/*
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <libcublasdx.h>
#include <array>
#include <vector>
#include "arch.hpp"
#include "macros.hpp"
using namespace examples;
/**
* Basic cuBLASDx GEMM example: configures a block-level matrix multiply descriptor,
* compiles to LTOIR, and reports the generated code size.
*/
int main() {
int m = 32;
int n = 8;
int k = 16;
int num_threads = 32;
arch_t dx_sm = get_dx_sm();
arch_t target_sm = get_target_sm();
/**
* Create a descriptor
* This is equivalent to `using BLAS = ...` in cuBLASDx C++
*/
cublasdxDescriptor h { 0 };
LIBMATHDX_CHECK(cublasdxCreateDescriptor(&h));
// This means we generate a Matmul ("MM")
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_FUNCTION, cublasdxFunction::CUBLASDX_FUNCTION_MM));
// COMMONDX_EXECUTION_BLOCK means we are generating a function with a "Block" API semantic. All threads in the CUDA
// block must participate and will cooperate to compute the matmul.
// COMMONDX_EXECUTION_THREAD would mean that each thread compute its own matmul independently from other threads.
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_EXECUTION, commondxExecution::COMMONDX_EXECUTION_BLOCK));
// CUBLASDX_API_SMEM means the function take inputs in shared memory and produce output in shared memory.
// The leading dimensions are fixed at compile time. The function signature is:
// void gemm(value_type, value_type*, value_type*, value_type, value_type)
// CUBLASDX_API_SMEM_DYNAMIC_LD would mean that the function takes runtime leading dimension values.
// Such function has the following signature:
// void gemm(value_type, value_type*, unsigned, value_type*, unsigned, value_type, value_type*, unsigned)
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_API, cublasdxApi::CUBLASDX_API_SMEM));
// COMMONDX_PRECISION_F16 means the matrices are filled with half-precision floating point numbers
// COMMONDX_PRECISION_F32 would be for single-precision
// COMMONDX_PRECISION_F64 would be for double precision
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_PRECISION, commondxPrecision::COMMONDX_PRECISION_F16));
// CUBLASDX_OPERATOR_SM indicates the target architecture (700 for SM70, etc)
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_SM, dx_sm.operator_sm()));
// CUBLASDX_TYPE_REAL means the matrices contain real type data
// CUBLASDX_TYPE_COMPLEX would be for matrices with complex type data
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_TYPE, cublasdxType::CUBLASDX_TYPE_REAL));
// CUBLASDX_OPERATOR_BLOCK_DIM indicates the block dimension
std::array<long long int, 3> block_dim = { num_threads, 1, 1 };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_BLOCK_DIM, block_dim.size(), block_dim.data()));
// CUBLASDX_OPERATOR_SIZE indicates the (M, N, K) of the problem.
std::array<long long int, 3> size = { m, n, k };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_SIZE, size.size(), size.data()));
// CUBLASDX_OPERATOR_TRANSPOSE_MODE is a tuple (ta, tb), where 'ta' and 'tb' can take one of the following
// values:
// - CUBLASDX_TRANSPOSE_MODE_NON_TRANSPOSED - the input matrix (A or B) is not transposed
// - CUBLASDX_TRANSPOSE_MODE_TRANSPOSED - the input matrix (A or B) is transposed
// - CUBLASDX_TRANSPOSE_MODE_CONJ_TRANSPOSED - the input matrix (A or B) is conjugated and transposed
std::array<long long int, 2> transpose_mode = { cublasdxTransposeMode_t::CUBLASDX_TRANSPOSE_MODE_NON_TRANSPOSED,
cublasdxTransposeMode_t::CUBLASDX_TRANSPOSE_MODE_NON_TRANSPOSED };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_TRANSPOSE_MODE, transpose_mode.size(), transpose_mode.data()));
// COMMONDX_OPTION_SYMBOL_NAME indicates the required name for the device function.
LIBMATHDX_CHECK(cublasdxSetOptionStr(h, commondxOption::COMMONDX_OPTION_SYMBOL_NAME, "my_gemm"));
/**
* Compile the device function
*/
commondxCode code;
LIBMATHDX_CHECK(commondxCreateCode(&code));
// Specify arch to compile to
LIBMATHDX_CHECK(commondxSetCodeOptionInt64(code, COMMONDX_OPTION_TARGET_SM, target_sm.operator_sm()));
LIBMATHDX_CHECK(cublasdxFinalizeCode(code, h));
size_t lto_size = 0;
LIBMATHDX_CHECK(commondxGetCodeLTOIRSize(code, <o_size));
std::vector<char> lto(lto_size);
LIBMATHDX_CHECK(commondxGetCodeLTOIR(code, lto.size(), lto.data()));
long long int isa = 0;
LIBMATHDX_CHECK(commondxGetCodeOptionInt64(code, COMMONDX_OPTION_CODE_ISA, &isa));
LIBMATHDX_CHECK(commondxDestroyCode(code));
printf("Successfully generated LTOIR (version %lld), %zu bytes for GEMM %d x %d x %d\n", isa, lto_size, m, n, k);
LIBMATHDX_CHECK(cublasdxDestroyDescriptor(h));
}
cuBLASDx “tensor API” example#
/*
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <libcommondx.h>
#include <libcublasdx.h>
#include <array>
#include <vector>
#include "arch.hpp"
#include "macros.hpp"
using namespace examples;
int main() {
long long int m = 256;
long long int n = 128;
long long int k = 16;
long long int num_threads = 128;
arch_t dx_sm = get_dx_sm();
arch_t target_sm = get_target_sm();
auto dx_sm_array = dx_sm.to_array();
auto target_sm_array = target_sm.to_array();
/**
* Create the cuBLASDx descriptor
*/
cublasdxDescriptor h { 0 };
LIBMATHDX_CHECK(cublasdxCreateDescriptor(&h));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_FUNCTION, cublasdxFunction::CUBLASDX_FUNCTION_MM));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_EXECUTION, commondxExecution::COMMONDX_EXECUTION_BLOCK));
// Using the Opaque Tensor API
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_API, cublasdxApi::CUBLASDX_API_TENSORS));
std::array<long long int, 3> prec = { COMMONDX_PRECISION_F32, COMMONDX_PRECISION_F32, COMMONDX_PRECISION_F32 };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_PRECISION, prec.size(), prec.data()));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_SM, dx_sm_array.size(), dx_sm_array.data()));
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_TYPE, cublasdxType::CUBLASDX_TYPE_REAL));
std::array<long long int, 3> block_dim = { num_threads, 1, 1 };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_BLOCK_DIM, block_dim.size(), block_dim.data()));
std::array<long long int, 3> size = { m, n, k };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_SIZE, size.size(), size.data()));
std::array<long long int, 3> arrangement = { CUBLASDX_ARRANGEMENT_COL_MAJOR,
CUBLASDX_ARRANGEMENT_COL_MAJOR,
CUBLASDX_ARRANGEMENT_COL_MAJOR };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_ARRANGEMENT, arrangement.size(), arrangement.data()));
LIBMATHDX_CHECK(cublasdxSetOptionStr(h, commondxOption::COMMONDX_OPTION_SYMBOL_NAME, "matmul"));
/**
* Define the input and output tensors
*/
cublasdxTensor smem_a { 0 };
cublasdxTensor smem_b { 0 };
cublasdxTensor acc_c { 0 };
cublasdxTensor smem_c { 0 };
cublasdxTensor rmem_c { 0 };
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_SMEM_A, &smem_a));
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_SMEM_B, &smem_b));
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C, &acc_c));
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_SMEM_C, &smem_c));
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_RMEM_C, &rmem_c));
cublasdxTensor big_gmem {};
std::vector<long long int> shape = { m, k };
std::vector<long long int> strides = { LIBMATHDX_RUNTIME, 1 };
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(
CUBLASDX_MEMORY_SPACE_GMEM, COMMONDX_R_32F, nullptr, shape.size(), shape.data(), strides.data(), &big_gmem));
std::array tensors = { smem_a, smem_b, acc_c, smem_c, rmem_c, big_gmem };
LIBMATHDX_CHECK(cublasdxFinalizeTensors(tensors.size(), tensors.data()));
for (auto t : tensors) {
long long int alignment = 0;
long long int size = 0;
size_t name_size = 0;
LIBMATHDX_CHECK(cublasdxGetTensorTraitInt64(t, CUBLASDX_TENSOR_TRAIT_ALIGNMENT_BYTES, &alignment));
LIBMATHDX_CHECK(cublasdxGetTensorTraitInt64(t, CUBLASDX_TENSOR_TRAIT_STORAGE_BYTES, &size));
LIBMATHDX_CHECK(cublasdxGetTensorTraitStrSize(t, CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME, &name_size));
std::vector<char> name(name_size);
LIBMATHDX_CHECK(cublasdxGetTensorTraitStr(t, CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME, name.size(), name.data()));
printf("Tensor %lld: name %s, storage size %lld B, alignment %lld B\n",
static_cast<long long int>(t),
name.data(),
size,
alignment);
}
/**
* Define a function operating on those input and output tensors.
*
* The device function output is an opaque and stateful accumulator.
*/
std::array gemm_tensors = { smem_a, smem_b, acc_c };
cublasdxDeviceFunction gemm_sa_sb_rc { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(
h, CUBLASDX_DEVICE_FUNCTION_EXECUTE, gemm_tensors.size(), gemm_tensors.data(), &gemm_sa_sb_rc));
cublasdxDeviceFunction init { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(h, CUBLASDX_DEVICE_FUNCTION_CREATE, 1, &acc_c, &init));
cublasdxDeviceFunction destroy { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 1, &acc_c, &destroy));
cublasdxDeviceFunction clear { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(h, CUBLASDX_DEVICE_FUNCTION_CLEAR, 1, &acc_c, &clear));
std::array copy_tensors_smem = { acc_c, smem_c };
cublasdxDeviceFunction copy_smem { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(
LIBMATHDX_NONE, CUBLASDX_DEVICE_FUNCTION_COPY, copy_tensors_smem.size(), copy_tensors_smem.data(), ©_smem));
LIBMATHDX_CHECK(
cublasdxSetDeviceFunctionOptionInt64(copy_smem, CUBLASDX_DEVICE_FUNCTION_OPTION_NUM_THREADS, num_threads));
std::array copy_tensors_rmem = { acc_c, rmem_c };
cublasdxDeviceFunction copy_rmem { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(
h, CUBLASDX_DEVICE_FUNCTION_COPY, copy_tensors_rmem.size(), copy_tensors_rmem.data(), ©_rmem));
std::array copy_tensors_big = { big_gmem, smem_a };
cublasdxDeviceFunction copy_big { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(
h, CUBLASDX_DEVICE_FUNCTION_COPY, copy_tensors_big.size(), copy_tensors_big.data(), ©_big));
{
size_t symbol_size { 0 };
LIBMATHDX_CHECK(
cublasdxGetDeviceFunctionTraitStrSize(gemm_sa_sb_rc, CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL, &symbol_size));
std::vector<char> symbol(symbol_size);
LIBMATHDX_CHECK(cublasdxGetDeviceFunctionTraitStr(
gemm_sa_sb_rc, CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL, symbol.size(), symbol.data()));
printf("Device function %lld: symbol: %s\n", static_cast<long long int>(gemm_sa_sb_rc), symbol.data());
}
std::vector<cublasdxDeviceFunction> functions = { gemm_sa_sb_rc, init, destroy, copy_smem,
copy_rmem, clear, copy_big };
/**
* Compile the device function to lto
*/
commondxCode code { 0 };
LIBMATHDX_CHECK(commondxCreateCode(&code));
LIBMATHDX_CHECK(
commondxSetCodeOptionInt64s(code, COMMONDX_OPTION_TARGET_SM, target_sm_array.size(), target_sm_array.data()));
LIBMATHDX_CHECK(cublasdxFinalizeDeviceFunctions(code, functions.size(), functions.data()));
/**
* Extract the LTOIR
*/
std::vector<char> lto;
size_t lto_size = 0;
LIBMATHDX_CHECK(commondxGetCodeLTOIRSize(code, <o_size));
lto.resize(lto_size);
LIBMATHDX_CHECK(commondxGetCodeLTOIR(code, lto_size, lto.data()));
printf("Generated LTOIR for GEMM device functions, %zu bytes\n", lto.size());
/**
* Destroy handles
*/
LIBMATHDX_CHECK(cublasdxDestroyTensor(smem_a));
LIBMATHDX_CHECK(cublasdxDestroyTensor(smem_b));
LIBMATHDX_CHECK(cublasdxDestroyTensor(acc_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(smem_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(rmem_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(big_gmem));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(gemm_sa_sb_rc));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(copy_smem));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(copy_rmem));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(clear));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(copy_big));
LIBMATHDX_CHECK(commondxDestroyCode(code));
LIBMATHDX_CHECK(cublasdxDestroyDescriptor(h));
}
cuBLASDx pipeline API example#
/*
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <cuda.h>
#include <cuda_runtime.h>
#include <libcublasdx.h>
#include <nvJitLink.h>
#include <nvrtc.h>
#include <array>
#include <random>
#include <string>
#include <unordered_map>
#include <vector>
#include "arch.hpp"
#include "common_examples.hpp"
#include "macros.hpp"
using namespace examples;
int main() {
long long int m = 512;
long long int n = 512;
long long int k = 512;
long long int tile_m = 128;
long long int tile_n = 128;
long long int tile_k = 32;
long long int num_threads = 128;
arch_t dx_sm = maybe_accelerated_dx(get_dx_cc());
arch_t target_sm = maybe_accelerated_target(get_target_cc());
auto dx_sm_array = dx_sm.to_array();
auto target_sm_array = target_sm.to_array();
auto [nvrtc_major, nvrtc_minor] = get_nvrtc_version();
if (nvrtc_major == 13 && nvrtc_minor == 0 && target_sm.cc >= cc_t { 10, 0 }) {
printf("Pipeline examples on SM100+ requires NVRTC 13.1 or above.\n");
return 0;
}
if (dx_sm.cc < cc_t { 7, 5 }) {
printf("Pipeline examples require SM75 or above.\n");
return 0;
}
/**
* Create the cuBLASDx descriptor
*/
cublasdxDescriptor h { 0 };
LIBMATHDX_CHECK(cublasdxCreateDescriptor(&h));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_FUNCTION, cublasdxFunction::CUBLASDX_FUNCTION_MM));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_EXECUTION, commondxExecution::COMMONDX_EXECUTION_BLOCK));
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_API, cublasdxApi::CUBLASDX_API_TENSORS));
std::array<long long int, 3> prec = { COMMONDX_PRECISION_I8, COMMONDX_PRECISION_I8, COMMONDX_PRECISION_I32 };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_PRECISION, prec.size(), prec.data()));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_SM, dx_sm_array.size(), dx_sm_array.data()));
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_TYPE, cublasdxType::CUBLASDX_TYPE_REAL));
std::array<long long int, 3> block_dim = { num_threads, 1, 1 };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_BLOCK_DIM, block_dim.size(), block_dim.data()));
std::array<long long int, 3> size = { tile_m, tile_n, tile_k };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_SIZE, size.size(), size.data()));
std::array<long long int, 3> arrangement = { CUBLASDX_ARRANGEMENT_ROW_MAJOR,
CUBLASDX_ARRANGEMENT_COL_MAJOR,
CUBLASDX_ARRANGEMENT_ROW_MAJOR };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_ARRANGEMENT, arrangement.size(), arrangement.data()));
std::array<long long int, 3> alignment = { 16, 16, 16 };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_ALIGNMENT, alignment.size(), alignment.data()));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_ENABLE_INPUT_STREAMING, 1));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_WITH_PIPELINE, 1));
LIBMATHDX_CHECK(cublasdxSetOptionStr(h, commondxOption::COMMONDX_OPTION_SYMBOL_NAME, "matmul"));
/**
* Define the input and output tensors
*/
int8_t* d_a {};
int8_t* d_b {};
int32_t* d_c {};
CUDA_CHECK(cudaMallocManaged(&d_a, m * k * sizeof(int8_t)));
CUDA_CHECK(cudaMallocManaged(&d_b, k * n * sizeof(int8_t)));
CUDA_CHECK(cudaMallocManaged(&d_c, m * n * sizeof(int32_t)));
CUDA_CHECK(cudaMemset(d_c, 0, m * n * sizeof(int32_t)));
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<> dis(0, 10);
for (int i = 0; i < m * k; i++) {
d_a[i] = static_cast<int8_t>(dis(gen));
}
for (int i = 0; i < k * n; i++) {
d_b[i] = static_cast<int8_t>(dis(gen));
}
cublasdxTensor matrix_a { 0 };
cublasdxTensor matrix_b { 0 };
cublasdxTensor tile_gemm_c { 0 };
std::vector<long long int> shape_a = { m, k };
std::vector<long long int> shape_b = { k, n };
// Importantly, the shape of the output tensor is the tile size, not the global size
// the raw pointer must be offset for each CTA tile as shown in the kernel code
std::vector<long long int> shape_c = { tile_m, tile_n };
std::vector<long long int> strides_a = { k, 1 };
std::vector<long long int> strides_b = { 1, k };
// while the shape of the tile gemm c is the tile size, the strides are the global size
std::vector<long long int> strides_c = { n, 1 };
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(CUBLASDX_MEMORY_SPACE_GMEM,
COMMONDX_R_8I,
nullptr,
shape_a.size(),
shape_a.data(),
strides_a.data(),
&matrix_a));
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(CUBLASDX_MEMORY_SPACE_GMEM,
COMMONDX_R_8I,
nullptr,
shape_b.size(),
shape_b.data(),
strides_b.data(),
&matrix_b));
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(CUBLASDX_MEMORY_SPACE_GMEM,
COMMONDX_R_32I,
nullptr,
shape_c.size(),
shape_c.data(),
strides_c.data(),
&tile_gemm_c));
cublasdxTensor acc_c { 0 };
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C, &acc_c));
cublasdxPipeline device_pipeline { 0 };
cublasdxPipeline tile_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDevicePipeline(h,
CUBLASDX_DEVICE_PIPELINE_SUGGESTED,
4,
CUBLASDX_BLOCK_SIZE_STRATEGY_FIXED,
matrix_a,
matrix_b,
&device_pipeline));
LIBMATHDX_CHECK(cublasdxCreateTilePipeline(h, CUBLASDX_TILE_PIPELINE_DEFAULT, device_pipeline, &tile_pipeline));
std::array tensors = { matrix_a, matrix_b, tile_gemm_c, acc_c };
std::array pipelines = { device_pipeline, tile_pipeline };
LIBMATHDX_CHECK(cublasdxFinalize(tensors.size(), tensors.data(), pipelines.size(), pipelines.data()));
// size and alignment of accumulator tensor
long long int acc_size = 0;
LIBMATHDX_CHECK(cublasdxGetTensorTraitInt64(acc_c, CUBLASDX_TENSOR_TRAIT_STORAGE_BYTES, &acc_size));
long long int acc_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetTensorTraitInt64(acc_c, CUBLASDX_TENSOR_TRAIT_ALIGNMENT_BYTES, &acc_alignment));
// size and alignment of the device pipeline
long long int device_pipeline_storage_size = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_BYTES, &device_pipeline_storage_size));
long long int device_pipeline_storage_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_ALIGNMENT_BYTES, &device_pipeline_storage_alignment));
// size and alignment of the shared memory buffer (on sm_90a and above (accelerated arches with suffix 'a') should
// be 128B aligned for TMA usage)
long long int shared_memory_buffer_size = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_BUFFER_SIZE, &shared_memory_buffer_size));
long long int shared_memory_buffer_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_BUFFER_ALIGNMENT_BYTES, &shared_memory_buffer_alignment));
// block dimension for GEMM kernel launch (might be different than operator block dim)
std::array<long long int, 3> device_pipeline_block_dim = { 0, 0, 0 };
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64s(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_BLOCK_DIM, 3, device_pipeline_block_dim.data()));
// size and alignment of the tile pipeline
long long int tile_pipeline_storage_size = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
tile_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_BYTES, &tile_pipeline_storage_size));
long long int tile_pipeline_storage_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
tile_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_ALIGNMENT_BYTES, &tile_pipeline_storage_alignment));
/**
* Define a function operating on those input and output tensors.
*
* The device function output is an opaque and stateful accumulator.
*/
cublasdxDeviceFunction init_acc { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_CREATE, 1, &acc_c, 1, &tile_pipeline, &init_acc));
cublasdxDeviceFunction init_device_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_CREATE, 0, nullptr, 1, &device_pipeline, &init_device_pipeline));
cublasdxDeviceFunction init_tile_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_CREATE, 0, nullptr, 1, &tile_pipeline, &init_tile_pipeline));
cublasdxDeviceFunction destroy_acc { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 1, &acc_c, &destroy_acc));
cublasdxDeviceFunction destroy_device_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 0, nullptr, 1, &device_pipeline, &destroy_device_pipeline));
cublasdxDeviceFunction destroy_tile_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 0, nullptr, 1, &tile_pipeline, &destroy_tile_pipeline));
cublasdxDeviceFunction execute_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_EXECUTE, 1, &acc_c, 1, &tile_pipeline, &execute_pipeline));
// Finalize the reusable accumulator after execute and before reading it back with copy
// (the non-epilogue read path does not finalize on its own).
cublasdxDeviceFunction finish_accumulation { 0 };
LIBMATHDX_CHECK(
cublasdxCreateDeviceFunction(h, CUBLASDX_DEVICE_FUNCTION_FINISH_ACCUMULATION, 1, &acc_c, &finish_accumulation));
std::array copy_tensors = { acc_c, tile_gemm_c };
cublasdxDeviceFunction copy_acc_big_gmem_c { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(
h, CUBLASDX_DEVICE_FUNCTION_COPY, copy_tensors.size(), copy_tensors.data(), ©_acc_big_gmem_c));
std::vector<cublasdxDeviceFunction> functions = {
execute_pipeline, finish_accumulation, init_acc,
init_device_pipeline, init_tile_pipeline, destroy_acc,
destroy_device_pipeline, destroy_tile_pipeline, copy_acc_big_gmem_c
};
std::unordered_map<cublasdxDeviceFunction, std::string> function_symbols;
for (auto f : functions) {
size_t symbol_size { 0 };
LIBMATHDX_CHECK(cublasdxGetDeviceFunctionTraitStrSize(f, CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL, &symbol_size));
std::vector<char> symbol(symbol_size, '\0');
LIBMATHDX_CHECK(
cublasdxGetDeviceFunctionTraitStr(f, CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL, symbol.size(), symbol.data()));
printf("Device function %lld: symbol: %s\n", static_cast<long long int>(f), symbol.data());
function_symbols[f] = std::string(symbol.data());
}
/**
* Compile the device functions to LTOIR
*/
commondxCode code { 0 };
LIBMATHDX_CHECK(commondxCreateCode(&code));
LIBMATHDX_CHECK(
commondxSetCodeOptionInt64s(code, COMMONDX_OPTION_TARGET_SM, target_sm_array.size(), target_sm_array.data()));
LIBMATHDX_CHECK(cublasdxFinalizeDeviceFunctions(code, functions.size(), functions.data()));
/**
* Extract the LTOIR
*/
size_t lto_size = 0;
LIBMATHDX_CHECK(commondxGetCodeLTOIRSize(code, <o_size));
std::vector<char> lto(lto_size);
LIBMATHDX_CHECK(commondxGetCodeLTOIR(code, lto_size, lto.data()));
printf("Generated LTOIR for GEMM pipeline, %zu bytes\n", lto.size());
const char kernel_template[] = R"(
struct libmathdx_tensor_0s_0s { void* ptr; };
struct libmathdx_pipeline { void* ptr; };
#define M %d
#define N %d
#define K %d
#define tile_m %d
#define tile_n %d
#define tile_k %d
#define acc_name libmathdx_tensor_0s_0s
#define ga_name libmathdx_tensor_0s_0s
#define gb_name libmathdx_tensor_0s_0s
#define gc_name libmathdx_tensor_0s_0s
#define device_pipeline_name libmathdx_pipeline
#define tile_pipeline_name libmathdx_pipeline
constexpr unsigned acc_size = %lld;
constexpr unsigned acc_alignment = %lld;
constexpr unsigned block_size = %lld;
constexpr unsigned smem_alignment = %lld;
constexpr unsigned tile_pipeline_size = %lld;
constexpr unsigned tile_pipeline_alignment = %lld;
#define execute_pipeline_acc %s
#define finish_acc %s
#define copy_acc_gc %s
#define create_dev_pipe %s
#define create_tile_pipe %s
#define create_acc %s
#define destroy_dev_pipe %s
#define destroy_tile_pipe %s
#define destroy_acc %s
using C_VALUE_TYPE = signed int;
extern "C" __device__ void execute_pipeline_acc(tile_pipeline_name, acc_name);
extern "C" __device__ void finish_acc(acc_name);
extern "C" __device__ void copy_acc_gc(acc_name, gc_name);
extern "C" __device__ void create_dev_pipe(device_pipeline_name, ga_name, gb_name);
extern "C" __device__ void create_tile_pipe(device_pipeline_name, tile_pipeline_name, char*, int*, int*);
extern "C" __device__ void create_acc(tile_pipeline_name, acc_name);
extern "C" __device__ void destroy_dev_pipe(device_pipeline_name);
extern "C" __device__ void destroy_tile_pipe(tile_pipeline_name);
extern "C" __device__ void destroy_acc(acc_name);
// Create the device pipeline
extern "C" __global__ void create_device_pipeline(void* device_pipeline_ptr, void* ga_storage, void* gb_storage) {
if(threadIdx.x == 0) {
auto ga = ga_name { ga_storage };
auto gb = gb_name { gb_storage };
auto device_pipeline = device_pipeline_name { device_pipeline_ptr };
create_dev_pipe(device_pipeline, ga, gb);
}
}
// Perform the GEMM
extern "C" __launch_bounds__(block_size, 1) __global__ void gemm(void* device_pipeline_ptr, void* gc_storage)
{
int row_id = blockIdx.x;
int col_id = blockIdx.y;
C_VALUE_TYPE* gmem_c = reinterpret_cast<C_VALUE_TYPE*>(gc_storage) + ((row_id * tile_m) * N + (col_id * tile_n));
// Allocate dynamic shared memory for tiles
extern __shared__ __align__(smem_alignment) char smem[];
// Allocate local memory
alignas(acc_alignment) char acc_storage[acc_size];
alignas(tile_pipeline_alignment) char tile_pipeline_storage[tile_pipeline_size];
// Create opaque types
auto acc = acc_name { acc_storage };
auto gc = gc_name { gmem_c };
auto device_pipeline = device_pipeline_name { device_pipeline_ptr };
auto tile_pipeline = tile_pipeline_name { tile_pipeline_storage };
// Perform the GEMM
create_tile_pipe(device_pipeline, tile_pipeline, smem, &row_id, &col_id);
create_acc(tile_pipeline, acc);
execute_pipeline_acc(tile_pipeline, acc);
finish_acc(acc);
copy_acc_gc(acc, gc);
// Destroy the accumulator and the tile pipeline
destroy_acc(acc);
destroy_tile_pipe(tile_pipeline);
}
// Destroy the device pipeline
extern "C" __global__ void destroy_device_pipeline(void* device_pipeline_ptr) {
if(threadIdx.x == 0) {
auto device_pipeline = device_pipeline_name { device_pipeline_ptr };
destroy_dev_pipe(device_pipeline);
}
}
)";
std::string cpp = strprintf(kernel_template,
m,
n,
k,
tile_m,
tile_n,
tile_k,
acc_size,
acc_alignment,
device_pipeline_block_dim[0],
shared_memory_buffer_alignment,
tile_pipeline_storage_size,
tile_pipeline_storage_alignment,
function_symbols[execute_pipeline].data(),
function_symbols[finish_accumulation].data(),
function_symbols[copy_acc_big_gmem_c].data(),
function_symbols[init_device_pipeline].data(),
function_symbols[init_tile_pipeline].data(),
function_symbols[init_acc].data(),
function_symbols[destroy_device_pipeline].data(),
function_symbols[destroy_tile_pipeline].data(),
function_symbols[destroy_acc].data());
std::vector<char> cubin = compile_and_link(cpp, lto, target_sm);
CUmodule module {};
CUfunction kernel, kernel2, kernel3 {};
CUDA_CHECK(cudaSetDevice(0));
CU_CHECK(cuModuleLoadDataEx(&module, cubin.data(), 0, 0, 0));
CU_CHECK(cuModuleGetFunction(&kernel, module, "create_device_pipeline"));
CU_CHECK(cuModuleGetFunction(&kernel2, module, "gemm"));
CU_CHECK(cuModuleGetFunction(&kernel3, module, "destroy_device_pipeline"));
// create the device pipeline using trait
uint8_t* device_pipeline_ptr {};
CUDA_CHECK(cudaMallocManaged(&device_pipeline_ptr, device_pipeline_storage_size));
CUDA_CHECK(cudaMemset(device_pipeline_ptr, 0, device_pipeline_storage_size));
{
std::vector<void*> kernel_args = { reinterpret_cast<void*>(&device_pipeline_ptr),
reinterpret_cast<void*>(&d_a),
reinterpret_cast<void*>(&d_b) };
CU_CHECK(cuLaunchKernel(kernel, 1, 1, 1, 1, 1, 1, 0, nullptr, kernel_args.data(), nullptr));
CUDA_CHECK(cudaDeviceSynchronize());
}
{
printf("Checking device pipeline bytes\n");
printf("Pointers to gmem A & B: %p %p\n", d_a, d_b);
printf("Device pipeline bytes: \n");
for (long long int i = 0; i < device_pipeline_storage_size; i++) {
if (i > 0 && i % 32 == 0) printf("\n");
printf("%02x ", device_pipeline_ptr[i]);
}
printf("\n");
}
{
CU_CHECK(cuFuncSetAttribute(
kernel2, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, static_cast<int>(shared_memory_buffer_size)));
std::vector<void*> kernel_args = { reinterpret_cast<void*>(&device_pipeline_ptr),
reinterpret_cast<void*>(&d_c) };
CU_CHECK(cuLaunchKernel(kernel2,
static_cast<unsigned int>(m / tile_m),
static_cast<unsigned int>(n / tile_n),
1,
static_cast<unsigned int>(device_pipeline_block_dim[0]),
static_cast<unsigned int>(device_pipeline_block_dim[1]),
static_cast<unsigned int>(device_pipeline_block_dim[2]),
static_cast<unsigned int>(shared_memory_buffer_size),
nullptr,
kernel_args.data(),
nullptr));
CUDA_CHECK(cudaDeviceSynchronize());
}
{
std::vector<void*> kernel_args = { reinterpret_cast<void*>(&device_pipeline_ptr) };
CU_CHECK(cuLaunchKernel(kernel3, 1, 1, 1, 1, 1, 1, 0, nullptr, kernel_args.data(), nullptr));
CUDA_CHECK(cudaDeviceSynchronize());
}
std::vector<int32_t> h_c_ref(m * n, 0);
for (int l = 0; l < k; l++) {
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
h_c_ref[i * n + j] += d_a[i * k + l] * d_b[j * k + l];
}
}
}
for (int i = 0; i < m * n; i++) {
if (h_c_ref[i] != d_c[i]) {
printf("Error at %d: h_c_ref[%d] = %d, d_c[%d] = %d\n", i, i, h_c_ref[i], i, d_c[i]);
abort();
}
}
printf("Successfully ran the kernel\n");
/**
* Destroy handles
*/
CU_CHECK(cuModuleUnload(module));
CUDA_CHECK(cudaFree(device_pipeline_ptr));
CUDA_CHECK(cudaFree(d_a));
CUDA_CHECK(cudaFree(d_b));
CUDA_CHECK(cudaFree(d_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(matrix_a));
LIBMATHDX_CHECK(cublasdxDestroyTensor(matrix_b));
LIBMATHDX_CHECK(cublasdxDestroyTensor(acc_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(tile_gemm_c));
LIBMATHDX_CHECK(cublasdxDestroyPipeline(device_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyPipeline(tile_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init_acc));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init_device_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init_tile_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy_acc));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy_device_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy_tile_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(execute_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(finish_accumulation));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(copy_acc_big_gmem_c));
LIBMATHDX_CHECK(commondxDestroyCode(code));
LIBMATHDX_CHECK(cublasdxDestroyDescriptor(h));
}
cuBLASDx pipelining callback example#
/*
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <cuda.h>
#include <cuda_runtime.h>
#include <libcublasdx.h>
#include <nvJitLink.h>
#include <nvrtc.h>
#include <array>
#include <random>
#include <string>
#include <unordered_map>
#include <vector>
#include "common_examples.hpp"
#include "macros.hpp"
using namespace examples;
int main() {
long long int m = 512;
long long int n = 512;
long long int k = 512;
long long int tile_m = 128;
long long int tile_n = 128;
long long int tile_k = 32;
long long int num_threads = 128;
arch_t dx_sm = maybe_accelerated_dx(get_dx_cc());
arch_t target_sm = maybe_accelerated_target(get_target_cc());
auto dx_sm_array = dx_sm.to_array();
auto target_sm_array = target_sm.to_array();
auto [nvrtc_major, nvrtc_minor] = get_nvrtc_version();
if (nvrtc_major == 13 && nvrtc_minor == 0 && target_sm.cc >= cc_t { 10, 0 }) {
printf("Pipeline examples on SM100+ requires NVRTC 13.1 or above.\n");
return 0;
}
if (dx_sm.cc < cc_t { 7, 5 }) {
printf("Pipeline examples require SM75 or above.\n");
return 0;
}
/**
* Create the cuBLASDx descriptor
*/
cublasdxDescriptor h { 0 };
LIBMATHDX_CHECK(cublasdxCreateDescriptor(&h));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_FUNCTION, cublasdxFunction::CUBLASDX_FUNCTION_MM));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_EXECUTION, commondxExecution::COMMONDX_EXECUTION_BLOCK));
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_API, cublasdxApi::CUBLASDX_API_TENSORS));
std::array<long long int, 3> prec = { COMMONDX_PRECISION_I8, COMMONDX_PRECISION_I8, COMMONDX_PRECISION_I32 };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_PRECISION, prec.size(), prec.data()));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_SM, dx_sm_array.size(), dx_sm_array.data()));
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_TYPE, cublasdxType::CUBLASDX_TYPE_REAL));
std::array<long long int, 3> block_dim = { num_threads, 1, 1 };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_BLOCK_DIM, block_dim.size(), block_dim.data()));
std::array<long long int, 3> size = { tile_m, tile_n, tile_k };
LIBMATHDX_CHECK(
cublasdxSetOperatorInt64s(h, cublasdxOperatorType::CUBLASDX_OPERATOR_SIZE, size.size(), size.data()));
std::array<long long int, 3> arrangement = { CUBLASDX_ARRANGEMENT_COL_MAJOR,
CUBLASDX_ARRANGEMENT_COL_MAJOR,
CUBLASDX_ARRANGEMENT_COL_MAJOR };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_ARRANGEMENT, arrangement.size(), arrangement.data()));
std::array<long long int, 3> alignment = { 16, 16, 16 };
LIBMATHDX_CHECK(cublasdxSetOperatorInt64s(
h, cublasdxOperatorType::CUBLASDX_OPERATOR_ALIGNMENT, alignment.size(), alignment.data()));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_ENABLE_INPUT_STREAMING, 1));
LIBMATHDX_CHECK(cublasdxSetOperatorInt64(h, cublasdxOperatorType::CUBLASDX_OPERATOR_WITH_PIPELINE, 1));
LIBMATHDX_CHECK(cublasdxSetOptionStr(h, commondxOption::COMMONDX_OPTION_SYMBOL_NAME, "matmul"));
/**
* Define the input and output tensors
*/
int8_t* d_a {};
int8_t* d_b {};
int32_t* d_c {};
CUDA_CHECK(cudaMallocManaged(&d_a, m * k * sizeof(int8_t)));
CUDA_CHECK(cudaMallocManaged(&d_b, k * n * sizeof(int8_t)));
CUDA_CHECK(cudaMallocManaged(&d_c, m * n * sizeof(int32_t)));
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<> dis(0, 100);
for (int i = 0; i < m * k; i++) {
d_a[i] = static_cast<int8_t>(dis(gen));
}
for (int i = 0; i < k * n; i++) {
d_b[i] = static_cast<int8_t>(dis(gen));
}
for (int i = 0; i < m * n; i++) {
d_c[i] = 0;
}
cublasdxTensor matrix_a { 0 };
cublasdxTensor matrix_b { 0 };
cublasdxTensor tile_gemm_c { 0 };
std::vector<long long int> shape_a = { m, k };
std::vector<long long int> shape_b = { k, n };
// Importantly, the shape of the output tensor is the tile size, not the global size
// the raw pointer must be offset for each CTA tile as shown in the kernel code
std::vector<long long int> shape_c = { tile_m, tile_n };
std::vector<long long int> strides_a = { 1, m };
std::vector<long long int> strides_b = { 1, k };
// while the shape of the tile gemm c is the tile size, the strides are the global size
std::vector<long long int> strides_c = { 1, m };
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(CUBLASDX_MEMORY_SPACE_GMEM,
COMMONDX_R_8I,
nullptr,
shape_a.size(),
shape_a.data(),
strides_a.data(),
&matrix_a));
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(CUBLASDX_MEMORY_SPACE_GMEM,
COMMONDX_R_8I,
nullptr,
shape_b.size(),
shape_b.data(),
strides_b.data(),
&matrix_b));
LIBMATHDX_CHECK(cublasdxCreateTensorStrided(CUBLASDX_MEMORY_SPACE_GMEM,
COMMONDX_R_32I,
nullptr,
shape_c.size(),
shape_c.data(),
strides_c.data(),
&tile_gemm_c));
cublasdxTensor acc_c { 0 };
LIBMATHDX_CHECK(cublasdxCreateTensor(h, CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C, &acc_c));
cublasdxPipeline device_pipeline { 0 };
cublasdxPipeline tile_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDevicePipeline(h,
CUBLASDX_DEVICE_PIPELINE_SUGGESTED,
LIBMATHDX_MAX_PIPELINE_DEPTH,
CUBLASDX_BLOCK_SIZE_STRATEGY_FIXED,
matrix_a,
matrix_b,
&device_pipeline));
LIBMATHDX_CHECK(cublasdxCreateTilePipeline(h, CUBLASDX_TILE_PIPELINE_DEFAULT, device_pipeline, &tile_pipeline));
std::array tensors = { matrix_a, matrix_b, tile_gemm_c, acc_c };
std::array pipelines = { device_pipeline, tile_pipeline };
LIBMATHDX_CHECK(cublasdxFinalize(tensors.size(), tensors.data(), pipelines.size(), pipelines.data()));
// size and alignment of accumulator tensor
long long int acc_size = 0;
LIBMATHDX_CHECK(cublasdxGetTensorTraitInt64(acc_c, CUBLASDX_TENSOR_TRAIT_STORAGE_BYTES, &acc_size));
long long int acc_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetTensorTraitInt64(acc_c, CUBLASDX_TENSOR_TRAIT_ALIGNMENT_BYTES, &acc_alignment));
// size and alignment of the device pipeline
long long int device_pipeline_storage_size = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_BYTES, &device_pipeline_storage_size));
long long int device_pipeline_storage_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_ALIGNMENT_BYTES, &device_pipeline_storage_alignment));
// size and alignment of the shared memory buffer (on sm_90a and above (accelerated arches with suffix 'a') should
// be 128B aligned for TMA usage)
long long int shared_memory_buffer_size = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_BUFFER_SIZE, &shared_memory_buffer_size));
long long int shared_memory_buffer_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_BUFFER_ALIGNMENT_BYTES, &shared_memory_buffer_alignment));
// block dimension for GEMM kernel launch (might be different than operator block dim)
std::array<long long int, 3> device_pipeline_block_dim = { 0, 0, 0 };
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64s(
device_pipeline, CUBLASDX_PIPELINE_TRAIT_BLOCK_DIM, 3, device_pipeline_block_dim.data()));
// size and alignment of the tile pipeline
long long int tile_pipeline_storage_size = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
tile_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_BYTES, &tile_pipeline_storage_size));
long long int tile_pipeline_storage_alignment = 0;
LIBMATHDX_CHECK(cublasdxGetPipelineTraitInt64(
tile_pipeline, CUBLASDX_PIPELINE_TRAIT_STORAGE_ALIGNMENT_BYTES, &tile_pipeline_storage_alignment));
/**
* Define a function operating on those input and output tensors.
*
* The device function output is an opaque and stateful accumulator.
*/
cublasdxDeviceFunction init_acc { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_CREATE, 1, &acc_c, 1, &tile_pipeline, &init_acc));
cublasdxDeviceFunction init_device_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_CREATE, 0, nullptr, 1, &device_pipeline, &init_device_pipeline));
cublasdxDeviceFunction init_tile_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_CREATE, 0, nullptr, 1, &tile_pipeline, &init_tile_pipeline));
cublasdxDeviceFunction destroy_acc { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 1, &acc_c, &destroy_acc));
cublasdxDeviceFunction destroy_device_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 0, nullptr, 1, &device_pipeline, &destroy_device_pipeline));
cublasdxDeviceFunction destroy_tile_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_DESTROY, 0, nullptr, 1, &tile_pipeline, &destroy_tile_pipeline));
cublasdxDeviceFunction execute_pipeline { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_EXECUTE, 1, &acc_c, 1, &tile_pipeline, &execute_pipeline));
cublasdxDeviceFunction epilogue { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunctionWithPipelines(
h, CUBLASDX_DEVICE_FUNCTION_EPILOGUE, 1, &acc_c, 1, &tile_pipeline, &epilogue));
// set the callback function for the epilogue function
LIBMATHDX_CHECK(
cublasdxSetDeviceFunctionOptionStr(epilogue, CUBLASDX_DEVICE_FUNCTION_OPTION_CALLBACK, "epilogue_callback"));
std::array copy_tensors = { acc_c, tile_gemm_c };
cublasdxDeviceFunction copy_acc_big_gmem_c { 0 };
LIBMATHDX_CHECK(cublasdxCreateDeviceFunction(
h, CUBLASDX_DEVICE_FUNCTION_COPY, copy_tensors.size(), copy_tensors.data(), ©_acc_big_gmem_c));
std::vector<cublasdxDeviceFunction> functions = { execute_pipeline,
init_acc,
init_device_pipeline,
init_tile_pipeline,
destroy_acc,
destroy_device_pipeline,
destroy_tile_pipeline,
copy_acc_big_gmem_c,
epilogue };
std::unordered_map<cublasdxDeviceFunction, std::string> function_symbols;
for (auto f : functions) {
size_t symbol_size { 0 };
LIBMATHDX_CHECK(cublasdxGetDeviceFunctionTraitStrSize(f, CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL, &symbol_size));
std::vector<char> symbol(symbol_size, '\0');
LIBMATHDX_CHECK(
cublasdxGetDeviceFunctionTraitStr(f, CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL, symbol.size(), symbol.data()));
printf("Device function %lld: symbol: %s\n", static_cast<long long int>(f), symbol.data());
function_symbols[f] = std::string(symbol.data());
}
/**
* Compile the device functions to LTOIR
*/
commondxCode code { 0 };
LIBMATHDX_CHECK(commondxCreateCode(&code));
LIBMATHDX_CHECK(
commondxSetCodeOptionInt64s(code, COMMONDX_OPTION_TARGET_SM, target_sm_array.size(), target_sm_array.data()));
LIBMATHDX_CHECK(cublasdxFinalizeDeviceFunctions(code, functions.size(), functions.data()));
/**
* Extract the LTOIR
*/
size_t lto_size = 0;
LIBMATHDX_CHECK(commondxGetCodeLTOIRSize(code, <o_size));
std::vector<char> lto(lto_size);
LIBMATHDX_CHECK(commondxGetCodeLTOIR(code, lto_size, lto.data()));
printf("Generated LTOIR for GEMM pipeline with epilogue callback, %zu bytes\n", lto.size());
const char kernel_template[] = R"(
struct my_user_data {
void* ptr;
};
struct libmathdx_tensor_0s_0s { void* ptr; };
struct libmathdx_pipeline { void* ptr; };
#define M %d
#define N %d
#define K %d
#define tile_m %d
#define tile_n %d
#define tile_k %d
#define acc_name libmathdx_tensor_0s_0s
#define ga_name libmathdx_tensor_0s_0s
#define gb_name libmathdx_tensor_0s_0s
#define gc_name libmathdx_tensor_0s_0s
#define device_pipeline_name libmathdx_pipeline
#define tile_pipeline_name libmathdx_pipeline
constexpr unsigned acc_size = %lld;
constexpr unsigned acc_alignment = %lld;
constexpr unsigned block_size = %lld;
constexpr unsigned smem_alignment = %lld;
constexpr unsigned tile_pipeline_size = %lld;
constexpr unsigned tile_pipeline_alignment = %lld;
#define execute_pipeline_acc %s
#define copy_acc_gc %s
#define create_dev_pipe %s
#define create_tile_pipe %s
#define create_acc %s
#define destroy_dev_pipe %s
#define destroy_tile_pipe %s
#define destroy_acc %s
#define epilogue %s
using C_VALUE_TYPE = signed int;
extern "C" __device__ void execute_pipeline_acc(tile_pipeline_name, acc_name);
extern "C" __device__ void copy_acc_gc(acc_name, gc_name);
extern "C" __device__ void create_dev_pipe(device_pipeline_name, ga_name, gb_name);
extern "C" __device__ void create_tile_pipe(device_pipeline_name, tile_pipeline_name, char*, int*, int*);
extern "C" __device__ void create_acc(tile_pipeline_name, acc_name);
extern "C" __device__ void destroy_dev_pipe(device_pipeline_name);
extern "C" __device__ void destroy_tile_pipe(tile_pipeline_name);
extern "C" __device__ void destroy_acc(acc_name);
extern "C" __device__ void epilogue(tile_pipeline_name, acc_name, void* user_data);
extern "C" __device__ void epilogue_callback(acc_name accumulator, my_user_data* user_data) {
auto gc = gc_name { user_data->ptr };
copy_acc_gc(accumulator, gc);
}
// Create the device pipeline
extern "C" __global__ void create_device_pipeline(void* device_pipeline_ptr, void* ga_storage, void* gb_storage) {
if(threadIdx.x == 0) {
auto ga = ga_name { ga_storage };
auto gb = gb_name { gb_storage };
auto device_pipeline = device_pipeline_name { device_pipeline_ptr };
create_dev_pipe(device_pipeline, ga, gb);
}
}
// Perform the GEMM
extern "C" __launch_bounds__(block_size, 1) __global__ void gemm(void* device_pipeline_ptr, void* gc_storage)
{
int row_id = blockIdx.x;
int col_id = blockIdx.y;
C_VALUE_TYPE* gmem_c = reinterpret_cast<C_VALUE_TYPE*>(gc_storage) + ((row_id * tile_m) + (col_id * tile_n) * M);
// Allocate dynamic shared memory for tiles
extern __shared__ __align__(smem_alignment) char smem[];
// Allocate local memory
alignas(acc_alignment) char acc_storage[acc_size];
alignas(tile_pipeline_alignment) char tile_pipeline_storage[tile_pipeline_size];
// Create opaque types
auto acc = acc_name { acc_storage };
auto device_pipeline = device_pipeline_name { device_pipeline_ptr };
auto tile_pipeline = tile_pipeline_name { tile_pipeline_storage };
// Perform the GEMM
create_tile_pipe(device_pipeline, tile_pipeline, smem, &row_id, &col_id);
create_acc(tile_pipeline, acc);
execute_pipeline_acc(tile_pipeline, acc);
auto user_data = my_user_data { gmem_c };
epilogue(tile_pipeline, acc, &user_data);
// Destroy the accumulator and the tile pipeline
destroy_acc(acc);
destroy_tile_pipe(tile_pipeline);
}
// Destroy the device pipeline
extern "C" __global__ void destroy_device_pipeline(void* device_pipeline_ptr) {
if(threadIdx.x == 0) {
auto device_pipeline = device_pipeline_name { device_pipeline_ptr };
destroy_dev_pipe(device_pipeline);
}
}
)";
std::string cpp = strprintf(kernel_template,
m,
n,
k,
tile_m,
tile_n,
tile_k,
acc_size,
acc_alignment,
device_pipeline_block_dim[0],
shared_memory_buffer_alignment,
tile_pipeline_storage_size,
tile_pipeline_storage_alignment,
function_symbols[execute_pipeline].data(),
function_symbols[copy_acc_big_gmem_c].data(),
function_symbols[init_device_pipeline].data(),
function_symbols[init_tile_pipeline].data(),
function_symbols[init_acc].data(),
function_symbols[destroy_device_pipeline].data(),
function_symbols[destroy_tile_pipeline].data(),
function_symbols[destroy_acc].data(),
function_symbols[epilogue].data());
std::vector<char> cubin = compile_and_link(cpp, lto, target_sm);
CUmodule module {};
CUfunction kernel, kernel2, kernel3 {};
CUDA_CHECK(cudaSetDevice(0));
CU_CHECK(cuModuleLoadDataEx(&module, cubin.data(), 0, 0, 0));
CU_CHECK(cuModuleGetFunction(&kernel, module, "create_device_pipeline"));
CU_CHECK(cuModuleGetFunction(&kernel2, module, "gemm"));
CU_CHECK(cuModuleGetFunction(&kernel3, module, "destroy_device_pipeline"));
// create the device pipeline using trait
char* device_pipeline_ptr {};
CUDA_CHECK(cudaMalloc(&device_pipeline_ptr, device_pipeline_storage_size));
CUDA_CHECK(cudaMemset(device_pipeline_ptr, 0, device_pipeline_storage_size));
{
std::vector<void*> kernel_args = { reinterpret_cast<void*>(&device_pipeline_ptr),
reinterpret_cast<void*>(&d_a),
reinterpret_cast<void*>(&d_b) };
CU_CHECK(cuLaunchKernel(kernel, 1, 1, 1, 1, 1, 1, 0, nullptr, kernel_args.data(), nullptr));
CUDA_CHECK(cudaDeviceSynchronize());
}
{
CU_CHECK(cuFuncSetAttribute(
kernel2, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, static_cast<int>(shared_memory_buffer_size)));
std::vector<void*> kernel_args = { reinterpret_cast<void*>(&device_pipeline_ptr),
reinterpret_cast<void*>(&d_c) };
CU_CHECK(cuLaunchKernel(kernel2,
static_cast<unsigned int>(m / tile_m),
static_cast<unsigned int>(n / tile_n),
1,
static_cast<unsigned int>(device_pipeline_block_dim[0]),
static_cast<unsigned int>(device_pipeline_block_dim[1]),
static_cast<unsigned int>(device_pipeline_block_dim[2]),
static_cast<unsigned int>(shared_memory_buffer_size),
nullptr,
kernel_args.data(),
nullptr));
CUDA_CHECK(cudaDeviceSynchronize());
}
{
std::vector<void*> kernel_args = { reinterpret_cast<void*>(&device_pipeline_ptr) };
CU_CHECK(cuLaunchKernel(kernel3, 1, 1, 1, 1, 1, 1, 0, nullptr, kernel_args.data(), nullptr));
CUDA_CHECK(cudaDeviceSynchronize());
}
std::vector<int32_t> h_c_ref(m * n, 0);
for (int l = 0; l < k; l++) {
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
h_c_ref[i + m * j] += d_a[i + m * l] * d_b[l + k * j];
}
}
}
for (int i = 0; i < m * n; i++) {
if (h_c_ref[i] != d_c[i]) {
printf("Error at %d: h_c_ref[%d] = %d, d_c[%d] = %d\n", i, i, h_c_ref[i], i, d_c[i]);
abort();
}
}
printf("Successfully ran the kernel\n");
/**
* Destroy handles
*/
CU_CHECK(cuModuleUnload(module));
CUDA_CHECK(cudaFree(device_pipeline_ptr));
CUDA_CHECK(cudaFree(d_a));
CUDA_CHECK(cudaFree(d_b));
CUDA_CHECK(cudaFree(d_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(matrix_a));
LIBMATHDX_CHECK(cublasdxDestroyTensor(matrix_b));
LIBMATHDX_CHECK(cublasdxDestroyTensor(acc_c));
LIBMATHDX_CHECK(cublasdxDestroyTensor(tile_gemm_c));
LIBMATHDX_CHECK(cublasdxDestroyPipeline(device_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyPipeline(tile_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init_acc));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init_device_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(init_tile_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy_acc));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy_device_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(destroy_tile_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(execute_pipeline));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(epilogue));
LIBMATHDX_CHECK(cublasdxDestroyDeviceFunction(copy_acc_big_gmem_c));
LIBMATHDX_CHECK(commondxDestroyCode(code));
LIBMATHDX_CHECK(cublasdxDestroyDescriptor(h));
}
API reference#
-
LIBMATHDX_RUNTIME#
Sentinel value used to mark runtime-provided shapes/strides.
Pass LIBMATHDX_RUNTIME in cublasdxCreateTensorStrided to indicate that the corresponding dimension or stride will be provided at runtime.
-
LIBMATHDX_MAX_PIPELINE_DEPTH#
Special pipeline depth requesting the maximal supported depth.
When creating pipelines with cublasdxCreateDevicePipeline pass LIBMATHDX_MAX_PIPELINE_DEPTH to select the maximum depth based on available shared memory.
-
LIBMATHDX_NONE#
Sentinel handle indicating that no descriptor is provided.
Some APIs accept LIBMATHDX_NONE in lieu of a cublasdxDescriptor to operate without a descriptor.
-
typedef long long int cublasdxDescriptor#
A handle to a cuBLASDx descriptor.
Equivalent to
using GEMM = ...in cuBLASDx CUDA C++.
-
typedef long long int cublasdxTensor#
A handle to an opaque device tensor.
-
typedef long long int cublasdxDeviceFunction#
A handle to a device function. A device function operators on tensors described by cublasdxTensor.
-
typedef long long int cublasdxPipeline#
A handle to an opaque device pipeline.
-
enum cublasdxApi_t#
Type of cublasDx API.
Handling problems with default or custom/dynamic leading dimensions. Check cublasdx::LeadingDimension operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#leadingdimension-operator)
Values:
-
enumerator CUBLASDX_API_SMEM#
Use API for problems with default leading dimensions. Function API is defined by its signature:
void (value_type_c* alpha, value_type_a* smem_a, value_type_b* smem_b, value_type_c* beta, value_type_c* smem_c)wheresmem_a,smem_bandsmem_care pointers to value of type given by the CUBLASDX_TRAIT_VALUE_TYPE a, b, and c trait.smem_a,smem_bandsmem_cmust be shared memory pointers.alphaandbetaare pointers to value of type CUBLASDX_TRAIT_VALUE_TYPE c.
Note that complex numbers must be over-aligned.
The function is
extern "C"and the symbol name can be queried using CUBLASDX_TRAIT_SYMBOL_NAME. See https://docs.nvidia.com/cuda/cublasdx/api/methods.html#shared-memory-api and in particular thePointer APIsection.
-
enumerator CUBLASDX_API_SMEM_DYNAMIC_LD#
Use API for problems with custom / dynamic leading dimensions. Function API is defined by its signature:
void (value_type_c alpha, value_type_a* smem_a, unsigned* lda, value_type_b *smem_b, unsigned* ldb, value_type_c* beta, value_type_c* smem_c, unsigned *ldc)wheresmem_a,smem_bandsmem_care pointers to value of type given by the CUBLASDX_TRAIT_VALUE_TYPE a, b, and c trait.smem_a,smem_bandsmem_cmust be shared memory pointers.alphaandbetaare pointers to value of type CUBLASDX_TRAIT_VALUE_TYPE c trait.lda,ldbandldcare pointers to unsigned 32 bits integer (unsigned)
Note that complex numbers must be over-aligned.
The function is
extern "C"and the symbol name can be queried using CUBLASDX_TRAIT_SYMBOL_NAME. See https://docs.nvidia.com/cuda/cublasdx/api/methods.html#shared-memory-api and in particular thePointer API, which allows providing runtime/dynamic leading dimensions for matrices A, B, and Csection.
-
enumerator CUBLASDX_API_TENSORS#
Use Tensor API. Function API is defined by the input and output tensors specified when calling cublasdxCreateDeviceFunction. The device functions are
extern "C". Tensors are trivial and passed by value. Scalars arevoid*.
-
enumerator CUBLASDX_API_SMEM#
-
enum cublasdxType_t#
Type of computation data.
Check cubladx::Type operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#type-operator)
Values:
-
enumerator CUBLASDX_TYPE_REAL#
Use for real matmuls
-
enumerator CUBLASDX_TYPE_COMPLEX#
Use for complex matmuls
-
enumerator CUBLASDX_TYPE_REAL#
-
enum cublasdxTransposeMode_t#
Tensor transpose mode.
The transpose mode depends on cubladx::TransposeMode operator which is deprecated since cublasDx 0.2.0 and might be removed in future versions of mathDx libraries
Check cubladx::TransposeMode operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#transposemode-operator)
Values:
-
enumerator CUBLASDX_TRANSPOSE_MODE_NON_TRANSPOSED#
Use matrix as-is in the matmul
-
enumerator CUBLASDX_TRANSPOSE_MODE_TRANSPOSED#
Use transposed matrix in the matmul
-
enumerator CUBLASDX_TRANSPOSE_MODE_CONJ_TRANSPOSED#
Use transposed and conjugate matrix in the matmul
-
enumerator CUBLASDX_TRANSPOSE_MODE_NON_TRANSPOSED#
-
enum cublasdxArrangement_t#
Data arrangement mode.
Defines data arrangements in tensors’ taking part in the calculation.
Check cubladx::TransposeMode operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#arrangement-operator)
Values:
-
enumerator CUBLASDX_ARRANGEMENT_COL_MAJOR#
Data is considered column-major
-
enumerator CUBLASDX_ARRANGEMENT_ROW_MAJOR#
Data is considered row-major
-
enumerator CUBLASDX_ARRANGEMENT_COL_MAJOR#
-
enum cublasdxFunction_t#
BLAS function.
Sets the BLAS function that will be executed.
Check cubladx::Function operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#function-operator)
Values:
-
enumerator CUBLASDX_FUNCTION_MM#
Matrix-multiply
-
enumerator CUBLASDX_FUNCTION_MM#
-
enum cublasdxOperatorType_t#
cublasDx operators
The set of supported cublasDx operators.
Check cublaDx description operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#function-operator)
Check cublasDx execution operator documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/operators.html#execution-operators)
Values:
-
enumerator CUBLASDX_OPERATOR_FUNCTION#
Operator data type: cublasdxFunction_t. Operator definition: required
-
enumerator CUBLASDX_OPERATOR_SIZE#
Operator data type: long long int * 3. Expected content:
<M, N, K>problem sizes. Operator definition: required
-
enumerator CUBLASDX_OPERATOR_TYPE#
Operator data type: cublasdxType_t. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_PRECISION#
Operator data type: commondxPrecision_t * 3. Expected content:
<A, B, C>precisions. Operator definition: required
-
enumerator CUBLASDX_OPERATOR_SM#
Operator data type: long long int. Expected content: 700 (Volta), 800 (Ampere), …. Operator definition: required
-
enumerator CUBLASDX_OPERATOR_EXECUTION#
Operator data type: commondxExecution_t. Operator definition: required
-
enumerator CUBLASDX_OPERATOR_BLOCK_DIM#
Operator data type: long long int * 3. Expected content: <x, y, z> block dimensions. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_LEADING_DIMENSION#
Operator data type: long long int * 3. Expected content: <LDA, LDB, LDC> leading dimensions. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_TRANSPOSE_MODE#
Operator data type: cublasdxTransposeMode_t * 2. Expected content: <A, B> transpose modes. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_API#
Operator data type: cublasdxApi_t. Operator definition: required
-
enumerator CUBLASDX_OPERATOR_ARRANGEMENT#
Operator data type: cublasdxArrangement_t * 3. Expected content: <A, B, C> data arrangements. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_ALIGNMENT#
Operator data type: long long int * 3. Expected content: <AAlign, BAlign, CAlign> tensors’ alignments. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_STATIC_BLOCK_DIM#
Operator data type: long long int. Expected content: 1, to enable cublasdx::experimental::StaticBlockDim. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_ENABLE_INPUT_STREAMING#
Operator data type: long long int. Expected content: 1, to enable cublasdx::EnableInputStreaming. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_WITH_PIPELINE#
Operator data type: long long int. Expected content: 1, to enable cublasdx::WithPipeline. Operator definition: optional
-
enumerator CUBLASDX_OPERATOR_FUNCTION#
-
enum cublasdxTraitType_t#
cublasDx traits
The set of supported types of traits that can be accessed from finalized sources that use cublasDx.
Check cublasDx Execution Block Traits documentation for more details (https://docs.nvidia.com/cuda/cublasdx/api/traits.html#block-traits)
Values:
-
enumerator CUBLASDX_TRAIT_VALUE_TYPE#
Trait data type: commondxValueType_t * 3. Expected content: <A, B, C> types.
-
enumerator CUBLASDX_TRAIT_SIZE#
Trait data type: long long int * 3. Expected content: <M, N, K> problem sizes.
-
enumerator CUBLASDX_TRAIT_BLOCK_SIZE#
Trait data type: long long int. Expected content: multiplication result of block dimensions (x * y * z).
-
enumerator CUBLASDX_TRAIT_BLOCK_DIM#
Trait data type: long long int * 3. Expected content: <x, y, z> block dimension.
-
enumerator CUBLASDX_TRAIT_LEADING_DIMENSION#
Trait data type: long long int * 3. Expected content: <LDA, LDB, LDC> leading dimensions.
-
enumerator CUBLASDX_TRAIT_SYMBOL_NAME#
Trait data type: C-string
-
enumerator CUBLASDX_TRAIT_ARRANGEMENT#
Trait data type: cublasdxArrangement_t * 3. Expected content: <A, B, C> data arrangements.
-
enumerator CUBLASDX_TRAIT_ALIGNMENT#
Trait data type: long long int * 3. Expected content: <AAlign, BAlign, CAlign> tensors’ alignments, in bytes.
-
enumerator CUBLASDX_TRAIT_SUGGESTED_LEADING_DIMENSION#
Trait data type: long long int * 3. Expected content: <LDA, LDB, LDC>.
-
enumerator CUBLASDX_TRAIT_SUGGESTED_BLOCK_DIM#
Trait data type: long long int * 3. Expected content: <X, Y, Z>.
-
enumerator CUBLASDX_TRAIT_MAX_THREADS_PER_BLOCK#
Trait data type: long long int. Expected content: the product of three elements in block dimension.
-
enumerator CUBLASDX_TRAIT_VALUE_TYPE#
-
enum cublasdxTensorType_t#
cuBLASDx desired tensor type
Tensor types are opaque (layout is unspecified), non-owning, and defined by
Memory space (global, shared or register memory)
Size & alignment (in bytes)
Tensor’s representation in-memory and in-device depends on their memory space. Shared & register tensors are defined as
struct tensor { void* ptr; }
where
ptrpoints to the associated data.Global memory tensors have an associated runtime leading dimension (64b signed integer), and their representation is
struct tensor { void* ptr; long long int strides[1]; }
where
ptrpoints to the associated data andstrides[0]is the leading dimension.In either case,
ptrmust point to some storage (with appropriate size and alignment, see below) and is not owning. The user is expected to keep memory allocated beyond any use of the tensor.strides[1]should be a signed, 64bit integer (long long) equal to the leading dimension of the global memory tensor. The leading dimension is the number of elements between two successive rows or columns (not bytes), depending on the context.All tensor APIs take their argument by value (not by pointer) and expect the struct to be passed as-is on the stack.
Each opaque tensor type is uniquely identified by a unique ID and name, see cublasdxTensorTrait_t .
Values:
-
enumerator CUBLASDX_TENSOR_SMEM_A#
A shared memory tensor for
A, in simple row or column layout In memory representation:struct { void* ptr; }withptra shared memory pointer. Corresponds to cuBLASDxmake_tensor(..., get_layout_smem_a());
-
enumerator CUBLASDX_TENSOR_SMEM_B#
A shared memory tensor for
B, in simple row or column layout. In memory representation:struct { void* ptr; }withptra shared memory pointer. Corresponds to cuBLASDxmake_tensor(..., get_layout_smem_b());
-
enumerator CUBLASDX_TENSOR_SMEM_C#
A shared memory tensor for
C, in simple row or column layout. In memory representation:struct { void* ptr; }withptra shared memory pointer. Corresponds to cuBLASDxmake_tensor(..., get_layout_smem_c());
-
enumerator CUBLASDX_TENSOR_SUGGESTED_SMEM_A#
A shared memory tensor for
A, in unspecified (could be swizzled, padded, etc) layout. In memory representation:struct { void* ptr; }withptra shared memory pointer. Corresponds to cuBLASDxmake_tensor(..., suggest_layout_smem_a());
-
enumerator CUBLASDX_TENSOR_SUGGESTED_SMEM_B#
A shared memory tensor for
B, in unspecified (could be swizzled, padded, etc) layout. In memory representation:struct { void* ptr; }withptra shared memory pointer. Corresponds to cuBLASDxmake_tensor(..., suggest_layout_smem_b());
-
enumerator CUBLASDX_TENSOR_SUGGESTED_SMEM_C#
A shared memory tensor for
C, in unspecified (could be swizzled, padded, etc) layout. In memory representation:struct { void* ptr; }withptra shared memory pointer. Corresponds to cuBLASDxmake_tensor(..., suggest_layout_smem_c());
-
enumerator CUBLASDX_TENSOR_SUGGESTED_RMEM_C#
A register tensor for
C, in unspecified layout. In memory representation:struct { void* ptr; }with ptr a stack (aka local or thread-private) memory pointer. Corresponds to cuBLASDxsuggest_accumulator().make_accumulator_fragment();
-
enumerator CUBLASDX_TENSOR_GMEM_A#
A global memory view for
A(typically a tile of a larger matrix) in row or column-major format, with a runtime leading dimension. In memory representation:struct { void* ptr; long long int[1] strides; }withptra global memory pointer andstrides[0]the leading dimension. Corresponds to cuBLASDxmake_tensor(a, get_layout_gmem_a(lda));
-
enumerator CUBLASDX_TENSOR_GMEM_B#
A global memory view for
B(typically a tile of a larger matrix) in row or column-major format, with a runtime leading dimension. In memory representation:struct { void* ptr; long long int[1] strides; }withptra global memory pointer andstrides[0]the leading dimension. Corresponds to cuBLASDxmake_tensor(a, get_layout_gmem_b(ldb));
-
enumerator CUBLASDX_TENSOR_GMEM_C#
A global memory view for
C(typically a tile of a larger matrix) in row or column-major format, with a runtime leading dimension. In memory representation:struct { void* ptr; long long int[1] strides; }withptra global memory pointer andstrides[0]the leading dimension. Corresponds to cuBLASDxmake_tensor(a, get_layout_gmem_c(ldc));
-
enumerator CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C#
An opaque, stateful accumulator for
C, in unspecified layout and in an unspecified memory space. Must be explicitly initialized and destroyed. In memory representation:struct { void* ptr; }with ptr a stack (aka local or thread-private) memory pointer pointing to an opaque state.ptrshould not be assumed to point to any specific data. Corresponds to cuBLASDxsuggest_accumulator();
-
enumerator CUBLASDX_TENSOR_RMEM_C#
A register tensor for
C, in unspecified layout. In memory representation:struct { void* ptr; }with ptr a stack (aka local, aka thread-private) memory pointer. Corresponds to cuBLASDxget_accumulator().make_accumulator_fragment();
-
enumerator CUBLASDX_TENSOR_ACCUMULATOR_C#
An opaque, stateful accumulator for
C, in unspecified layout and in an unspecified memory space. Must be explicitly initialized and destroyed. In memory representation:struct { void* ptr; }where ptr is a stack (aka local or thread-private) memory pointer pointing to an opaque state.ptrshould not be assumed to point to any specific data. Corresponds to cuBLASDxget_accumulator();
-
enum cublasdxTensorTrait_t#
Tensor traits, used to query informations.
Values:
-
enumerator CUBLASDX_TENSOR_TRAIT_STORAGE_BYTES#
The size of the underlying storage, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_TENSOR_TRAIT_ALIGNMENT_BYTES#
The alignment of the underlying storage, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_TENSOR_TRAIT_UID#
The tensor type UID. Tensor types with the same UID are identical and can be passed through various cuBLASDx device functions. UIDs are only well defined within a process. Note: This trait has been deprecated. use
CUBLASDX_TENSOR_TRAIT_OPAQUE_NAMEinstead to identify device tensors. Trait data type: long long int.
-
enumerator CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME#
A human readable but unspecified C-string representing the opaque tensor type name. Names are stable and unique per tensor type, and tensor types with the same name can be used interchangeably. Opaque names are not C++ type name identifiers. Trait data type: C-string.
-
enumerator CUBLASDX_TENSOR_TRAIT_LOGICAL_SIZE#
The logical number of elements owned by the tensor. Note this is not related to the size in memory. Trait data type: int.
-
enumerator CUBLASDX_TENSOR_TRAIT_MEMORY_SPACE#
Returns the memory space of the underlying data of this tensor. Trait data type: cublasdxMemorySpace.
-
enumerator CUBLASDX_TENSOR_TRAIT_STORAGE_BYTES#
-
enum cublasdxDeviceFunctionTrait_t#
Device function traits, used to query informations.
Values:
-
enumerator CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL#
The symbol name of the device function. Trait data type: C-string
-
enumerator CUBLASDX_DEVICE_FUNCTION_TRAIT_SYMBOL#
-
enum cublasdxDeviceFunctionType_t#
Device functions supported by the library.
Values:
-
enumerator CUBLASDX_DEVICE_FUNCTION_EXECUTE#
Execute the device function (matmul).
When the input is a tile pipeline, and the output is a accumulator tensor, the device function API is
execute(tile_pipeline, C)which computesC += A x B.When the output is a register tensor, the device function API is
execute(A, B, C)which computesC += A x B.When the output is a shared memory tensor, the device function API is
execute(alpha, A, B, beta, C)which computesC = alpha A x B + beta C.A,BandCare tensors, while alpha and beta are scalars with the same type ofC(passed byvoid*pointers), and tile_pipeline is a reference to tiles of A and B.Different
executegenerated from distinctcublasdxDescriptorare generally different and cannot be used interchangeably, even with an identical set of input and output tensors.cublasdxCreateDeviceFunction must be called with three tensors:
tile_pipeline, an instance ofA, an instance ofB, an instance ofC, an instance of
The resulting function has the following device API:
void execute(void* alpha, TA A, TB B, void* beta, TC C)whenCis a shared memory tensor,void execute(TA A, TB B, TC C)whenCis a register memory tensors.void execute(TP tile_pipeline, TC C)whenCis an accumulator tensor.
The names for
TA,TB,TCandTPcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .
-
enumerator CUBLASDX_DEVICE_FUNCTION_COPY#
Copies from one tensor to another.
copy(S, D)copies fromStoD.Different
copygenerated from distinctcublasdxDescriptorare in general different and cannot be used interchangeably, even with identical input and output tensors.cublasdxCreateDeviceFunction must be called with two tensors:
S, an instance ofD, an instance of
SandDcan be in different memory spaces but must correspond to the same A, B or C matrix.The resulting function has one of the following device API:
- `void copy(TS S, TD D)` if a BLAS descriptor is provided, - `void copy(int* tid, TS S, TD D)` if \ref LIBMATHDX_NONE is used in lieu of a BLAS descriptor. In this case, `tid` should be the thread ID (e.g. `threadIdx.x`).
The names for
TSandTDcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .When the BLAS descriptor was created with CUBLASDX_OPERATOR_WITH_PIPELINE, a copy between a register-fragment tensor (CUBLASDX_TENSOR_SUGGESTED_RMEM_C / CUBLASDX_TENSOR_RMEM_C) and an SMEM/GMEM tensor must be created with cublasdxCreateDeviceFunctionWithPipelines, passing the tile pipeline. The tile pipeline supplies the register partitioning (a plain rmem fragment copy is rejected at compile time in pipelined mode). The resulting device API gains the tile pipeline as its first argument:
void copy(TilePipeline tp, TS S, TD D).
-
enumerator CUBLASDX_DEVICE_FUNCTION_COPY_WAIT#
Wait on all previously issued copies to complete.
wait_all()waits on all previously issued copies to complete.Different
wait_allfrom distinctcublasdxDescriptorare identical and may used interchangeably. They will have the same symbol name and implementation.cublasdxCreateDeviceFunction must be called without any tensors.
The resulting function has the following device API:
void copy_wait()LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
-
enumerator CUBLASDX_DEVICE_FUNCTION_CLEAR#
Zeroes out a tensor.
clear(C)zeroes outC.Different
cleargenerated from distinctcublasdxDescriptorare in general different and cannot be used interchangeably, even with identical input and output tensors.cublasdxCreateDeviceFunction must be called with one tensors:
C, an instance of
The resulting function has the following device API:
void clear(TC C)The name for
TCcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
-
enumerator CUBLASDX_DEVICE_FUNCTION_AXPBY#
Computes
D = alpha * C + beta * D.The
cublasdxDescriptormaybe me LIBMATHDX_NONE and is effectively ignored.cublasdxCreateDeviceFunction must be called with two tensors:
C, an instance ofD, an instance of
The resulting function has the following device API:
void axpby(void* alpha, TC C, void* beta, TD D)whereCandDare tensors andalpha,betaare pointers to scalars, wherealphahas type ofTCandbetathe type ofTD.
The name for
TCandTDcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .
-
enumerator CUBLASDX_DEVICE_FUNCTION_MAP_IDX2CRD#
Iterates over a 2D tensor.
Supports iterative logical access to different layouts in underlying CuTe Tensor.
cublasdxCreateDeviceFunction must be called with one tensor:
A, an instance of
The resulting function has the following device API:
void map_idx2crd(T A, int* lin_idx, int* i, int* j, void* ptr)where inputAis a tensor and inputlin_idx, outputi, outputj, and outputptrare pointers of type, int, int, int, and void respectively.
The arguement defitions are as follows:
T Ais the tensor to iterate over.int* lin_idxis the provided linearized index.int* iis the returned row index value for the physical tensor element stored inptr.int* jis the returned column index value for the physical tensor element stored inptr.void* ptris the pointer to the physical tensor element.
The name for
Tcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
The following six functions are based on this link: https://docs.nvidia.com/cuda/cublasdx/api/other_tensors.html#partitioner
-
enumerator CUBLASDX_DEVICE_FUNCTION_MAP_IDX2CRD_PARTITIONER#
Iterates over a 2D tensor.
Supports iterative logical access to different layouts in underlying implicit CuTe Tensor.
cublasdxCreateDeviceFunction must be called with one tensor:
A, an instance of
The resulting function has the following device API:
void map_idx2crd_partitioner(int* lin_idx, int* i, int* j)where inputlin_idx, outputi, and outputjare pointers of type, int, int, int, respectively.
The argument definitions are as follows:
int* lin_idxis the provided linearized index.int* iis the returned row index value for the physical tensor element stored inptr.int* jis the returned column index value for the physical tensor element stored inptr.
LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
The name for
Tcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .
-
enumerator CUBLASDX_DEVICE_FUNCTION_MAP_CRD2IDX#
Random access to a specific location in a 2D tensor.
Supports random logical access to different layouts in underlying CuTe tensor.
cublasdxCreateDeviceFunction must be called with one tensor:
A, an instance of
The resulting function has the following device API:
void map_crd2idx(T A, int* i, int* j, int* lin_idx, void* ptr)where inputAis a tensor and inputi, inputj, outputlin_idx, and outputptrare pointers of type, int, int, int, and void respectively.
The argument definitions are as follows:
T Ais the tensor to iterate over.int* iis the provided row index value.int* jis the provided column index value.int* lin_idxis the returned linearized index offset for the physical tensor element stored inptr.void* ptris the pointer to the physical tensor element.
LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
The name for
Tcan be retreived using CUBLASDX_TENSOR_TRAIT_OPAQUE_NAME .
-
enumerator CUBLASDX_DEVICE_FUNCTION_IS_THREAD_ACTIVE#
Returns true if the current thread is part of the GEMM execution.
cublasdxCreateDeviceFunction must be called with one tensor:
A, an instance of
The resulting function has the following device API:
void is_thread_active(void* yes_or_no)where outputyes_or_nois a pointer of type int.
LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
-
enumerator CUBLASDX_DEVICE_FUNCTION_IS_PREDICATED#
Returns true if any threads within the active BlockDim set of threads are predicated, ie if it does not participate in the GEMM calculations.
cublasdxCreateDeviceFunction must be called with one tensor:
A, an instance of
The resulting function has the following device API:
void is_predicated(void* yes_or_no)where outputyes_or_nois a pointer of type int.
LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
-
enumerator CUBLASDX_DEVICE_FUNCTION_IS_INDEX_IN_BOUNDS#
Determines if memory access is in bounds.
cublasdxCreateDeviceFunction must be called with one tensor:
A, an instance of
The resulting function has the following device API:
void is_index_in_bounds(int* lin_idx, void* yes_or_no)where inputlin_idxand outputyes_or_noare pointers of type int.
LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
-
enumerator CUBLASDX_DEVICE_FUNCTION_CREATE#
Initialize an opaque stateful tensor.
cublasdxCreateDeviceFunction must be called with one tensor:
C, an instance of
The resulting function has the following device API:
void create(T)whereTis the tensor to initialize.void create(DP, TA, TB)whereDPis the device pipeline,TAis the tensor A, andTBis the tensor B.void create(DP, TP, char* smem, int* idx, int* idy)whereDPis the device pipeline,TPis the tile pipeline,smemis the shared memory pointer,idxis offset(s) in the first dimension,idyis offset(s) in the second dimension. The dimension count is determined by the tensor rank.
-
enumerator CUBLASDX_DEVICE_FUNCTION_DESTROY#
Destroys an opaque stateful tensor.
cublasdxCreateDeviceFunction must be called with one tensor:
C, an instance of
The resulting function has the following device API:
void destroy(T)whereTis the tensor to destroy.
-
enumerator CUBLASDX_DEVICE_FUNCTION_RESET#
Resets an opaque stateful tile pipeline.
cublasdxCreateDeviceFunction must be called with either one or two pipelines:
DP, an instance ofTP, an instance of
The resulting function has the following device APIs depending on the number of pipelines:
void reset(DP, TP, int* idx, int* idy)whereDPis the device pipeline,TPis the tile pipeline,idxis offsets in the first dimension, andidyis offsets in the second dimension. This function is only valid for 3D input tensors.void reset(TP)whereTPis the tile pipeline. This function is only valid for 2D input tensors.
-
enumerator CUBLASDX_DEVICE_FUNCTION_EPILOGUE#
Epilogue callback function for cublasdx pipelining API.
cublasdxCreateDeviceFunction must be called with one pipeline, one accumulator tensor:
pipeline, an instance ofaccumulator_tensor, an instance of
The resulting function has the following device API:
void epilogue(TP, TC, void*)whereTPis the pipeline,TCis the accumulator tensor, andvoid*is any user required external state to be passed to the callback function. If no external state is required,void*can beNULL. The callback function is expected to be defined by the user and passed as a C-string to cublasdxSetDeviceFunctionOptionStr.
LIBMATHDX_NONE may be used in lieu of a BLAS descriptor, and has the same effect.
-
enumerator CUBLASDX_DEVICE_FUNCTION_FINISH_ACCUMULATION#
Finalizes a reusable accumulator after pipelined execution.
Must be called once after the last CUBLASDX_DEVICE_FUNCTION_EXECUTE into a given accumulator and before its results are read by any non-epilogue path (e.g. CUBLASDX_DEVICE_FUNCTION_COPY from the accumulator). The CUBLASDX_DEVICE_FUNCTION_EPILOGUE path performs this finalization internally, so it must not be combined with an explicit finish_accumulation on the same accumulator.
cublasdxCreateDeviceFunction must be called with one accumulator tensor and no pipelines:
C, an instance of
The resulting function has the following device API:
void finish_accumulation(TC)whereTCis the accumulator tensor.
-
enumerator CUBLASDX_DEVICE_FUNCTION_EXECUTE#
-
enum cublasdxTensorOption_t#
Tensor options.
Values:
-
enumerator CUBLASDX_TENSOR_OPTION_ALIGNMENT_BYTES#
The alignment of the underlying storage, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_TENSOR_OPTION_ALIGNMENT_BYTES#
-
enum cublasdxDeviceFunctionOption_t#
Device function options.
Values:
-
enumerator CUBLASDX_DEVICE_FUNCTION_OPTION_SYMBOL_NAME#
Specify an optional symbol name for the device function. Trait data type: const char*
-
enumerator CUBLASDX_DEVICE_FUNCTION_OPTION_COPY_ALIGNMENT#
Specify an optional alignment for copy and copy_fragment functions. Default is taken from the input tensors. Trait data type: long long int.
-
enumerator CUBLASDX_DEVICE_FUNCTION_OPTION_CALLBACK#
Specify a callback for the device function. This option is only supported for epilogue device function, and is required. The callback function signature is strictly defined as
callback(<internal_tensor_type> accumulator, void* user_data).<internal_tensor_type>details can be found in CUBLASDX_TENSOR_SUGGESTED_ACCUMULATOR_C and CUBLASDX_TENSOR_ACCUMULATOR_C . Trait data type: const char*
-
enumerator CUBLASDX_DEVICE_FUNCTION_OPTION_NUM_THREADS#
Specify the number of threads for the operation. This option is only supported for copy functions without BLAS descriptors (aka using LIBMATHDX_NONE ). Trait data type: long long int.
-
enumerator CUBLASDX_DEVICE_FUNCTION_OPTION_SYMBOL_NAME#
-
enum cublasdxMemorySpace_t#
Memory space.
Values:
-
enumerator CUBLASDX_MEMORY_SPACE_RMEM#
Register (aka stack) memory space
-
enumerator CUBLASDX_MEMORY_SPACE_SMEM#
Shared memory space
-
enumerator CUBLASDX_MEMORY_SPACE_GMEM#
Global memory space
-
enumerator CUBLASDX_MEMORY_SPACE_ANY#
Unspecified, aka any of the above or any other memory space (like Tensor memory - TMEM)
-
enumerator CUBLASDX_MEMORY_SPACE_RMEM#
-
enum cublasdxDevicePipelineType_t#
Type of device pipeline.
Values:
-
enumerator CUBLASDX_DEVICE_PIPELINE_SUGGESTED#
A suggested device pipeline.
-
enumerator CUBLASDX_DEVICE_PIPELINE_SUGGESTED#
-
enum cublasdxTilePipelineType_t#
Type of tile pipeline.
Values:
-
enumerator CUBLASDX_TILE_PIPELINE_DEFAULT#
A tile pipeline.
-
enumerator CUBLASDX_TILE_PIPELINE_DEFAULT#
-
enum cublasdxBlockSizeStrategy_t#
Type of block size strategy used in device pipeline and tile pipeline creation. This enum affects register usage in pipelining GEMM kernels and should be considered for different architectures.
Values:
-
enumerator CUBLASDX_BLOCK_SIZE_STRATEGY_HEURISTIC#
A heuristic block size strategy where cuBLASDx will pick the “best” strategy using internal heuristics.
-
enumerator CUBLASDX_BLOCK_SIZE_STRATEGY_FIXED#
A fixed block size strategy where cuBLASDx will use the user-specified block size.
-
enumerator CUBLASDX_BLOCK_SIZE_STRATEGY_HEURISTIC#
-
enum cublasdxPipelineTrait_t#
Pipeline traits, used to query informations.
Values:
-
enumerator CUBLASDX_PIPELINE_TRAIT_STORAGE_BYTES#
The size of the underlying storage, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_PIPELINE_TRAIT_STORAGE_ALIGNMENT_BYTES#
The alignment of the underlying storage, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_PIPELINE_TRAIT_BUFFER_SIZE#
The size of the buffer memory, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_PIPELINE_TRAIT_BUFFER_ALIGNMENT_BYTES#
The alignment of the buffer memory, in bytes. Trait data type: long long int.
-
enumerator CUBLASDX_PIPELINE_TRAIT_OPAQUE_NAME#
A human readable but unspecified C-string representing the opaque pipeline type name. Names are stable and unique per pipeline type, and pipeline types with the same name can be used interchangeably. Opaque names are not C++ type name identifiers. Trait data type: C-string.
-
enumerator CUBLASDX_PIPELINE_TRAIT_BLOCK_DIM#
The block dimension of the pipeline. Trait data type: dim3.
-
enumerator CUBLASDX_PIPELINE_TRAIT_STORAGE_BYTES#
- commondxStatusType cublasdxGetVersion(
- int *major,
- int *minor,
- int *patch
Returns the major.minor.patch version of cuBLASDx.
- Parameters:
major – [out] The major version
minor – [out] The minor version
patch – [out] The patch version
- Returns:
COMMONDX_SUCCESS
- commondxStatusType cublasdxCreateDescriptor(
- cublasdxDescriptor *handle
Creates a cuBLASDx descriptor.
- Parameters:
handle – [out] A pointer to a handle
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetOptionStr(
- cublasdxDescriptor handle,
- commondxOption option,
- const char *value
Sets a C-string option on a cuBLASDx descriptor.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
option – [in] An option to set the descriptor to.
value – [in] A pointer to a C-string. Cannot be
NULL.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetOptionStrs(
- cublasdxDescriptor handle,
- commondxOption option,
- size_t count,
- const char **values
Sets one or more C-string options on a cuBLASDx descriptor.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
option – [in] An option to set the descriptor to.
count – [in] The number of options.
values – [in] A pointer to an array of
countC-strings.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetOperatorInt64(
- cublasdxDescriptor handle,
- cublasdxOperatorType op,
- long long int value
Set an operator on a cuBLASDx descriptor to an integer value.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
op – [in] An operator to set the descriptor to.
value – [in] The operator’s value
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetOperatorInt64s(
- cublasdxDescriptor handle,
- cublasdxOperatorType op,
- size_t count,
- const long long int *array
Set an operator on a cuBLASDx descriptor to an integer array.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
op – [in] An option to set the descriptor to.
count – [in] The size of the operator array, as indicated by the cublasdxOperatorType_t documentation
array – [in] A pointer to an array containing the operator’s value. Must point to at least
countelements.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxCreateTensor(
- cublasdxDescriptor handle,
- cublasdxTensorType tensor_type,
- cublasdxTensor *tensor
Create a tensor handle.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
tensor_type – [in] The tensor type to bind to the handle
tensor – [out] A valid tensor handle
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxCreateTensorStrided(
- cublasdxMemorySpace memory_space,
- commondxValueType value_type,
- void *ptr,
- long long int rank,
- long long int *shape,
- long long int *stride,
- cublasdxTensor *tensor
Create a tensor handle for a N-dimensional strided tensor.
The resulting tensor has the following in-memory, on-device representation
struct tensor { void* ptr; long long int shapes[n_runtime_shapes]; long long int strides[n_runtime_strides]; }
where
ptrpoints to the data in the appropriate memory space,shapes[n_runtime_shapes]is an array holding the runtime (aka not static) shapes,strides[n_runtime_strides]is an array holding the runtime (aka not static) shapes.
Runtime shapes and strides are marked by passing LIBMATHDX_RUNTIME in
shapeandstride. Static shapes and strides should be specified as-is inshapeandstrideand should not be provided again at runtime.- Parameters:
memory_space – [in] The memory space for the tensor.
value_type – [in] The datatype of the individual elements.
ptr – [in] A pointer to the data. Currently, only
NULLis supported.rank – [in] The rank of of tensor.
shape – [in] An array of size
rankindicating the tensor shape. LIBMATHDX_RUNTIME can be used to indicate a runtime shape.stride – [in] An array of size
rankindicating the tensor stride. LIBMATHDX_RUNTIME can be used to indicate a runtime stride.tensor – [out] The tensor handle
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxMakeTensorLike(
- cublasdxTensor input,
- commondxValueType value_type,
- cublasdxTensor *output
Create an opaque tensor with a identical layout (smem/gmem) or partitioner (rmem), but with a different datatype.
The resulting tensor in-memory and on-device representation is identical to
input’s representation, except that the memory pointer must point to data of the appropriate type.- Parameters:
input – [in] An opaque tensors
value_type – [in] The new datatype
output – [out] The output tensor
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxCreateDevicePipeline(
- cublasdxDescriptor handle,
- cublasdxDevicePipelineType device_pipeline_type,
- long long int pipeline_depth,
- cublasdxBlockSizeStrategy block_size_strategy,
- cublasdxTensor tensor_a,
- cublasdxTensor tensor_b,
- cublasdxPipeline *device_pipeline
Create a device pipeline handle.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
device_pipeline_type – [in] The type of the device pipeline
pipeline_depth – [in] The depth of the pipeline. If LIBMATHDX_MAX_PIPELINE_DEPTH is passed, the pipeline depth will be set to the maximal depth based on available shared memory.
block_size_strategy – [in] The block size strategy to use, fixed (you use the number of threads specified) or heuristic (cuBLASDx can use more threads if needed)
tensor_a – [in] The tensor handle for global matrix A
tensor_b – [in] The tensor handle for global matrix B
device_pipeline – [out] A valid device pipeline handle
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxCreateTilePipeline(
- cublasdxDescriptor handle,
- cublasdxTilePipelineType tile_pipeline_type,
- cublasdxPipeline device_pipeline,
- cublasdxPipeline *tile_pipeline
Create a tile pipeline handle.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
tile_pipeline_type – [in] The type of the tile pipeline
device_pipeline – [in] The device pipeline handle this tile pipeline is associated with
tile_pipeline – [out] A valid tile pipeline handle
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxFinalizeTensors(
- size_t count,
- const cublasdxTensor *array
Finalize the tensors. This is required before traits can be queried.
- Parameters:
count – [in] The number of tensors to finalized
array – [out] The array of tensors
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxFinalizePipelines(
- size_t count,
- const cublasdxPipeline *array
Finalize the pipelines. This is required before traits can be queried.
- Parameters:
count – [in] The number of pipelines to finalized
array – [out] The array of pipelines
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxFinalize(
- size_t countTensors,
- const cublasdxTensor *tensors,
- size_t countPipelines,
- const cublasdxPipeline *pipelines
Finalize both tensors and pipelines. This is required before traits can be queried. Internally calls cublasdxFinalizeTensors and cublasdxFinalizePipelines.
- Parameters:
countTensors – [in] The number of tensors to finalized
tensors – [out] The array of tensors
countPipelines – [in] The number of pipelines to finalized
pipelines – [out] The array of pipelines
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTensorTraitInt64(
- cublasdxTensor tensor,
- cublasdxTensorTrait trait,
- long long int *value
Query an integer trait value from a finalized tensor.
- Parameters:
tensor – [in] A finalized tensor handle, output of cublasdxCreateTensor
trait – [in] The trait to query
value – [out] The trait value
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTensorTraitStrSize(
- cublasdxTensor tensor,
- cublasdxTensorTrait trait,
- size_t *size
Query an C-string trait’s size from a finalized tensor.
- Parameters:
tensor – [in] A finalized tensor handle, output of cublasdxCreateTensor
trait – [in] The trait to query
size – [out] The C-string size (including the
\0)
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTensorTraitStr(
- cublasdxTensor tensor,
- cublasdxTensorTrait trait,
- size_t size,
- char *value
Query a C-string trait value from a finalized tensor.
- Parameters:
tensor – [in] A finalized tensor handle, output of cublasdxCreateTensor
trait – [in] The trait to query
size – [in] The C-string size, as returned by cublasdxGetTensorTraitStrSize
value – [out] The C-string trait value
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetPipelineTraitInt64(
- cublasdxPipeline pipeline,
- cublasdxPipelineTrait trait,
- long long int *value
Query an integer trait value from a finalized pipeline.
- Parameters:
pipeline – [in] A finalized pipeline handle, output of cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline
trait – [in] The trait to query
value – [out] The trait value
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetPipelineTraitInt64s(
- cublasdxPipeline pipeline,
- cublasdxPipelineTrait trait,
- size_t count,
- long long int *array
Returns an array trait’s value from a finalized pipeline.
- Parameters:
pipeline – [in] A finalized pipeline handle, output of cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline
trait – [in] The trait to query
count – [in] The number of values to query
array – [out] The array of trait values. Must point to exactly
countelements of typelong long int.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetPipelineTraitStrSize(
- cublasdxPipeline pipeline,
- cublasdxPipelineTrait trait,
- size_t *size
Query an C-string trait’s size from a finalized pipeline.
- Parameters:
pipeline – [in] A finalized pipeline handle, output of cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline
trait – [in] The trait to query
size – [out] The C-string size (including the
\0)
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetPipelineTraitStr(
- cublasdxPipeline pipeline,
- cublasdxPipelineTrait trait,
- size_t size,
- char *value
Query a C-string trait value from a finalized pipeline.
- Parameters:
pipeline – [in] A finalized pipeline handle, output of cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline
trait – [in] The trait to query
size – [in] The C-string size, as returned by cublasdxGetPipelineTraitStrSize
value – [out] The C-string trait value
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxCreateDeviceFunction(
- cublasdxDescriptor handle,
- cublasdxDeviceFunctionType device_function_type,
- size_t count,
- const cublasdxTensor *array,
- cublasdxDeviceFunction *device_function
Create a device function from a set of tensor.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor or LIBMATHDX_NONE if no descriptor is required.
device_function_type – [in] The device function to create.
count – [in] The number of input & output tensors to the device function.
array – [in] The array of input & output tensors.
device_function – [out] The device function.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxCreateDeviceFunctionWithPipelines(
- cublasdxDescriptor handle,
- cublasdxDeviceFunctionType device_function_type,
- size_t tensor_count,
- const cublasdxTensor *tensors,
- size_t pipeline_count,
- const cublasdxPipeline *pipelines,
- cublasdxDeviceFunction *device_function
Binds (aka create) a device function from a set of tensor.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
device_function_type – [in] The device function to create
tensor_count – [in] The number of input & output tensors to the device function
tensors – [in] The array of input & output tensors
pipeline_count – [in] The number of pipelines to the device function
pipelines – [in] The array of pipelines
device_function – [out] The device function
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxFinalizeDeviceFunctions(
- commondxCode code,
- size_t count,
- const cublasdxDeviceFunction *array
Finalize (aka codegen) a set of device function into a code handle.
After this, LTOIR can be extracted from
codeusing commondxGetCodeLTOIR- Parameters:
code – [out] A code handle, output from commondxCreateCode
count – [in] The number of device functions to codegen
array – [in] The array of device functions
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetDeviceFunctionTraitStrSize(
- cublasdxDeviceFunction device_function,
- cublasdxDeviceFunctionTrait trait,
- size_t *size
Query a device function C-string trait value size.
- Parameters:
device_function – [in] A device function handle, output from cublasdxFinalizeDeviceFunctions
trait – [in] The trait to query the device function
size – [out] The size of the trait value C-string, including the
\0
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetDeviceFunctionTraitStr(
- cublasdxDeviceFunction device_function,
- cublasdxDeviceFunctionTrait trait,
- size_t size,
- char *value
Query a device function C-string trait value.
- Parameters:
device_function – [in] A device function handle, output from cublasdxFinalizeDeviceFunctions
trait – [in] The trait to query the device function
size – [in] The size of the trait value C-string as returned by cublasdxGetDeviceFunctionTraitStrSize
value – [out] The trait value as a C-string. Must point to at least
sizebytes.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetDeviceFunctionOptionInt64(
- cublasdxDeviceFunction function,
- cublasdxDeviceFunctionOption option,
- long long int opt
Set an integer option on a device function.
- Parameters:
function – [in] A device function handle.
option – [in] The option to set on the device function.
opt – [in] The value for the option.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetDeviceFunctionOptionStr(
- cublasdxDeviceFunction function,
- cublasdxDeviceFunctionOption option,
- const char *opt
Set a string option on a device function.
- Parameters:
function – [in] A device function handle.
option – [in] The option to set on the device function.
opt – [in] The string value for the option. Must be a null terminated C-string.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetLTOIRSize(
- cublasdxDescriptor handle,
- size_t *lto_size
Returns the LTOIR size, in bytes.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
lto_size – [out] The size of the LTOIR.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetLTOIR(
- cublasdxDescriptor handle,
- size_t size,
- void *lto
Returns the LTOIR.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
size – [in] The size, in bytes, of the LTOIR, as returned by cublasdxGetLTOIRSize
lto – [out] A pointer to at least
sizebytes containing the LTOIR
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTraitStrSize(
- cublasdxDescriptor handle,
- cublasdxTraitType trait,
- size_t *size
Returns the size of a C-string trait.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
trait – [in] The trait to query the size of
size – [out] The size of the C-string value, including the
\0.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTraitStr(
- cublasdxDescriptor handle,
- cublasdxTraitType trait,
- size_t size,
- char *value
Returns a C-string trait’s value.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
trait – [in] The trait to query on the descriptor
size – [in] The size of the C-string (including the
\0)value – [out] The C-string trait value. Must point to at least
sizebytes.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTraitInt64(
- cublasdxDescriptor handle,
- cublasdxTraitType trait,
- long long int *value
Returns an integer trait’s value.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
trait – [in] A trait to query the handle for
value – [out] The trait value.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxGetTraitInt64s(
- cublasdxDescriptor handle,
- cublasdxTraitType trait,
- size_t count,
- long long int *array
Returns an array trait’s value.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
trait – [in] A trait to query handle for
count – [in] The number of elements in the trait array, as indicated in the cublasdxTraitType_t documentation.
array – [out] A pointer to at least count integers. As output, an array filled with the trait value.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
-
const char *cublasdxOperatorTypeToStr(cublasdxOperatorType op)#
Convert an operator enum to a human readable C-string.
- Parameters:
op – [in] The operator enum to convert
- Returns:
The C-string
-
const char *cublasdxTraitTypeToStr(cublasdxTraitType trait)#
Convert a trait enum to a human readable C-string.
- Parameters:
trait – [in] The trait enum to convert
- Returns:
The C-string
- commondxStatusType cublasdxFinalizeCode(
- commondxCode code,
- cublasdxDescriptor handle
Fill an instance of commondxCode with the code from the cuBLASDx descriptor.
- Parameters:
code – [out] A commondxCode code
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxDestroyDescriptor(
- cublasdxDescriptor handle
Destroy a cuBLASDx descriptor.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
- Returns:
COMMONDX_SUCCESSon success, or an error code.
-
commondxStatusType cublasdxDestroyTensor(cublasdxTensor tensor)#
Destroys a tensor handle created using cublasdxCreateTensor or cublasdxMakeTensorLike.
- Parameters:
tensor – [in] The tensor to destroy.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
-
commondxStatusType cublasdxDestroyPipeline(cublasdxPipeline pipeline)#
Destroys a pipeline handle created using cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline.
- Parameters:
pipeline – [in] The pipeline to destroy.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxDestroyDeviceFunction(
- cublasdxDeviceFunction device_function
Destroys a device function handle.
- Parameters:
device_function – [in] A cuBLASDx device function, output of cublasdxCreateDeviceFunction.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
-
const char *cublasdxApiToStr(cublasdxApi api)#
Convert an API enum to a human readable C-string.
- Parameters:
api – [in] The API enum to convert
- Returns:
The C-string
-
const char *cublasdxTypeToStr(cublasdxType type)#
Convert a type enum to a human readable C-string.
- Parameters:
type – [in] The type enum to convert
- Returns:
The C-string
-
const char *cublasdxTransposeModeToStr(cublasdxTransposeMode mode)#
Convert a transpose mode enum to a human readable C-string.
- Parameters:
mode – [in] The transpose mode enum to convert
- Returns:
The C-string
- const char *cublasdxArrangementToStr(
- cublasdxArrangement arrangement
Convert an arrangement enum to a human readable C-string.
- Parameters:
arrangement – [in] The arrangement enum to convert
- Returns:
The C-string
-
const char *cublasdxFunctionToStr(cublasdxFunction function)#
Convert a function enum to a human readable C-string.
- Parameters:
function – [in] The function enum to convert
- Returns:
The C-string
-
const char *cublasdxTensorTypeToStr(cublasdxTensorType type)#
Convert a tensor type enum to a human readable C-string.
- Parameters:
type – [in] The tensor type enum to convert
- Returns:
The C-string
-
const char *cublasdxTensorOptionToStr(cublasdxTensorOption option)#
Convert a tensor option enum to a human readable C-string.
- Parameters:
option – [in] The tensor option enum to convert
- Returns:
The C-string
-
const char *cublasdxTensorTraitToStr(cublasdxTensorTrait trait)#
Convert a tensor trait enum to a human readable C-string.
- Parameters:
trait – [in] The tensor trait enum to convert
- Returns:
The C-string
- const char *cublasdxDeviceFunctionTraitToStr(
- cublasdxDeviceFunctionTrait trait
Convert a device function trait enum to a human readable C-string.
- Parameters:
trait – [in] The device function trait enum to convert
- Returns:
The C-string
- const char *cublasdxDeviceFunctionOptionToStr(
- cublasdxDeviceFunctionOption option
Convert a device function option enum to a human readable C-string.
- Parameters:
option – [in] The device function option enum to convert
- Returns:
The C-string
- const char *cublasdxDeviceFunctionTypeToStr(
- cublasdxDeviceFunctionType type
Convert a device function type enum to a human readable C-string.
- Parameters:
type – [in] The device function type enum to convert
- Returns:
The C-string
- const char *cublasdxMemorySpaceToStr(
- cublasdxMemorySpace memory_space
Convert a memory space enum to a human readable C-string.
- Parameters:
memory_space – [in] The memory space enum to convert
- Returns:
The C-string
- const char *cublasdxBlockSizeStrategyToStr(
- cublasdxBlockSizeStrategy strategy
Convert a block size strategy enum to a human readable C-string.
- Parameters:
strategy – [in] The block size strategy enum to convert
- Returns:
The C-string
- const char *cublasdxDevicePipelineTypeToStr(
- cublasdxDevicePipelineType type
Convert a device pipeline type enum to a human readable C-string.
- Parameters:
type – [in] The device pipeline type enum to convert
- Returns:
The C-string
- const char *cublasdxTilePipelineTypeToStr(
- cublasdxTilePipelineType type
Convert a tile pipeline type enum to a human readable C-string.
- Parameters:
type – [in] The tile pipeline type enum to convert
- Returns:
The C-string
-
const char *cublasdxPipelineTraitToStr(cublasdxPipelineTrait trait)#
Convert a pipeline trait enum to a human readable C-string.
- Parameters:
trait – [in] The pipeline trait enum to convert
- Returns:
The C-string
- commondxStatusType cublasdxGetTraitCommondxDataTypes(
- cublasdxDescriptor handle,
- cublasdxTraitType trait,
- size_t count,
- commondxValueType *array
Returns an array of CommonDxValueType values representing value type traits.
- Parameters:
handle – [in] A cuBLASDx descriptor, output of cublasdxCreateDescriptor
trait – [in] A trait to query handle for
count – [in] The number of elements in the trait array, as indicated in the cublasdxTraitType_t documentation.
array – [out] A pointer to at least count commondxValueType. As output, an array filled with the trait value.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetTensorOptionInt64(
- cublasdxTensor tensor,
- cublasdxTensorOption option,
- long long int value
Set an option on a tensor. This must be called before the tensor is finalized.
- Parameters:
tensor – [in] A cuBLASDx tensor, output of cublasdxCreateTensor.
option – [in] The option to set on the tensor.
value – [in] A value for the option.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetTensorOptionStr(
- cublasdxTensor tensor,
- commondxOption option,
- const char *value
Set an option on a tensor. This must be called before the tensor is finalized.
- Parameters:
tensor – [in] A cuBLASDx tensor, output of cublasdxCreateTensor.
option – [in] The commondx option to set on the tensor.
value – [in] A C-string value to set on the tensor.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetTensorOptionStrs(
- cublasdxTensor tensor,
- commondxOption option,
- size_t count,
- const char **values
Set one or more options on a tensor. This must be called before the tensor is finalized.
- Parameters:
tensor – [in] A cuBLASDx tensor, output of cublasdxCreateTensor.
option – [in] The commondx option to set on the tensor.
count – [in] The number of options to set.
values – [in] A pointer to an array of
countC-strings.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetPipelineOptionStr(
- cublasdxPipeline pipeline,
- commondxOption option,
- const char *value
Set an option on a pipeline. This must be called before the pipeline is finalized.
- Parameters:
pipeline – [in] A cuBLASDx pipeline, output of cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline.
option – [in] The commondx option to set on the pipeline.
value – [in] A C-string value to set on the pipeline.
- Returns:
COMMONDX_SUCCESSon success, or an error code.
- commondxStatusType cublasdxSetPipelineOptionStrs(
- cublasdxPipeline pipeline,
- commondxOption option,
- size_t count,
- const char **values
Set one or more options on a pipeline. This must be called before the pipeline is finalized.
- Parameters:
pipeline – [in] A cuBLASDx pipeline, output of cublasdxCreateDevicePipeline or cublasdxCreateTilePipeline.
option – [in] The commondx option to set on the pipeline.
count – [in] The number of options to set.
values – [in] A pointer to an array of
countC-strings.
- Returns:
COMMONDX_SUCCESSon success, or an error code.