CUDA Shared Memory Swizzling - Lei Mao's Log Book

CUDA Shared Memory Swizzling<br>05-14-2024 07-31-2024 blog 26 minutes read (About 3899 words) visits

Introduction<br>When we write CUDA kernels that use shared memory, we have to be careful about shared memory bank conflicts. Having severe shared memory bank conflicts can introduce a significant performance penalty.

One simple way to deal with shared memory bank conflicts is to use padding. However, padding can waste shared memory and can have other drawbacks.

In this blog post, I would like to discuss how to deal with shared memory bank conflicts using swizzling. Swizzling is a more complicated technique that can be used to avoid shared memory bank conflicts without wasting shared memory.

CUDA Shared Memory Swizzling<br>Swizzling Example<br>When we use CUDA shared memory to cache data without using padding, it’s very common that either reading from or writing to shared memory by a warp can cause shared memory bank conflicts. Swizzling is a technique that rearranges the mapping of the shared memory index to avoid shared memory bank conflicts. Matrix transpose is a perfect example that can have shared memory bank conflicts if the implementation does not use padding or swizzling.

In the above example, the shared memory is a 2D array of float with size 32 × 16. In terms of matrix transpose, each warp reads a row of 32 values from the global memory and writes them to the shared memory with swizzling. There will be no shared memory bank conflicts when writing to the shared memory. To perform matrix transpose, each wrap reads two swizzled “columns” of 32 values from the shared memory and writes them to the global memory. For example, the swizzled column 0 and 1 are colored in yellow and cyan, respectively. In this way, there will be only one shared memory bank conflict when reading from the shared memory. Without using swizzling, there will be 16 (2-way) shared memory bank conflicts when reading from the shared memory. Of course, obviously, if the shared memory is a 2D array of float with size 32 × 32, there will be no shared memory bank conflicts when writing to the shared memory and reading from the shared memory.

Swizzling Formula<br>Given an array of T array[][NX] on shared memory, we define NX × sizeof(T) == SWIZZLE_SIZE. The allowed values of SWIZZLE_SIZE are powers of 2 that is larger than or equal to 32, such as 32, 64, 128, 256, …, etc.

Given the index [y][x] in T array[][NX], we can compute the swizzled index x_swz as follows:

Compute the index of the TC-byte chunk within SWIZZLE_SIZE-byte segment:<br>i_chunk = (y × NX + x) × sizeof(T) / sizeof(TC)<br>y_chunk = i / (SWIZZLE_SIZE / sizeof(TC))<br>x_chunk = i % (SWIZZLE_SIZE / sizeof(TC))

Compute the swizzled index of TC-byte chunk using XOR operation:<br>x_chunk_swz = y_chunk ^ x_chunk

Compute swizzled index:<br>x_swz = x_chunk_swz × sizeof(TC) / sizeof(T) % NX + x % (sizeof(TC) / sizeof(T))

Swizzling Properties<br>This swizzling formula has the following properties:

The index before and after swizzling must be one to one mapped.

$\text{NX}$ must be a power of 2.

Given any $x$ and any $\{y, y+1, y+2, \cdots, y+31\}$, the number of unique swizzled index $x_{\text{swz}}$ should be maximized.

The property 1 ensures that there will be no data loss during swizzling. The property 2 ensures that the index before and after swizzling will be one to one mapped.

Here I am going to show some informal mathematical proofs for the properties from the swizzling formula.

Proof

We will first prove the property 1.

$$<br>\begin{align}<br>x_{\text{chunk}}<br>&= i_{\text{chunk}} \% (\text{SWIZZLE_SIZE} / \text{sizeof}(\text{TC})) \\<br>&= \left(\left(y × \text{NX} + x\right) × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})\right) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \\<br>&= \left(y × \text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) + x × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})\right) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \\<br>&= \left(y × \text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) + x × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \right) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \\<br>&= \left(x × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \right) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \\<br>&= x × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) \% (\text{NX} × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC})) \\<br>&= \left( x \% \text{NX} \right) × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) \\<br>&= x × \text{sizeof}(\text{T}) / \text{sizeof}(\text{TC}) \\<br>\end{align}<br>$$

It seems that we have derived another equivalent formula for $x_{\text{chunk}}$. Note that...