共享内存管理#

cuBLASDx 的 GEMM 要求输入矩阵位于共享内存中，并可以选择将 C 矩阵在线程之间分区。共享内存的使用施加了某些规则（对齐）和限制（有限的共享内存空间）。cuBLASDx 提供了共享内存管理工具，以方便所有用户在共享内存上操作。

警告

从 cuBLASDx 0.3.0 版本开始，::shared_memory_size 特性和 ::shared_memory_size() 方法不再存在，并已被本章中描述的新 API 所取代。

共享存储大小实用工具#

CUDA 要求用户在内核启动期间预先指定已使用的动态共享内存量。对于执行 GEMM，此大小将取决于

  1. 问题大小 (Size<M, N, K>)。

  2. 选择的输入类型（由于 cuBLASDx 0.3.0 输入精度可以与计算精度不同，请参阅 精度运算符）。

  3. 选择的矩阵对齐方式 (Alignment<A, B, C>)。

  4. 将用于执行 GEMM 的 API（寄存器或共享内存）。

由于此信息既不包含在 BLAS 类型中（因为输入精度与计算精度解耦），也不包含在张量中（它们缺少对齐信息），因此创建了辅助实用工具来帮助简化此过程

// Shared memory API
template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type,
                     class CValueType = typename BLAS::c_value_type,
                     class ALayout, class BLayout, class CLayout>
constexpr unsigned get_shared_storage_size(ALayout const& a_layout, BLayout const& b_layout, CLayout const& c_layout);

template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type,
                     class CValueType = typename BLAS::c_value_type>
__host__ __device__ __forceinline__ constexpr unsigned
get_shared_storage_size(unsigned lda, unsigned ldb, unsigned ldc);

// Register API
template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type,
                     class ALayout, class BLayout>
constexpr unsigned get_shared_storage_size_ab(ALayout const& a_layout, BLayout const& b_layout);

template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type>
constexpr unsigned get_shared_storage_size_ab(unsigned lda, unsigned ldb);

结果值是以字节为单位的共享内存大小，用于分配输入和输出矩阵并执行计算。

该值由 值类型、矩阵大小 和 对齐方式 的大小决定。

请注意，BLAS::get_shared_memory_size 接受任意 CuTe 布局。上述函数原型中的类 ALayout、BLayout 和 CLayout 可以是 cute::Layout 或 cute::ComposedLayout。

这些函数可以按如下方式使用

// Shared API - Regular execution
auto shared_size = cublasdx::get_shared_storage_size<BLAS>();

// Shared API - Decoupled input precision execution
auto shared_size = cublasdx::get_shared_storage_size<BLAS, InputTypeA, InputTypeB, InputTypeC>();


// Shared API - Regular execution
auto shared_size = cublasdx::get_shared_storage_size_ab<BLAS>();

// Shared API - Decoupled input precision execution
auto shared_size = cublasdx::get_shared_storage_size_ab<BLAS, InputTypeA, InputTypeB, InputTypeC>();

对于特殊情况，cuBLASDx 提供了一个共享内存大小计算器

shared_storage_calc make_shared_storage_calc();

它公开了以下 API

template<class Layout>
shared_storage_calculator& add(unsigned alignment, unsigned elem_size, const Layout& layout);

__host__ __device__ __forceinline__ constexpr
shared_storage_calculator& add(unsigned alignment, unsigned matrix_size_bytes);

__host__ __device__ __forceinline__ constexpr
shared_storage_calculator& add(unsigned alignment, unsigned elem_size, unsigned num_elements);

它可以用于计算流水线式寄存器 API 执行的共享内存需求

// 2 Stage pipelined register memory execution
auto shared_memory_size =
cublasdx::make_shared_storage_calc()
.add(cublasdx::alignment_of_v_a<BLAS>, sizeof(AInputType), BLAS::suggest_layout_smem_a())
.add(cublasdx::alignment_of_v_b<BLAS>, sizeof(BInputType), BLAS::suggest_layout_smem_b())
.add(cublasdx::alignment_of_v_a<BLAS>, sizeof(AInputType), BLAS::suggest_layout_smem_a())
.add(cublasdx::alignment_of_v_b<BLAS>, sizeof(BInputType), BLAS::suggest_layout_smem_b())
.get();

共享内存切片#

警告

从 cuBLASDx 0.3.0 版本开始，::slice_shared_memory() 方法不再存在，并已移至下面描述的 API。

如果 is_complete_blas_execution 特性 对其为 true，则共享内存切片自由函数与 BLAS 一起使用。

// #1 Slice shared memory with default leading dimensions and default matrices layouts
template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type,
                     class CValueType = typename BLAS::c_value_type>
cute::tuple<AValueType*, BValueType*, CValueType*>
cublasdx::slice_shared_memory(void* smem_ptr)

// #2: Slice shared memory with dynamic leading dimensions
template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type,
                     class CValueType = typename BLAS::c_value_type>
cute::tuple<AValueType*, BValueType*, CValueType*>
cublasdx::slice_shared_memory(void* smem_ptr,
                              unsigned int lda,
                              unsigned int ldb,
                              unsigned int ldc)

// #3: Slice shared memory with custom matrices layouts
template<class BLAS, class AValueType = typename BLAS::a_value_type,
                     class BValueType = typename BLAS::b_value_type,
                     class CValueType = typename BLAS::c_value_type,
                     class ALayout, class BLayout, class CLayout>
cute::tuple<AValueType*, BValueType*, CValueType*>
cublasdx::slice_shared_memory(void* smem_ptr,
                              ALayout a_layout,
                              BLayout b_layout,
                              CLayout c_layout)

方法 cublasdx::slice_shared_memory(...) 将共享内存切片成块，每个矩阵一块。

返回值是指向 A、B 和 C 矩阵切片第一个元素的指针。它们遵循 BLAS 描述中的 对齐方式，同时，没有过度对齐，即，两个切片之间的字节数小于对齐方式。

请注意，BLAS::slice_shared_memory 接受任意 CuTe 布局。上述函数原型中的类 ALayout、BLayout 和 CLayout 可以是 cute::Layout 或 cute::ComposedLayout。

示例

using BLAS = decltype(...);

extern __shared__ __align__(16) char smem[];

// use structured binding
auto [smem_a, smem_b, smem_c] = BLAS::slice_shared_memory();

// or
auto smem_slices = BLAS::slice_shared_memory();
auto smem_a = cute::get<0>(sliced_smem);
auto smem_b = cute::get<1>(sliced_smem);
auto smem_c = cute::get<2>(sliced_smem);

警告

当使用 NVRTC 时，会使用 libcu++ (libcudacxx)，并且 std:: 被替换为 cuda::std::。

using BLAS = decltype(...);

extern __shared__ __align__(16) char smem[];

// Structured bindings support for cuda::std::tuple was added in 2.1.0 version of libcu++
#if _LIBCUDACXX_CUDA_API_VERSION >= 2001000
auto [smem_a, smem_b, smem_c] = BLAS::slice_shared_memory();
#else
auto smem_slices = BLAS::slice_shared_memory();
auto smem_a = std::get<0>(sliced_smem);
auto smem_b = std::get<1>(sliced_smem);
auto smem_c = std::get<2>(sliced_smem);
#endif

cuBLASDx 还提供了一个高级通用切片 API，允许任意数量的输入矩阵

// memory_descriptor is just a cute::tuple<unsigned, unsigned> containing:
// 1. Cosize of tensor
// 2. Alignment of tensor

template<class ... PointerTypes, class ... Tuples>
__host__ __device__ __forceinline__ auto
slice_shared_memory_generic(void* smem, Tuples const& ... memory_descriptors);

它可以按以下方式使用，以便为 2 阶段流水线式 GEMM 执行正确地切片共享内存

auto [smem_a, smem_b, smem_a_n, smem_b_n] =
 cublasdx::slice_shared_memory_generic<AInputType, BInputType, AInputType, BInputType>(
     smem,
     cute::make_tuple(cublasdx::cosize(BLAS::suggest_layout_smem_a()), cublasdx::alignment_of_v_a<BLAS>),
     cute::make_tuple(cublasdx::cosize(BLAS::suggest_layout_smem_b()), cublasdx::alignment_of_v_b<BLAS>),
     cute::make_tuple(cublasdx::cosize(BLAS::suggest_layout_smem_a()), cublasdx::alignment_of_v_a<BLAS>),
     cute::make_tuple(cublasdx::cosize(BLAS::suggest_layout_smem_b()), cublasdx::alignment_of_v_b<BLAS>)
 );