CUDA Shared Memory Swizzling NVIDIA's CUDA programming guide introduces a swizzling technique to eliminate shared memory bank conflicts without wasting memory, using an XOR-based formula that rearranges shared memory indices. The method, detailed in a blog post, ensures one-to-one mapping and maximizes unique swizzled indices for warps, reducing bank conflicts from 16 to 1 in a 32×16 float matrix transpose example. CUDA Shared Memory Swizzling Introduction When we write CUDA kernels that use shared memory, we have to be careful about shared memory bank conflicts /blog/CUDA-Shared-Memory-Bank/ . 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 Swizzling Example 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 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: i chunk = y × NX + x × sizeof T / sizeof TC y chunk = i / SWIZZLE SIZE / sizeof TC x chunk = i % SWIZZLE SIZE / sizeof TC Compute the swizzled index of TC -byte chunk using XOR operation https://en.wikipedia.org/wiki/XOR gate : x chunk swz = y chunk ^ x chunk Compute swizzled index: x swz = x chunk swz × sizeof TC / sizeof T % NX + x % sizeof TC / sizeof T Swizzling Properties 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. $$ \begin{align} x {\text{chunk}} &= i {\text{chunk}} \% \text{SWIZZLE SIZE} / \text{sizeof} \text{TC} \\ &= \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} \\ &= \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} \\ &= \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} \\ &= \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} \\ &= x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \% \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \\ &= \left x \% \text{NX} \right × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \\ &= x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \\ \end{align} $$ It seems that we have derived another equivalent formula for $x {\text{chunk}}$. Note that $\text{sizeof} \text{T} / \text{sizeof} \text{TC} $ is a bit right shifting operation when $\text{sizeof} \text{TC} / \text{sizeof} \text{T} $ is a power of 2. $$ \begin{align} y {\text{chunk}} &= i {\text{chunk}} / \text{SWIZZLE SIZE} / \text{sizeof} \text{TC} \\ &= \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} \\ &= \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} \\ &= 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} \\ &= y + x / \text{NX} \\ &= y \\ \end{align} $$ It seems that we have also derived another equivalent formula for $y {\text{chunk}}$. $$ \begin{align} x {\text{chunk swz}} &= y {\text{chunk}} \oplus x {\text{chunk}} \\ &= y \oplus \left x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \right \\ &= y / \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} × \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} + \left y \% \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \right \oplus \left x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \right \\ &= y / \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} × \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} + \left \left y \% \text{NX} \right × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \right \oplus \left x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \right \\ &= y / \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} × \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} + \left y \% \text{NX} \right \oplus x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \\ \end{align} $$ Note that $\oplus$ is the bitwise XOR operation. If either $y {\text{chunk}}$ or $x {\text{chunk}}$ is a constant, the mapping is a one to one mapping. $$ \begin{align} x {\text{swz}} &= x {\text{chunk swz}} × \text{sizeof} \text{TC} / \text{sizeof} \text{T} \% \text{NX} + x \% \text{sizeof} \text{TC} / \text{sizeof} \text{T} \\ \end{align} $$ Here the proof becomes a little bit informal. Because a consecutive of $\text{sizeof} \text{TC} / \text{sizeof} \text{T} $ $x$ values will be mapped to one unique trunk index $x {\text{chunk}}$, the mapping between $x {\text{chunk}}$ and $x {\text{chunk swz}}$ is a one to one mapping, one $x {\text{chunk swz}}$ value will map to one unique $x {\text{chunk swz}} × \text{sizeof} \text{TC} / \text{sizeof} \text{T} \% \text{NX}$ value. To create the one to one mapping between the swizzled index $x {\text{swz}}$ and the original index $x$, the offset $x \% \text{sizeof} \text{TC} / \text{sizeof} \text{T} $ is added. Therefore, the index before and after swizzling must be one to one mapped. The property 2 is trivial to show. The property 3 might be somewhat easier to see given the following expression for $x {\text{swz}}$. $$ \begin{align} x {\text{swz}} &= x {\text{chunk swz}} × \text{sizeof} \text{TC} / \text{sizeof} \text{T} \% \text{NX} + x \% \text{sizeof} \text{TC} / \text{sizeof} \text{T} \\ &= \left y / \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} × \text{NX} × \text{sizeof} \text{T} / \text{sizeof} \text{TC} + \left y \% \text{NX} \right \oplus x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} \right × \text{sizeof} \text{TC} / \text{sizeof} \text{T} \% \text{NX} + x \% \text{sizeof} \text{TC} / \text{sizeof} \text{T} \\ &= \left y \% \text{NX} \right \oplus x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} × \text{sizeof} \text{TC} / \text{sizeof} \text{T} \% \text{NX} + x \% \text{sizeof} \text{TC} / \text{sizeof} \text{T} \\ &= \left y \% \text{NX} \right \oplus x × \text{sizeof} \text{T} / \text{sizeof} \text{TC} × \text{sizeof} \text{TC} / \text{sizeof} \text{T} + x \% \text{sizeof} \text{TC} / \text{sizeof} \text{T} \\ \end{align} $$ Given any $x$ and any $\{y, y+1, y+2, \cdots, y+\text{NX}-1\}$, the number of unique swizzled index $x {\text{swz}}$ is $\text{NX}$ which is maximized. Examples Matrix Transpose In this example, we implemented matrix transpose CUDA kernels using shared memory in three different ways: - Transpose with shared memory bank conflict. - Transpose without shared memory bank conflict via padding. - Transpose without shared memory bank conflict via swizzling. 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367 | include