9.7.15.5.1.2. Shared Memory Matrix Layout

If the argument imm-trans-a / imm-trans-b of the instruction wgmma.mma_async{.sp} is 0, then K-major is used for matrix A / B respectively. If the value of argument imm-trans-a is 1 then M-major is used for matrix A. If the value of the argument imm-trans-b is 1, then N-major is used for matrix B.

In a column-major default BLAS library such as cuBLAS, the matrices A and B with and without transpose can be classified as either K-Major or M-or-N-Major as shown in the following table:

To avoid confusion with A, B, row-major, col-major, transpose, and non-transpose, we will use MN-Major and K-Major throughout this section.

The matrices in the shared memory are made up of one or more "swizzle layout atom". The exact layout of these swizzle atoms depends on the swizzling mode, swizzle-atomicity, and the leading dimension. The layout of the swizzle are shown in Table 38.

Table 38 Various combinations of swizzling mode, leading dimension and swizzle-atom layout

The above shapes are for elements of size 128 bits. For smaller elements sizes, the same shapes would get multiplied along the leading dimension by a factor of 128/sizeof_bits(Element). For example, 128B MN major swizzle atom would have a shape of (8*(128/32))x8 = 32x8 for tf32 tensor core inputs.

Examples

The following are some example layouts of MxK or KxN matrices with various swizzling modes, and are in units of 128b elements as shown by each colored cell as shown in Figure 156, Figure 157, Figure 158, Figure 159, Figure 160, Figure 161, Figure 162, Figure 163.

!MN major 128B swizzling

Figure 156 MN major 128B swizzling

!K major 128B swizzling

Figure 157 K major 128B swizzling

!MN major 64B swizzling

Figure 158 MN major 64B swizzling

!K major 64B swizzling

Figure 159 K major 64B swizzling

!MN major 32B swizzling

Figure 160 MN major 32B swizzling

!K major 32B swizzling

Figure 161 K major 32B swizzling

!MN major interleaved

Figure 162 MN major interleaved

!K major interleaved

Figure 163 K major interleaved

Following are some of the examples of the 128B swizzling layout for tf32 element type.

K-Major: Figure 164

!K major

Figure 164 K major

MN-Major: Figure 165

!MN major

Figure 165 MN major

9.7.15.5.1.2.1. Major-ness supported by Strides

There are two strides involved while accessing a matrix from shared memory:

• Leading dimension byte offset
• Stride dimension byte offset

9.7.15.5.1.2.1.1. Leading Dimension Byte Offset

The leading dimension byte offset is defined differently for transposed and non-transposed matrices. The leading byte offset is defined as follows for matrices whose element types are normalized to 128-bits:

9.7.15.5.1.2.1.2. Stride Dimension Byte Offset

The stride dimension byte offset is defined differently for transposed and non-transposed matrices. The stride dimension byte offset is defined as follows for matrices whose element types are normalized to 128-bits:

9.7.15.5.1.2.1.3. Canonical Layouts

In terms of CuTe layouts the canonical layout can be expressed as follows:

where

• T = 128 / sizeof-elements-in-bits — T represents scale factor which normalizes matrix element types to 128-bits.
• m represents the number of repeating patterns across rows.
• k represents the number of repeating patterns across columns.

Examples

K-Major, no-swizzling and tf32 type: Figure 166

!K major, no-swizzling and tf32 type

Figure 166 K major, no-swizzling and tf32 type

the strides and related details are as follows:

Exact layout : Swizzle<0,4,3> o ((8,2),(4,4)):((4,32),(1,64))

Canonical Layout :Swizzle<0,4,3> o ((8,m),(T,2k)):((1T,SBO),(1,LBO))

K-Major, 32B swizzling and tf32 type: Figure 167

!K major, 32B swizzling and tf32 type

Figure 167 K major, 32B swizzling and tf32 type

the strides and related details are as follows:

Exact layout : Swizzle<1,4,3> o ((8,2),(4,4)):((8,64),(1,4))

Canonical Layout :Swizzle<1,4,3> o ((8,m),(T,2k)):((2T,SBO),(1,T))

MN-Major, no-swizzling and bf16 type: Figure 168

!MN major, no-swizzling and bf16 type

Figure 168 MN major, no-swizzling and bf16 type

the strides and related details are as follows:

Exact layout : Swizzle<0,4,3> o ((8,1,2),(8,2)):((1,8,64),(8,128))

Canonical Layout :Swizzle<0,4,3> o ((T,1,m),(8,k)):((1,T,SBO),(1T,LBO))

MN-Major, 32B swizzling and bf16 type: Figure 169

!MN major, 32B swizzling and bf16 type

Figure 169 MN major, 32B swizzling and bf16 type

the strides and related details are as follows:

Exact layout : Swizzle<1,4,3> o ((8,2,2),(8,2)):((1,8,128),(16,256))

Canonical Layout :Swizzle<1,4,3> o ((T,2,m),(8,k)):((1,T,LBO),(2T,SBO))

MN-Major, 64B swizzling and bf16 type: Figure 170

!MN major, 64B swizzling and bf16 type

Figure 170 MN major, 64B swizzling and bf16 type

the strides and related details are as follows:

Exact layout : Swizzle<2,4,3> o ((8,4,2),(8,2)):((1,8,256),(32,512))

Canonical Layout :Swizzle<2,4,3> o ((T,4,m),(8,k)):((1,T,LBO),(4T,SBO))