张量#

cuBLASDx 张量#

cuBLASDx 公开了 cublasdx::tensor 类，它是来自 CuTe 库的 cute::Tensor 类的别名（参见 CUTLASS）。

张量创建#

// #1: Wrapper of cute::make_tensor
template<class Iterator, class... Args>
__device__ __host__
constexpr cublasdx::tensor make_tensor(const Iterator& iter, Args const&... args);

// #2: With pointer layout returned by the get_<smem/gmem>_layout_<a/b/c>, suggest_layout_smem_<a/b/c> method from the BLAS description.
template<class T, class PointerLayout>
__device__ __host__
constexpr cublasdx::tensor make_tensor(T* ptr, const PointerLayout& pl);

cublasdx::make_tensor 是用于创建 cublasdx::tensor 对象的辅助函数。

这里有两种变体。第一种只是 cute::make_tensor(…) 的一个包装器，它通常需要手动标记原始指针及其内存空间。另一种与获取内存布局和建议的共享内存布局方法一起使用。它使用返回的指针布局创建一个全局或共享内存张量。与第一种变体相比，它将从指针布局中获取内存空间信息，并相应地标记原始指针。

示例

using BLAS = decltype(Size<128, 128, 128>() + Type<type::real>() + Precision<float, float, double>() + Block() + ...);

// tensor with global memory data
auto a_global_tensor = cublasdx::make_tensor(a, BLAS::get_layout_gmem_a());

// tensor with shared memory data
auto a_shared_tensor = cublasdx::make_tensor(smem_a, BLAS::get_layout_smem_a());
auto a_shared_tensor = cublasdx::make_tensor(smem_a, BLAS::suggest_layout_smem_a());

分区器和寄存器片段张量#

cuBLASDx 0.3.0 引入了 2 个新的内存对象：寄存器片段张量和分区器。

由于寄存器内存限制于一个线程，它不用于存储整个张量，仅用于存储它们的小部分。将全局/共享内存张量划分为参与 BLAS 的线程称为分区。在 cuBLASDx 中，BLAS 操作是分区模式的来源。

在每次 BLAS 执行中，每个线程都被分配了结果矩阵 C 的一些元素，它会计算这些元素。这种选择取决于 3 个因素

线程数，
选择的 MMA 指令，以及
为 MMA 指令选择的平铺。

cuBLASDx 的用户只能直接控制这 3 个选项中的 1 个，因此分区信息必须是不透明的，并且只能与定义它的 BLAS 对象一起使用。cuBLASDx 通过 partitioner 的方式公开此不透明信息。

分区器#

可以从 BLAS 类型以 2 种方式获取分区器对象

// #1a As pre-requested default partitioner
auto partitioner = BLAS::get_partitioner();

// #1b As pre-requested suggested partitioner
auto partitioner = BLAS::suggest_partitioner();

// #2 As result of register API without accumulator call
auto [c_register_fragment, partitioner] = BLAS().execute(a_shared_tensor, b_shared_tensor);

注意

选项 #1a 和 #1b 要求用户选择与使用的布局一致的版本。Suggested 分区器仅适用于建议的布局，而 default 分区器适用于任何其他布局。混合使用它们可能会导致性能下降。

分区器是一个不透明的对象，公开以下 API

__device__ bool is_predicated();
__device__ bool is_thread_active();
__device__ constexpr auto make_accumulator_fragment();

template<class CTensor>
__forceinline__ __device__
auto partition_like_C(CTensor && ctensor) const;

template<class ... Coords>
__forceinline__ __device__
auto map_fragment_index(Coords&& ... coords) const;

template<class ... Coords>
__forceinline__ __device__
bool is_index_in_bounds(Coords&& ... coords) const;

有 3 种方法可用于获取有关分区模式的基本信息，或检索寄存器片段。

// Partitioning properties
__device__ bool is_predicated();
__device__ bool is_thread_active();

// Accumulator creation, creates a register cublasdx::tensor
__device__ constexpr auto make_accumulator_fragment();

如分区器和寄存器片段张量中提到的，线程之间元素的划分是通过在 BLAS 定义的问题大小上平铺 MMA 指令来实现的。每个指令负责计算特定形状，并且问题形状可能无法被原始指令形状整除。在这种情况下，“额外”的指令元素将填充 0，而不是从内存中读取，并在存储结果时跳过。这称为元素谓词化。is_predicated() 成员函数可用于检索该信息。

由于 cuBLASDx 支持在 CUDA 线程块大小与 BlockDim 运算符不匹配的内核中执行，因此并非所有线程都将参与 GEMM 运算。这意味着某些线程可能没有任何元素分配给它们。is_thread_active_result() 成员函数可用于准确检查调用线程是否属于这种情况。

返回的布尔值是在编译时还是运行时已知将取决于所使用的特定 BLAS 配置。

最后一个无参数方法 - make_accumulator_fragment() 可用于检索与从中检索分区器的 BLAS 执行相对应的未初始化的寄存器片段张量。

请参阅下面的代码，了解所有 3 种方法的使用示例。

template<class CTensor>
__forceinline__ __device__
auto partition_like_C(CTensor && ctensor) const;

此方法将返回其参数的子张量的非拥有视图，该子张量分配给调用线程，对应于其本地寄存器片段。

所有这些方法都可用于手动复制非谓词化 GEMM，如下所示

分区器示例 #1#

auto c_global_tensor = ...;

// Get partitioner
auto partitioner = BLAS::get_partitioner();
// Create local register fragment
auto c_register_fragment = partitioner.make_accumulator_fragment();
// Get view of this thread's global memory subtensor
auto c_global_partition = partitioner.partition_like_C(c_global_tensor);

// Ensure that all elements are in-bounds and
// no predication is necessary
static_assert(not partitioner.is_predicated());

// If this thread takes part in GEMM
if(partitioner.is_thread_active()) {
   // For each element of register fragment
   for(int i = 0; i < cublasdx::size(c_register_fragment); ++i) {
      // Copy respective global element into it
      c_register_fragment(i) = c_global_partition(i);
   }
}

// C += A * B
BLAS().execute(a_shared_tensor, b_shared_tensor, c_register_fragment);

如果 GEMM 是谓词化的，则需要更多信息

template<class ... Coords>
__forceinline__ __device__
auto map_fragment_index(Coords&& ... coords) const;

template<class ... Coords>
__forceinline__ __device__
bool is_index_in_bounds(Coords&& ... coords) const;

这两个函数扩展了 is_predicated() 的功能，允许将本地寄存器片段索引映射到其源（全局或共享）张量索引，并检查此索引是否在界限内。

所有这些方法都可用于执行从全局张量到寄存器张量的边界检查加载元素

分区器示例 #2#

auto c_global_tensor = ...;

// Get partitioner
auto partitioner = BLAS::get_partitioner();
// Create local register fragment
auto c_register_fragment = partitioner.make_accumulator_fragment();

// If this thread takes part in GEMM
if(partitioner.is_thread_active()) {
   // For each element of register fragment
   for(int i = 0; i < cublasdx::size(c_register_fragment); ++i) {
      auto global_index = partitioner.map_fragment_index(i);
      if((not partitioner.is_predicated()) or partitioner.is_index_in_bounds(i)) {
         // Copy respective global element into it
         c_register_fragment(i) = load_op(c_global_tensor(global_index));
      }
   }
}

// C += A * B
BLAS().execute(a_shared_tensor, b_shared_tensor, c_register_fragment);

每次都必须手动执行此类加载将非常繁琐，因此 cuBLASDx 提供了优化的和自动向量化的全局/共享 ⟷ 寄存器片段复制函数。有关更多信息，请参阅复制寄存器张量

导入的张量实用程序#

由于 cuBLASDx 集成了 CuTe 库，因此它公开了其一些张量和布局功能，有时会添加兼容性层以支持其自己的类型。目前支持以下函数

// Aliased
using cute::clear;
using cute::transform;
using cute::make_fragment_like;

// With compatibility layers
cublasdx::size;   // (cute::size)
cublasdx::cosize; // (cute::cosize)
cublasdx::axpby;  // (cute::axpby)

cublasdx::clear 接受寄存器片段作为输入，并将其所有值设置为 0。

auto partitioner = BLAS::get_partitioner();
auto c_register_fragment = partitioner.make_accumulator_fragment();

cublasdx::clear(c_register_fragment);

cublasdx::transform 将元素级仿函数应用于传递的寄存器张量的所有元素。

auto c_global_tensor = ...

auto partitioner = BLAS::get_partitioner();
auto c_register_fragment = partitioner.make_accumulator_fragment();

cublasdx::copy_fragment<alignment_of<BLAS>::c>(c_global_tensor, c_register_fragment, partitioner);

// in-place
cublasdx::transform(c_register_fragment, [](auto v) { return 2 * v; });

// out-of-place
auto d_register_fragment = partitioner.make_accumulator_fragment();
cublasdx::transform(c_register_fragment, d_register_fragment, [](auto v) { return 2 * v; });

cublasdx::make_fragment_like 可用于创建具有与其参数张量相同的布局和类型的寄存器片段。

auto c_global_tensor = ...

auto partitioner = BLAS::get_partitioner();
auto c_global_partition = partitioner.partition_like_C(c_global_tensor);

// Same type
auto c_register_fragment = cublasdx::make_fragment_like(c_global_partition);

// Other type
using new_type = double;
auto c_register_fragment = cublasdx::make_fragment_like<new_type>(c_global_partition);

警告

以下函数应仅以其 cuBLASDx 形式使用，因为它提供了兼容性层和自动转换。

cublasdx::axpby 执行 \(\mathbf{D}_{m\times n} = {\alpha} \times \mathbf{C}_{m\times n} + {\beta} \times \mathbf{D}_{m\times n}\)

auto c_global_tensor = ...

auto a_shared_tensor = ...
auto b_shared_tensor = ...

auto partitioner = BLAS::get_partitioner();
auto c_register_fragment = partitioner.make_accumulator_fragment();
auto d_register_fragment = partitioner.make_accumulator_fragment();
cublasdx::copy_fragment<alignment_of<BLAS>::c>(c_global_tensor, c_register_fragment, partitioner);

// These 2 functions combined perform the classic GEMM: C = alpha * A * B + beta * C;
BLAS().execute(a_shared_tensor, b_shared_tensor, c_register_fragment);
cublasdx::axpby(alpha, c_register_fragment, beta, d_register_fragment);

cublasdx::copy_fragment<alignment_of<BLAS>::c>(d_register_fragment, c_global_tensor, partitioner);

cublasdx::size 返回张量中有效元素的数量。这只是所有 shape 维度的乘积。

auto a_shared_layout = BLAS::get_layout_smem_a();

// Same as size_of<BLAS>::m * size_of<BLAS>::k
constexpr auto layout_size = cublasdx::size(a_shared_layout);

cublasdx::cosize 返回从张量的最后一个元素到其第一个元素的距离。它描述了参数布局跨越了多少个元素。对于紧凑或没有“孔”（例如，作为额外的前导维度元素的結果）的布局，这与 cublasdx::size 相同

// Default leading dimension
auto BLAS = decltype(... + Size<M, N, K>() + LeadingDimension<M, K, M>());
auto a_shared_layout = BLAS::get_layout_smem_a();

constexpr auto layout_size = cublasdx::size(a_shared_layout);
constexpr auto layout_cosize = cublasdx::cosize(a_shared_layout);

static_assert(layout_size == layout_cosize);

// Extra elements leading dimension
constexpr auto lda = M + 1; // Add padding to avoid shared memory bank conflicts
auto BLAS = decltype(... + Size<M, N, K>() + LeadingDimension<lda, K, M>());
auto a_shared_layout = BLAS::get_layout_smem_a();

constexpr auto layout_size = cublasdx::size(a_shared_layout);
constexpr auto layout_cosize = cublasdx::cosize(a_shared_layout);

static_assert(layout_size != layout_cosize);
static_assert(layout_size < layout_cosize);

复制张量#

协同全局 ⟷ 共享复制#

template<uint32_t NumThreads,       // Number of threads performing copy operation
         uint32_t AlignmentInBytes, // Pointer alignment of src and dst tensor (minimum of them if they are different)
         class SrcEngine,
         class SrcLayout,
         class DstEngine,
         class DstLayout>
__forceinline__ __device__
void copy(const unsigned int                            tid, // Thread index in CUDA block
          const cublasdx::tensor<SrcEngine, SrcLayout>& src,
          cublasdx::tensor<DstEngine, DstLayout>&       dst)

// Assumes pointers in both dst and src tensors are not extra aligned
template<uint32_t NumThreads, // Number of threads performing copy operation
         class SrcEngine,
         class SrcLayout,
         class DstEngine,
         class DstLayout>
__forceinline__ __device__
void copy(const unsigned int                            tid, // Thread index in CUDA block
          const cublasdx::tensor<SrcEngine, SrcLayout>& src,
          cublasdx::tensor<DstEngine, DstLayout>&       dst)

template<class BLAS,                // BLAS description which provides the number of threads
         uint32_t AlignmentInBytes, // Pointer alignment of src and dst tensor (minimum of them if they are different)
         class SrcEngine,
         class SrcLayout,
         class DstEngine,
         class DstLayout>
__forceinline__ __device__
void copy(const cublasdx::tensor<SrcEngine, SrcLayout>& src,
          cublasdx::tensor<DstEngine, DstLayout>&       dst)

cublasdx::copy 是用于在共享内存或全局内存中的张量之间复制数据的辅助函数。

复制是协同完成的。所有线程（由 NumThreads 或 BLAS::block_dim 指示）都将参与复制。该函数考虑了给定的对齐方式，并尝试在可能的情况下向量化加载和存储指令。

要求

张量中的数据必须在共享内存或全局内存中。不支持从寄存器复制或复制到寄存器。
src 和 dst 张量都必须表示相同底层元素类型的张量（cublasdx::tensor<Engine, Layout>::value_type, Engine::value_type）。
src 和 dst 张量都必须具有相同的大小，即元素数量。
AlignmentInBytes 必须是张量的底层元素类型对齐方式的倍数。
AlignmentInBytes 必须等于 1、2、4、8 或 16，或等于张量的底层元素类型的对齐方式。
src 和 dst 张量中的底层指针必须与 AlignmentInBytes 字节对齐。

// Synchronization step required after cublasdx::copy and before the use of dst tensor
__forceinline__ __device__ void copy_wait();

cublasdx::copy_wait 创建同步点。它必须在 cublasdx::copy 操作之后、在任何后续读取或写入 dst 张量之前，以及在任何后续写入 src 张量之前调用。否则，复制操作的结果是未定义的。重要的是要注意，它始终不是 __syncthreads() 的一对一等效项，因为它还处理异步数据复制（参见 cp.async 指令系列）

示例

从全局内存复制 A 矩阵到共享内存并返回的示例。

using BLAS = decltype(Size<128, 128, 128>() + Type<type::real>() + Precision<float, float, double>() + Block() + ...);
extern __shared__ __align__(16) char smem[];

// Slice shared memory
auto [smem_a, smem_b, smem_c] = cublasdx::slice_shared_memory<BLAS>(smem);

auto gmem_tensor_a = cublasdx::make_tensor(a_gmem_pointer, BLAS::get_layout_gmem_a());
auto smem_tensor_a = cublasdx::make_tensor(smem_a, BLAS::suggest_layout_smem_a());

// Copy from global to shared
using alignment = cublasdx::alignment_of<BLAS>;
cublasdx::copy<BLAS, alignment::a>(gmem_tensor_a, smem_tensor_a);
cublasdx::copy_wait();

// Copy from shared to global
cublasdx::copy<BLAS, alignment::a>(smem_tensor_a, gmem_tensor_a);
cublasdx::copy_wait();

复制寄存器张量#

可以使用 partitioner API 或使用 cublasdx::copy_fragment 将寄存器片段手动复制到全局或共享内存张量。

分区器 API 允许 2 种不同的复制方法

创建调用线程元素的子张量视图，并使用寄存器片段索引空间进行复制（参见分区器示例 #1）。
将本地寄存器片段索引映射到全局索引空间，并直接寻址全局/共享张量，从而实现位置感知处理（参见分区器示例 #2）。

第二种方法是基于双向复制方法 cublasdx::copy_fragment，它按线程工作

#1 Store fragment: partition and copy from register fragment to global / shared memory tensor
template<unsigned AlignmentInBytes,    // Alignment of source tensor pointer
         class TRC, class CFragLayout, // Register Memory Fragment Tensor
         class TC, class CLayout,      // Global or Shared Memory tensor
         class Partitioner>
__forceinline__ __device__
copy_fragment(tensor<TRC, CFragLayout> const& tS, // Entire non-partitioned global / shared tensor
              tensor<TC, CLayout>           & tD, // Calling thread's register fragment tensor
              Partitioner              const& p);

#2 Load fragment: partition and copy from global / shared memory tensor to register fragment
template<unsigned AlignmentInBytes,    // Alignment of source tensor pointer
         class TRC, class CFragLayout, // Register Memory Fragment Tensor
         class TC, class CLayout,      // Global or Shared Memory tensor
         class Partitioner>
__forceinline__ __device__
copy_fragment(tensor<TC, CLayout>      const& tS,
              tensor<TRC, CFragLayout>      & tD,
              Partitioner              const& p);

根据定义的语义，这些可以被认为是 gather 和 scatter 函数

从全局/共享张量角度来看：load fragment 是 scatter，store fragment 是 gather。
从线程角度来看：load fragment 是 gather，store fragment 是 scatter。

重要的是要注意，这些函数中的每一个都接受

一个模板参数 unsigned integer，描述共享/全局内存张量的对齐方式，
此线程的寄存器片段张量，
一个完整的非分区全局/共享内存张量，以及
一个分区器。

它们都可以安全地用于线程和值谓词化片段 GEMM，并将自动向量化到传递的对齐方式所描述的程度。

使用示例，执行 register API GEMM with accumulator

auto partitioner = GEMM::get_partitioner();
auto c_register_fragment = partitioner.make_accumulator_fragment();

// Load fragment
cublasdx::copy_fragment<alignment::c>(c_global_tensor, c_register_fragment, partitioner);

GEMM().execute(a_shared_tensor, b_shared_tensor, c_register_fragment);

// Store fragment
cublasdx::copy_fragment<alignment::c>(c_register_fragment, c_global_tensor, partitioner);