Tensors

cuBLASDx Tensor

cuBLASDx exposes cublasdx::tensor class which is an alias of a cute::Tensor class from CuTe library (see CUTLASS).

Tensor Creation

// #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_layout_<smem/gmem>_<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 is a helper function for creating cublasdx::tensor objects.

There are two variants. The first one is simply a wrapper of cute::make_tensor(…), which usually requires manually tagging the raw pointer with its memory space. The other one works together with the Get Memory Layout and the Suggested Shared Memory Layout methods. It creates a global or shared memory tensor with the returned pointer layout. In contrast to the first variant, it will pick up the memory space information from the pointer layouts and tag the raw pointer correspondingly.

Example

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());

// or, with the suggested layout
auto a_suggested_shared_tensor = cublasdx::make_tensor(smem_a, BLAS::suggest_layout_smem_a());

Accumulator And Register Fragment Tensors

cuBLASDx 0.3.0 introduces 2 new memory objects: register fragment tensors and accumulators.

Since register memory is limited to a thread, it’s not used to store entire tensors, only small parts of them. A division of a global / shared memory tensor into threads taking part in a BLAS is called a partitioning. In cuBLASDx the BLAS operation is a source of partitioning pattern.

In every BLAS execution each thread is assigned some elements of result matrix C which it calculates. This selection is dependent upon 3 things:

1. number of threads,

2. chosen MMA instruction, and

3. tiling chosen for the MMA instruction.

User of cuBLASDx has direct control over only 1 of these 3 options, hence the partitioning information must be opaque and used only with BLAS object which defined it. cuBLASDx exposes this opaque information by means of an accumulator.

Accumulator

Accumulator objects can be obtained from BLAS type in 2 ways:

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

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

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

Note

Options #1a and #1b require the user to choose version coherent with used layouts. Suggested accumulator only with suggested layouts and default accumulator with any other layouts.

Accumulator is an opaque object exposing the following API:

// Both return a value contextually convertible to bool; may be a compile-time constant
__device__ constexpr auto is_predicated() const;
__device__ auto is_thread_active() const;

// Returns a new zero-initialized register fragment with the same shape as the results
__device__ auto make_empty_fragment() const;

// Returns a copy of this thread's results as a new register fragment
__device__ auto get_results() const;

// Copies results into a user-provided register tensor
template<class Engine, class Layout>
__device__ void get_results(cublasdx::tensor<Engine, Layout>& out) const;
template<class Engine, class Layout>
__device__ void get_results(cublasdx::tensor<Engine, Layout>&& out) const;

// Resets the accumulator so that the next execute() starts from zero
__device__ void clear();

// Partitioning interfaces

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;

template<class Coord>
__forceinline__ __device__
bool is_coord_in_bounds(Coord const& coord) const;

// Partitioning and data copying interfaces

template<class Fragment, class CTensor>
__forceinline__ __device__
void partition_and_copy(Fragment const& fragment, CTensor && ctensor) const;

template<class CTensor, class Fragment>
__forceinline__ __device__
void partition_and_copy(CTensor const& ctensor, Fragment && fragment) const;

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

template<class CTensor>
__forceinline__ __device__
void partition_and_store(CTensor && ctensor) const;

// Epilogue interfaces

template<class CTensor, class ProcessOp>
__forceinline__ __device__
void process_and_store(CTensor && ctensor, ProcessOp && process_op) const;

template<class CTensor, class Lambda>
__forceinline__ __device__
void reduce_and_store(CTensor && ctensor, Lambda && lambda) const;

template<class Alpha, class Beta, class CTensor,
         class LoadOp = cublasdx::identity, class StoreOp = cublasdx::identity>
__forceinline__ __device__
void axpby(Alpha const& alpha, Beta const& beta, CTensor && ctensor,
           LoadOp const& load_op = {}, StoreOp const& store_op = {}) const;

accumulator.axpby(alpha, beta, ctensor, load_op, store_op) computes ctensor = store_op(alpha * results + beta * load_op(ctensor)). When beta is zero, ctensor is not read (it may be uninitialized; its contents, including NaNs, do not affect the result), so load_op has no effect; store_op is always applied.

When partition_and_store(...), process_and_store(...), reduce_and_store(...), or axpby(...) write to a shared memory tensor, each thread writes only its own elements and no synchronization is performed. __syncthreads() is required before other threads read that tensor (including a block-wide cublasdx::copy of the tile).

Warning

An epilogue operation that reads the accumulator may be performed only once on a ready accumulator. These are the reading operations:

• get_results()

• process_and_store()

• reduce_and_store()

• partition_and_store()

• axpby()

When an accumulator becomes ready depends on which result-storage mode owns it:

• With cublasdx::internal_accumulator the pipeline owns it, and it arrives ready in the epilogue functor passed to execute(). The register API is the same in spirit: BLAS().execute(...) returns a ready accumulator directly.

• With cublasdx::reusable_accumulator, where the accumulator is user-owned, finish_accumulation() must be called first - or it must be read through tile_pipeline.epilogue(), which calls finish_accumulation() internally.

To read it again, run another accumulation cycle.

The restriction exists because on architectures with a tensor memory accumulator (UTCMMA, SM 100 and newer) every one of these calls advances the accumulator’s TMEM pipeline stage and signals the compute-side barrier. A second call therefore reads a different stage and arrives at a barrier the producer has not filled, which yields wrong results or hangs - it is not merely redundant work. Register-backed accumulators have no such state and tolerate repeated reads, so code that breaks this rule can still pass on pre-SM-100 hardware.

To apply several transforms to one result, read the accumulator once into a register fragment and work on that fragment:

// Correct: one read, then operate on the fragment
auto results = accumulator.get_results();
// ... any number of operations on `results` ...

// Incorrect: two reads of the same ready accumulator
auto results = accumulator.get_results();
accumulator.reduce_and_store(c_global_tensor, lambda);

The following methods do not read the accumulator and are unrestricted: make_empty_fragment(), partition_and_copy(), make_partition_and_copy(), partition_like_C(), map_fragment_index(), is_predicated(), is_thread_active(), is_index_in_bounds() and is_coord_in_bounds().

There are 4 methods that can be used to get basic information about the partitioning pattern, or retrieve a register fragment.

// Partitioning properties
// (values contextually convertible to bool; may be compile-time constants)
__device__ constexpr auto is_predicated() const;
__device__ auto is_thread_active() const;

// Copies results to a register fragment and returns it to the user
__device__ auto get_results() const;
__device__ auto make_empty_fragment() const;

// Copies results to a user-provided register fragment
template<class Engine, class Layout>
__device__ void get_results(cublasdx::tensor<Engine, Layout>& out) const;

As mentioned in Accumulator And Register Fragment Tensors the division of elements among threads happens by tiling an MMA instruction over the BLAS defined problem size. Each instruction is responsible for computing certain shape, and it’s possible that the problem shape is not divisible by primitive instruction shape. In such cases the “extra” instruction elements are filled with 0s instead of reading from memory, and skipped when storing results. This is called element predication. The is_predicated() member function can be used to retrieve that information; it is also true when some threads of the block are inactive (see below).

Due to cuBLASDx supporting execution within kernels with CUDA thread block sizes not matching the BlockDim Operator, not always will all threads take part in a GEMM operation. This implies that some threads may not have any elements assigned to them. is_thread_active() member function can be used to check exactly whether this is the case for the calling thread.

The values returned by is_predicated() and is_thread_active() can be used directly in conditions (they are contextually convertible to bool). Whether the value can be known at compile time or only at runtime will depend on the specific BLAS configuration used.

The number of threads used by a GEMM can be smaller than the block size. is_thread_active() is false for every thread outside of the tiled MMA, and is_predicated() is true whenever the C matrix is not divisible by the tiled instruction shape or some threads of the declared block do not take part in the GEMM. Complex integer GEMMs are not capped and keep every block thread in the tiled MMA; real integer and floating-point matrix-vector products use fewer threads than the block.

For register accumulators, the memory operations (partition_and_copy(), partition_and_store(), process_and_store(), axpby(), reduce_and_store()) are no-ops for inactive threads, so they need no is_thread_active() guard (tensor memory accumulators exist only in pipelines, where all participating threads are active, and still require active threads); active threads use the regular vectorized copies, with per-element predication only when C is not divisible. get_results() returns unspecified values on inactive threads of register accumulators. is_index_in_bounds() is false for every index of an inactive thread of a register accumulator. User functors passed to process_and_store() and reduce_and_store() may run on partially active warps, so they must not use full-mask warp collectives (for example __shfl_sync(0xffffffff, ...)).

The clear() (returns void) will effectively clear the accumulator for the next MMA execution. A reusable (user-owned) accumulator keeps its running sum across executes, but the result is kept and can be reused only until finish_accumulation() is called: each finish_accumulation() requires a subsequent clear() before the next accumulation. Accumulators owned by a pipeline (internal accumulation) start from zero for every execution.

Warning

.clear() method may just set a special flag to next performed multiplication, it’s valid only in connection with subsequent .execute() call.

The final argumentless method - make_empty_fragment() can be used to retrieve a non-initialized register fragment tensor corresponding to BLAS execution from which the accumulator was retrieved.

Please see Accumulator Example #1 below for an example of use of these methods.

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

This method will return a non-owning view of its argument’s subtensor assigned to the calling thread, corresponding to its local register fragment.

All these methods can be used for manual copying of unpredicated GEMM as follows:

Accumulator Example #1

auto c_global_tensor = ...;

// Get accumulator
auto accumulator = BLAS::get_accumulator();
// Create local register fragment
auto d_register_fragment = accumulator.make_empty_fragment();
// Get view of this thread's global memory subtensor
auto c_global_partition = accumulator.partition_like_C(c_global_tensor);

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

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

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

cublasdx::axpby(alpha, accumulator.get_results(), beta, d_register_fragment);

If a GEMM is predicated, more information will be necessary:

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;

template<class Coord>
__forceinline__ __device__
bool is_coord_in_bounds(Coord const& coord) const;

These functions extend is_predicated() by mapping a local register-fragment index to its source (global or shared) tensor index, as well as checking whether a fragment index (is_index_in_bounds) or an already-mapped (m, n) coordinate (is_coord_in_bounds) lies inside the (M, N) problem shape.

map_fragment_index(value) preserves compile-time information when value is a compile-time fragment index, allowing the implementation to fold the per-value part of the mapping. The same call also accepts runtime fragment indices. Both forms are portable across RMEM, emulated RMEM, and TMEM accumulator backends and preserve the exact coordinates of padded slots. Such coordinates remain outside the problem shape and can be checked with is_coord_in_bounds() or is_index_in_bounds().

All these methods can be used to perform a bound-checked loading of elements from global to register tensor:

Accumulator Example #2

auto c_global_tensor = ...;

// Get accumulator
auto accumulator = BLAS::get_accumulator();
// Create local register fragment
auto d_register_fragment = accumulator.make_empty_fragment();

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

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

cublasdx::axpby(alpha, accumulator.get_results(), beta, d_register_fragment);

It would be very tedious to have to do such loading manually every time, hence cuBLASDx offers optimized and auto-vectorizing global/shared ⟷ register fragment copying functions. See Copying Register Tensors for more information

Reduce and Store

template<class CTensor, class Lambda>
__forceinline__ __device__
void reduce_and_store(CTensor && ctensor, Lambda && lambda) const;

reduce_and_store(output, lambda) combines the accumulator results with an existing output tensor element by element, without the caller having to partition, load, or index anything manually. The output can be a global or shared CTA-wide (M, N) tensor, or a register fragment matching the accumulator’s fragment.

For every element owned by the calling thread, the lambda is invoked as:

new_output_value = lambda(result_element, output_element, coord);

where coord is the element’s (m, n) coordinate within the CTA output tile - the same value map_fragment_index(...) returns for the corresponding fragment index. Its components can be accessed with structured bindings or cublasdx::get<0> / cublasdx::get<1>. For CTA-wide memory outputs, elements that fall outside the (M, N) problem shape of a predicated GEMM are skipped: they are never handed to the lambda and their output slots are left untouched. A register output must match the accumulator fragment size and the lambda runs over the whole fragment - including padded slots, whose coordinates lie outside (M, N); lambdas that index by coordinate can guard against them with is_coord_in_bounds(coord).

Accumulator Example #3

auto c_global_tensor = ...;

// Get accumulator and run GEMM
auto accumulator = BLAS::get_accumulator();
BLAS().execute(a_shared_tensor, b_shared_tensor, accumulator);

// C = alpha * (A * B) + beta * C with a per-element scale lookup
accumulator.reduce_and_store(c_global_tensor, [&](auto const& result, auto const& c, auto const& coord) {
    const auto [row, column] = coord;
    return alpha * row_scale(row) * result + beta * c;
});

Imported Tensor Utilities

Since cuBLASDx integrated CuTe library, it exposes some of its tensor and layout functionality, sometimes adding compatibility layers to enable its own types support. Currently following functions are supported:

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

// Wrapper (forwards its arguments to cute::transform)
cublasdx::transform_fragment;

// Layout and tensor size queries, provided by commonDx
cublasdx::size;
cublasdx::cosize;

// Native cuBLASDx implementation
cublasdx::axpby;

cublasdx::clear takes a register fragment as input and sets all its values to 0 in place.

auto accumulator = BLAS::get_accumulator();
auto c_register_fragment = accumulator.make_empty_fragment();

cublasdx::clear(c_register_fragment);

cublasdx::transform_fragment applies an elementwise functor to all elements of passed register tensor. It forwards its arguments to cute::transform, and supports the following forms:

// In-place unary: tensor(i) = op(tensor(i))
cublasdx::transform_fragment(tensor, op);

// Out-of-place unary: dst(i) = op(src(i))
cublasdx::transform_fragment(src, dst, op);

// Binary: dst(i) = op(src_a(i), src_b(i))
cublasdx::transform_fragment(src_a, src_b, dst, op);

auto c_global_tensor = ...

auto accumulator = BLAS::get_accumulator();
auto c_register_fragment = accumulator.make_empty_fragment();

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

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

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

cublasdx::make_fragment_like can be used to create a register fragment with the same layout and type as its argument tensor.

auto c_global_tensor = ...

auto accumulator = BLAS::get_accumulator();
auto c_global_partition = accumulator.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_other = cublasdx::make_fragment_like<new_type>(c_global_partition);

Warning

The following functions should be used only in their cuBLASDx form, as it offers compatibility layers and automatic conversions.

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

It operates per thread on the whole tensor passed to it, so it is intended for per-thread register fragments. Calling it from several threads on the same shared (or global) memory tensor is a data race. For a memory-backed C matrix, use accumulator.axpby(alpha, beta, c_tensor) instead, which partitions the tensor among threads.

When beta is zero, the contents of D are not read (it may be uninitialized; its contents, including NaNs, do not affect the result).

auto c_global_tensor = ...

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

auto accumulator = BLAS::get_accumulator();
auto d_register_fragment = accumulator.make_empty_fragment();
cublasdx::copy_fragment<alignment_of<BLAS>::c>(c_global_tensor, d_register_fragment, accumulator);

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

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

cublasdx::size returns a number of valid elements in a layout or tensor. This is simply a product of all shape dimensions. It accepts a layout (including the layouts returned by get_layout_* and suggest_layout_* methods) or a tensor.

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 returns the size of the codomain of a layout, i.e. the largest index it maps to plus one. It describes how many elements does the argument layout span. Like cublasdx::size, it accepts a layout (including the layouts returned by get_layout_* and suggest_layout_* methods) or a tensor. This is the same as cublasdx::size for layouts which are compact, or have no “holes” (e.g. as result of extra leading dimension elements)

// Default leading dimension
using 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
using BLAS_padded = decltype(... + Size<M, N, K>() + LeadingDimension<lda, K, M>());
auto a_padded_layout = BLAS_padded::get_layout_smem_a();

constexpr auto padded_size = cublasdx::size(a_padded_layout);
constexpr auto padded_cosize = cublasdx::cosize(a_padded_layout);

static_assert(padded_size != padded_cosize);
static_assert(padded_size < padded_cosize);

Copying Tensors

Cooperative Global ⟷ Shared Copying

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 is a helper function for copying data between tensors that are either in shared or global memory.

The copy is done cooperatively. All threads, indicated either by NumThreads or by BLAS::block_dim, will participate in the copy. The function accounts for the given alignments and attempts to vectorize load and store instructions when possible.

Requirements:

• Data in tensors has to be in shared or global memory. Copying from or to registers is not supported.

• Both src and dst tensors must represent tensors of the same underlying element types (cublasdx::tensor<Engine, Layout>::value_type, Engine::value_type).

• Both src and dst tensors must have the same size, i.e. the number of elements.

• AlignmentInBytes must be a multiple of the alignment of the underlying element type of tensors.

• AlignmentInBytes must be equal to 1, 2, 4, 8 or 16, or equal to the alignment of the underlying element type of tensors.

• Underlying pointers in src and dst tensors must be aligned to AlignmentInBytes bytes.

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

cublasdx::copy_wait creates a synchronization point. It has to be called after cublasdx::copy operation, before any consequent read or write to dst tensor, and before any consequent write to src tensor. Otherwise, the result of copying operation is undefined. It is not a plain __syncthreads(): in addition it waits for all outstanding asynchronous copies (see cp.async family of instructions)

If a copy is the final operation that writes a result to global memory and no later instruction in the same kernel reads or overwrites the source or destination tensor, examples may rely on kernel completion instead of an explicit trailing copy_wait(). Add copy_wait() whenever the copied data participates in more work inside the same kernel.

Example

Example of copying A matrix from global to shared memory and back.

using BLAS = decltype(Size<128, 128, 128>() + Type<type::real>() + Precision<float, float, double>() + Block() + ...);
extern __shared__ __align__(16) cublasdx::byte 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();

Copying Register Tensors

Register fragments can be copied to either global or shared memory tensors manually using accumulator API, or with use of cublasdx::copy_fragment.

Accumulator API allows for 2 distinct approaches to copying:

1. Create a subtensor view of calling thread’s elements and copy using register fragment index space (see Accumulator Example #1).

2. Map local register fragment index to global index space and address global / shared tensor directly, allowing for location-aware processing (see Accumulator Example #2).

The second way is based on a bidirectional copying method cublasdx::copy_fragment, which works per-thread:

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

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

Depending on the defined semantics, these can be thought of as gather and scatter functions:

1. From global / shared tensor perspective: load fragment is scatter, store fragment is gather.

2. From thread perspective: load fragment is gather, store fragment is scatter.

It’s important to notice that each of these functions takes:

1. a template argument unsigned integer describing alignment of shared / global memory tensor,

2. this thread’s register fragment tensor and an entire non-partitioned global / shared memory tensor, ordered by copy direction, and

3. an accumulator.

They are both safe to use with thread and value predicated fragment GEMMs, and will auto-vectorize to the extent described by passed alignment.

Note

cublasdx::copy_fragment is a per-thread operation and contains no synchronization. After storing a fragment into a shared memory tensor, __syncthreads() is required before other threads read that tensor, including a block-wide cublasdx::copy of the tile. The same applies to the accumulator methods which write to shared memory tensors (partition_and_store(...), partition_and_copy(...), process_and_store(...), reduce_and_store(...), axpby(...)).

Example of use, to perform register API GEMM with accumulator:

auto accumulator = GEMM::get_accumulator();

// Load fragment
auto d_fragment = accumulator.make_empty_fragment();
cublasdx::copy_fragment<alignment::c>(c_global_tensor, d_fragment, accumulator);
// alternatively
// auto d_fragment = accumulator.make_partition_and_copy(c_global_tensor);

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

// Store fragment
cublasdx::axpby(alpha, accumulator.get_results(), beta, d_fragment);
cublasdx::copy_fragment<alignment::c>(d_fragment, c_global_tensor, accumulator);

// Alternatively, partition, load, axpby and store in one call
// accumulator.axpby(alpha, beta, c_global_tensor);