Clarification on template BLOCK_DIM_Y parameter

[johnmansell-dyndrite]

Problem

I'm having trouble understanding how to utilize the BLOCK_DIM_Y parameter in cub templates. From what I can see, all the examples only use 1D templates. The documentation states:

BLOCK_DIM_Y – [optional] The thread block length in threads along the Y dimension (default: 1)

However it is not behaving the way I would expect, so I think I must have a misunderstanding of what is supposed to be happening, or how one should process 2D data.

Example

I've created an example which is the simplest case demonstrating my confusion.

• Create a 2D matrix (32 x 64)
• Copy to device with cudaMemcpy2D
• block size = (32, 32, 1)
• grid size = (1, 1, 1)
• Items / thread = 2
• Load global memory into thread memory with cub::BlockLoad
• Copy data to a new matrix

I've created the matrix with int2 values so each (x,y) pair can be tracked to see where it ends up, or what's missing. The example works as is, but there seems to be some unknown constraints or something I'm not understanding. If I change the matrix dimensions, one of the following happens:

• The data ends up in the wrong location in the result matrix
• The result matrix is all zeros
• Some of the result matrix is blank

The copy seems to work correctly if

• BLOCK_DIM_X == ROWS * ITEMS * (constant)
  and
• BLOCK_DIM_Y == COLS * ITEMS * (constant)

Question

• Can you please explain how BLOCK_DIM_Y is used in the templates, and if it's possible to iterate over the data in this manner?
• What am I not getting about the templates when extended from 1D to 2D?
• Are there any examples of using the cub templates in a 2D block rather than a 1D block?

[pauleonix]

Have you noted the

• For multi-dimensional blocks, threads are linearly ranked in row-major order.

in the documentation, e.g. the last point here. I.e. the algorithm seems to allow for multi-dimensional blocks, but they are flattened. The algorithm itself is always one-dimensional.

[pauleonix]

Also note (emphasis mine)

The BlockLoad class provides collective data movement methods for loading a linear segment of items from memory into a blocked arrangement across a CUDA thread block.

Loading more than one row from 2D data is generally not a "linear segment".

[johnmansell-dyndrite]

I did notice that. Unfortunately that only adds to my confusion.

1. If the block is flattened, why include a template parameter for x, y, z dimensions rather than just a flat total block size?
2. The block load doesn't seem to load flattened data.

Given a 128 x 128 matrix, I would expect data loading to either happen like

Flattened

// Block Row 0
    data = [0, 1] [2, 3] [4, 5] .... [63, 63]

// Block Row 1
    data = [64, 65], [66, 67] ... [126, 127]

Not Flattened

// Block Row 0
    data = [0, 1] [2, 3] [4, 5] .... [63, 63]

// Block Row 1
    data = [128, 129] [130, 131] ... [191, 192]

The problem is that neither of these seems to happen consistently from what I can tell. In the example I gave, if I alter it to do only one read / write cycle, the data is loaded from each row. (not flattened).

/* 
ROWS = 128
COLS = 128
*/

__global__ void copy_matrix(int2 *input, int2 *result, size_t pitch_in_bytes)
{
    ...
    // --- Tile Dims
    constexpr auto tile_width = BLOCK_SIZE.x * ITEMS;  // 32 * 2
    constexpr auto tile_height = BLOCK_SIZE.y; // 32

    // --- Iterate over tiles in the X direction
    for (int tile=0; tile < 1; tile++) {

        // --- Load Data
        auto valid_items = ROWS * COLS;
        mem_loader.Load(&input_tile[tile * tile_width], thread_data, valid_items);

        // --- Write
        auto offset = (tile_width * tile) + (x * ITEMS);
        for (int i=0; i<ITEMS; i++) {
            result_row[offset + i] = thread_data[i];
        }
    }
}

Output

0   [ (0, 0)  (1, 0)  (2, 0)  (3, 0) ... ( 62, 0)  ( 63, 0)  (0, 0) (0, 0) ... (0, 0) ]
1   [ (0, 1)  (1, 1)  (2, 1)  (3, 1) ... ( 62, 1)  ( 63, 1)  (0, 0) (0, 0) ... (0, 0) ]
2   [ (0, 2)  (1, 2)  (2, 2)  (3, 2) ... ( 62, 2)  ( 63, 2)  (0, 0) (0, 0) ... (0, 0) ]

• Note that each row loads tile_width items from it's corresponding row in data.
• Columns greater than tile_width are (0, 0)
• All column 0 values are (0, y)

[pauleonix]

@johnmansell-dyndrite Well, to compute the 1D index from threadIdx.x, threadIdx.y and threadIdx.z you need to know their extents (i.e. (threadIdx.z * BLOCK_DIM_Y + threadIdx.y) * BLOCK_DIM_X + threadIdx.x). The total number of threads is not enough. Sure one could query blockDim at runtime, but I guess the engineers wanted to have as much optimization potential as possible by making the values available at compile time for all block-algorithms.

Your "hack" might be working because these BLOCK_DIM_ dimensions are only use to calculate offsets into global memory so if you specify big numbers intermediate values are just not loaded by their corresponding non-existent threads (i.e. it is trying to load flattened data but relying on threads that are not there). I would not assume that this kind of usage is officially supported. It might be considered UB, i.e. behavior could change in future releases. On the other hand CUB could easily check these parameters at compile time and error out on values that are impossible to launch as actual block sizes. Let's wait what the devs are saying.

[gevtushenko]

@johnmansell-dyndrite thank you for reaching out! We should provide examples illustrating CUB usage in the context of multi-dimensional thread blocks. I'll file an issue to make sure this is tracked somewhere.

As to your question, I believe there are multiple factors at play.
Let's start with the block load. Block load accepts an iterator and reads N consecutive elements from it. For instance, let's consider the following input:

int input[] = { 0, 1, 2, 3, 4, ... }

If you load from input using single-dimensional block load with 2 items per thread and, you'll get:

thread{x=0}: thread_data[2] = {0, 1};
thread{x=1}: thread_data[2] = {2, 3};
...

Now, let's see what happens with multi-dimensional thread block.
When you add Y dimension, block load internals linearize thread indexing as follows:

linear_tid(RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))

As @pauleonix rightfully noted, having block dimensions at compile time allows us to avoid unnecessary linearized thread index computations. If the block is defined as BLOCK_DIM_X = 32, BLOCK_DIM_Y = 1, ..., linearized index computation will not include the y and z terms. This is an optimization.

From the functional point of view, nothing else changes, CUB internals use linear_tid as thread "index", so that:

thread{x= 0, y=0}: linear_tid{0}: thread_data[2] = {0, 1};
thread{x= 1, y=0}: linear_tid{1}: thread_data[2] = {2, 3};
...
thread{x=31, y=0}: linear_tid{31}: thread_data[2] = {62, 63};
thread{x= 0, y=1}: linear_tid{32}: thread_data[2] = {64, 65};
...

I believe that the main issue that introduces confusion in your case is cudaMallocPitch. It allocates data as:

        COL0    COL1    COL2    COL3
ROW0   {0, 0}, {1, 0}, {2, 0}, {3, 0}   /* -- padding -- */
ROW1   {0, 1}, {1, 1}, {2, 1}, {3, 1}   /* -- padding -- */
ROW2   {0, 2}, {1, 2}, {2, 2}, {3, 2}   /* -- padding -- */

The linearized data looks like:

{0, 0}, {1, 0}, {2, 0}, {3, 0}, { -- padding -- }, {0, 1}, {1, 1} ...

In the code you attached, the padding is 0, because pitch is equal to 512 bytes, int2 is 8 bytes, giving us 512 / 8 = 64, which matches COLS = 64. That means that block load reads the full matrix in one call, and the rest of the for (int tile=0; tile < n_tiles; tile++) { is loading the same data 32 times. When the padding is present, results are undefined. The padding bits can contain zeroes or some left-over data from previous computations.

I hope this helps describe the behavior you observe. Please, let us know if you find other gaps in our documentation. This would greatly help us improve it.

[johnmansell-dyndrite]

@gevtushenko Thank you for responding and clearing that all up.

• I do think extra documentation would be great. The snippets cited by @pauleonix are helpful, but at least for me still left some ambiguity with how I expected it to work based on the inclusion of the BLOCK_DIM_Y parameter.
• I really appreciate the power and portability of CUB as compared to Thrust or Modern GPU, but the template nature does make it a bit more opaque and harder to discern for myself what's happening by stepping through the code. So documentation and examples would be really helpful.

Follow Up Request

Is there a way in CUB or another library to do the kind of iteration I was attempting?

I often find myself working on problems that require iterating over the data one tile at a time. I've always used my own kernels for load, scan, reduce, etc. I'd assumed I was just using the CUB parameters incorrectly, but now I know it's not designed for what I was attempting.

I believe it'd be a very powerful option to add a 2D version of some of the CUB block primitives for cases like this. If the CUB kernels / functions had parallels to cudaMemcpy and cudaMemcpy2D it would add a lot more flexibility to the library. An equivalent pitch_in_bytes parameter probably wouldn't be known at compile time, but if it's passed in to the load function similar to valid_items then each row in the Y direction could know the offset to begin loading or storing memory.

Alternatively, extending the warp wide primitives to cover more of the algorithms included in the block wide options would do essentially the same thing. A warp load function could take in a pointer to the beginning of the row in a 2D matrix and iterate over it. Though I'm guessing the nature of how structs work in CUDA will be the biggest hurdle. My understanding is that the whole block shares a single copy of the struct, so re-purposing it to do warp-wide loads rather than block loads might have it's own hurdles to overcome.

Documentation

I'd be happy to contribute to the documentation with examples or otherwise if that'd be helpful. However, since I had a fundamental misunderstanding of the internals, and got the right output the wrong way, I might not be the best candidate to write up an explanation.

[pauleonix]

@johnmansell-dyndrite

A warp load function could take in a pointer to the beginning of the row in a 2D matrix and iterate over it.

cub::WarpLoad does exist.

Though I'm guessing the nature of how structs work in CUDA will be the biggest hurdle. My understanding is that the whole block shares a single copy of the struct, so re-purposing it to do warp-wide loads rather than block loads might have it's own hurdles to overcome.

That is not true, structs work the same as in C++, every thread has it's own copy as long as the struct is not placed in shared memory. I.e. for BlockLoad every thread has its own instantiation, but - if used correctly - those instantiations share (i.e. have the same value for) the reference or pointer to the TempStorage in shared memory. The abstraction just relies on the user instantiating these structs with the same data for all threads.