I recently set to explore the efficency of mma instruction within a single precision GEMM kernel. I recall that using this warp-level instruction requires abhering to specific data layout rules in regesiter. This means each thread in a warp must contain a particular set of data to ensure the mma operation produces the correct computational result (Of course, the output data layout also must follow some rules). To investgate, I consulted the PTX documentation aiming to review the data layout specifications from NVIDA. There, I noted that the mma.m16n8k8.row.col instruction can be usd for tf32 data type, and the data layout for the left matrix ($A$) and the right matrix ($B$) detailed in the following figure:
Obviously, for the matrix A, each thread must hold 4 floats in its local memory. The 4 floats is not stored at contiguous addresses; instead, they are distributed across 4 sub-matrices. Consequently, the 32 threads in a warp collectively hold a single sub-matirx in a row-major order and this pattern is repeated 4 times😎. Similarly, the pattern for matrix B is repeated twice, so each thread must contain 2 floats. The key distinction is that the elements within each sub-matrix are arranged in column-major order.
As is well known, NVIDA provides a function for the data movement from shared memory to registers called ldmatirx. Readers can consult the office documentation for mare information. The primary limitation of this function is that it supports a maximum of 16-bit data, meaning the operational unit is 16 bits. This function is executed within a warp, and the “x1”, “x2”, and “x4” modifiers are used to repeat the execution. Each repetition builds upon the data layout of the “x1” operation. Providing the correct shared memory address for each thread is sufficient, as the instruction automatically gets the corresponding data defined by NVIDA. The data layout for “x1” is shown in the following figure:
It is clear that the data is stored in row-major in an 8*8 matrix for half precision. However, ldmatirx supports a maximum of 16-bit data movement while the register size on chip is 32 bits. This leads me to wonder if I can use this function for single precision data movement My reasoning is that the chip stores 2 half-precision values into a single 32-bit register in row-major. Therefore, when I use the “x1” for fp32 data, I would theoretically get an 8x4 matrix. In fact, It does work as I expected. I write a simple demo kernel to verify the data layout and the each thread gets the correct data.
So for the mma instruction with tf32 data type, I can use the ldmatrix with “x4” to load the left matrix A. But for the right matrix B, the approach differs. Since the data layout for matirx B requires a column-major in the shared memory, I have to instruct ldmatrix to load the data accordingly. But how ? Fortunately, this can be achieve using trans modifier. While using the trans, ldmatrix will work in a column-major way. So clearly, for the matrix B, I can use the ldmatirx with “x2” and trans.
But is the method truly work ?
No, it is not.
Since the ldmatrix is designed for 16-bit data movement, when I use it for non-trans 32-bit data layout, it does work. This is because the loading order is consistent with the data layout, which is row-major. Therefore, adjacent 16-bit data can be viewed as a single 32-bit unit, and the final column dimension of data layout is reduced by half. However, when I use trans, each thread will load the wrong 16-bit data and fially packs them into incorrect 32-bit registers. This leads to wrong computation results when using the mma instruction. Look at the following figure for more details:
Even though the ldmatrix get the correct 16-bit data, the packing into 32-bit register is wrong. Due to the maximum data movement unit is 16-bit, ldmatrix loads value in a column-major way using trans. This causes each thread to source half of its required data from other values, leading to incorrect results. For example, a thread (thread one shown in the figure) might combine the high 16 bits of one value with the high 16 bits of another, assembling a 32-bit number that is entirely wrong !
The solution is to modify the shared memory layout to column-major order, which allows the data to be loaded correctly without using the transpose parameter. To solve this issue, one effective method is to transpose the data layout when writing to shared memory from global memory. Or, I can manually map each thread to load the correct data from shared memory so that it can obey the data layout rules.
Now, we can answer the original question: when you use ldmatrix for 32-bit data, it only works correctly when the data layout is row-major. If the data layout is column-major, using the trans modifier will lead to incorrect results due to the 16-bit loading unit limitation. You must adjust the shared memory layout accordingly to ensure correct data loading.😊