CUDA Matrix Transpose only with warp shuffle instructions not working - cuda

Question

Code: running example
#include <iostream>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
__global__ void transpose_block(float* mat, size_t sx, size_t sy)
{
constexpr size_t size = 8;
auto tile = cg::tiled_partition<size>(cg::this_thread_block());
auto ix = tile.thread_rank();
float col[size];
for (size_t iy = 0; iy < size; ++iy)
col[iy] = mat[ix + iy * sx];
auto val = [&tile, &col](int ix, int iy) { return tile.shfl(col[iy], ix); };
for (size_t iy = 0; iy < size; ++iy)
mat[ix + iy * sx] = val(iy, ix);
}
void print_mat(float* mat, size_t sx, size_t sy)
{
printf("{\n");
for (size_t iy = 0; iy < sy; ++iy)
{
printf("\t{ ");
for (size_t ix = 0; ix < sx; ++ix)
printf("%6.1f, ", mat[ix + iy * sx]);
printf("},\n");
}
printf("}\n");
}
int main()
{
constexpr size_t sx = 8;
constexpr size_t sy = 8;
float* mat;
cudaMallocManaged(&mat, sx * sy * sizeof(float));
for (size_t iy = 0; iy < sy; ++iy)
for (size_t ix = 0; ix < sx; ++ix)
mat[ix + sx * iy] = ix + sx * iy;
print_mat(mat, sx, sy);
transpose_block<<<1, 8>>>(mat, sx, sy);
cudaDeviceSynchronize();
print_mat(mat, sx, sy);
cudaFree(mat);
}
I'm trying to transpose an 8x8 block by saving the columns locally for each of the 8 threads and then writing a lambda that would essentially give me access to the entire block using warp shuffle:
// tile = the cooperative group that holds all 8 threads
// col = local vector that holds all 8 column values of column[tile.thread_rank()]
auto val = [&tile, &col](int ix, int iy) { return tile.shfl(col[iy], ix); };
This doesn't work and I just cannot find the reason why.

Answer 1

When I run your code I get output like this:
{
{ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, },
{ 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, },
{ 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, },
{ 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, },
{ 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0, },
{ 40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0, },
{ 48.0, 49.0, 50.0, 51.0, 52.0, 53.0, 54.0, 55.0, },
{ 56.0, 57.0, 58.0, 59.0, 60.0, 61.0, 62.0, 63.0, },
}
{
{ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, },
{ 9.0, 9.0, 9.0, 9.0, 9.0, 9.0, 9.0, 9.0, },
{ 18.0, 18.0, 18.0, 18.0, 18.0, 18.0, 18.0, 18.0, },
{ 27.0, 27.0, 27.0, 27.0, 27.0, 27.0, 27.0, 27.0, },
{ 36.0, 36.0, 36.0, 36.0, 36.0, 36.0, 36.0, 36.0, },
{ 45.0, 45.0, 45.0, 45.0, 45.0, 45.0, 45.0, 45.0, },
{ 54.0, 54.0, 54.0, 54.0, 54.0, 54.0, 54.0, 54.0, },
{ 63.0, 63.0, 63.0, 63.0, 63.0, 63.0, 63.0, 63.0, },
}
Let's start by making sure you understand how a shuffle op works so that we can at least explain that output. Your shuffle op exchanges data like this:
tile.shfl(col[iy], ix)
^ ^
| The lane that I go to to get the value I will return.
The value I publish to other lanes.
So when you do that shuffle op in the for-loop:
for (size_t iy = 0; iy < size; ++iy)
mat[ix + iy * sx] = val(iy, ix);
Let's first note that although the order of parameters in your lambda definition is (ix, iy) you are passing val(iy, ix). Confusing! Maybe you like programming that way. I find it confusing.
Anyway, what we observe is that the first index (you are passing iy) will be used to select the lane to copy the value from. Therefore all 8 lanes will target the same source lane on each iteration of the for loop. This explains why the values are constant in the output across a row. The second index (you are passing ix) will be used to select the column entry to publish to other lanes. Since ix is zero for thread 0, 1 for thread 1, etc., and also loop-invariant, we observe that thread 0 always publishes its col[0], thread 1 always publishes its col[1], and so on. This explains the pattern as we look down a column in the output.
You might think at this point that it was a mistake to pass iy for ix and ix for iy, but if we reverse the order of values that we pass, then the output matrix simply duplicates the input matrix (exercise left to the reader.)
We're going to have to be more inventive than that. If we are going to keep the loop variable iy unmodified, on the first pass of the for loop, we would like thread 0 to keep its value in col[0], we would like thread 1 to retrieve its value from col[1] of thread 0, and we would like thread 2 to get the value of col[2] of thread 0, and so on. On the first pass of the for-loop, we would like the entire column from thread 0 to be passed to the other threads. But this is impossible with a shuffle - we can only publish one value (per thread) at a time.
Therefore we reach the conclusion that the naive approach of processing the matrix row-by-row will not work. We must be more inventive than that. After some headscratching, you may realize that a suitable pattern can be obtained using an xor operation on the thread index. Effectively, the xor indexing pattern creates a set of swaps or exchanges around the main diagonal. Rather than try and describe it in great detail, I refer you here (and here). The original posting there in the "suggested solution" block shows the xor pattern, and how it is used to specify both the column value to publish, as well as the source lane, for each thread's shuffle op.
If we apply that concept to your code, we could come up with this modification to your kernel:
$ cat t2170.cu
#include <iostream>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
__global__ void transpose_block(float* mat, size_t sx, size_t sy)
{
constexpr size_t size = 8;
auto tile = cg::tiled_partition<size>(cg::this_thread_block());
auto ix = tile.thread_rank();
float col[size];
for (size_t iy = 0; iy < size; ++iy)
col[iy] = mat[ix + iy * sx];
auto val = [&tile, &col](int ix, int iy) { return tile.shfl(col[iy], ix); };
for (size_t idx = 1; idx < size; ++idx){
size_t iy = threadIdx.x^idx;
mat[ix + iy * sx] = val(iy, iy);}
}
void print_mat(float* mat, size_t sx, size_t sy)
{
printf("{\n");
for (size_t iy = 0; iy < sy; ++iy)
{
printf("\t{ ");
for (size_t ix = 0; ix < sx; ++ix)
printf("%6.1f, ", mat[ix + iy * sx]);
printf("},\n");
}
printf("}\n");
}
int main()
{
constexpr size_t sx = 8;
constexpr size_t sy = 8;
float* mat;
cudaMallocManaged(&mat, sx * sy * sizeof(float));
for (size_t iy = 0; iy < sy; ++iy)
for (size_t ix = 0; ix < sx; ++ix)
mat[ix + sx * iy] = ix + sx * iy;
print_mat(mat, sx, sy);
transpose_block<<<1, 8>>>(mat, sx, sy);
cudaDeviceSynchronize();
print_mat(mat, sx, sy);
cudaFree(mat);
}
$ nvcc -arch=sm_70 -o t2170 t2170.cu
$ compute-sanitizer ./t2170
========= COMPUTE-SANITIZER
{
{ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, },
{ 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, },
{ 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, },
{ 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, },
{ 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0, },
{ 40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0, },
{ 48.0, 49.0, 50.0, 51.0, 52.0, 53.0, 54.0, 55.0, },
{ 56.0, 57.0, 58.0, 59.0, 60.0, 61.0, 62.0, 63.0, },
}
{
{ 0.0, 8.0, 16.0, 24.0, 32.0, 40.0, 48.0, 56.0, },
{ 1.0, 9.0, 17.0, 25.0, 33.0, 41.0, 49.0, 57.0, },
{ 2.0, 10.0, 18.0, 26.0, 34.0, 42.0, 50.0, 58.0, },
{ 3.0, 11.0, 19.0, 27.0, 35.0, 43.0, 51.0, 59.0, },
{ 4.0, 12.0, 20.0, 28.0, 36.0, 44.0, 52.0, 60.0, },
{ 5.0, 13.0, 21.0, 29.0, 37.0, 45.0, 53.0, 61.0, },
{ 6.0, 14.0, 22.0, 30.0, 38.0, 46.0, 54.0, 62.0, },
{ 7.0, 15.0, 23.0, 31.0, 39.0, 47.0, 55.0, 63.0, },
}
========= ERROR SUMMARY: 0 errors
$