[QST] Using tile sizes in generator.py can't pass compilation?

[umiswing]

Hi! I'm trying to change the tile size in examples/36_gather_scatter_fusion to find the best performance. The A, B, C and D in my code is all fp32. According to #612 , I read generator.py and find GenerateSM80_TensorOp_1688(manifest, cuda_version). I use the tile size in its TileDescription list, but get compilation error like this:

/home/me/cutlass/include/cutlass/transform/pitch_linear_thread_map.h(295): error: static assertion failed with "Number of iterations must be non-zero"

The tile settings in my source code:

// This code section describes the tile size a thread block will compute
using ShapeMMAThreadBlock =
    cutlass::gemm::GemmShape<256, 128, 16>;  
// This code section describes tile size a warp will compute
using ShapeMMAWarp = cutlass::gemm::GemmShape<4, 2, 1>;  
// This code section describes the size of MMA op
using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 8>;  
...
// Number of pipelines you want to use
constexpr int NumStages = 3;

[jackkosaian]

Hi, @umiswing.

The [4, 2, 1] in the TileDescription fields you mention represents the number of warps in each dimension of the ThreadblockShape. Thus, the WarpShape for a ThreadblockShape = <256, 128, 16> with number of warps per dimension equal to [4, 2, 1] should be: <256/4, 128/2, 16/1> = <64, 64, 16>.

You can investigate the kernels generated by CUTLASS under build/tools/library/generated/gemm to see this. You can also see how this is calculated here.

Hope this helps.
- Jack

[umiswing]

@jackkosaian Thanks! It works now! :) And what determines the best tile size? Is this related to both the size of the problem and the hardware? If so, does this mean I should test the tile size one-by-one to find the best one? Are there any rules or experience I can refer to to set optimal tile size for different problem size?

[jackkosaian]

You're correct that the best tile size will depend both on the problem size and on the hardware. Your best bet is likely to be trying different combinations of ThreadblockShape, WarpShape, and Stages. This is what the CUTLASS profiler does when it generates many TileDescriptions, as you've seen.

- Jack

[umiswing]

@jackkosaian Thanks! Unfortunately I get an Internal Error when using a tile size from the TileDescriptions list.
The tile size and num stages are:

// This code section describes the tile size a thread block will compute
using ShapeMMAThreadBlock =
    cutlass::gemm::GemmShape<256,64,32>; 
// This code section describes tile size a warp will compute
using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 32>;  
// This code section describes the size of MMA op
using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 8>;  
...
// Number of pipelines you want to use
constexpr int NumStages = 4;

The origin code is examples/36_gather_scatter_fusion, I modify the problem size (m, n, k, indices_size) to (20500000 ,16, 4, 205000). A, B, C and D are all row major and type fp32.
I set the CUTLASS_DEBUG_TRACE_LEVEL to 1 and it reports:

/home/me/cutlass/include/cutlass/gemm/kernel/gemm_universal.h:225  GemmUniversal::Arguments::Arguments() - problem_size: cutlass::gemm::GemmCoord {205000,16,4}
/home/me/cutlass/include/cutlass/gemm/device/gemm_universal_base.h:157  GemmUniversalBase::get_workspace_size()
/home/me/cutlass/include/cutlass/gemm/device/gemm_universal_base.h:182    workspace_bytes: 0
/home/me/cutlass/include/cutlass/gemm/kernel/gemm_universal.h:386  GemmUniversal::can_implement()
/home/me/cutlass/include/cutlass/gemm/kernel/gemm_universal.h:456    returning kSuccess
/home/me/cutlass/include/cutlass/gemm/device/gemm_universal_base.h:283  GemmUniversalBase::initialize() - workspace 0, stream: null
/home/me/cutlass/include/cutlass/gemm/device/gemm_universal_base.h:157  GemmUniversalBase::get_workspace_size()
/home/me/cutlass/include/cutlass/gemm/device/gemm_universal_base.h:182    workspace_bytes: 0
/home/me/cutlass/include/cutlass/gemm/device/gemm_universal_base.h:288    workspace_bytes: 0
Got cutlass error: Error Internal at: 172

I compile with CUDA Compilation Architectures: 80, which is less than max compute of this tile size.

[jackkosaian]

Does it work for other settings of tile sizes and number of stages (in particular, when each are smaller than what you've currently set)?

Which GPU are you running on?

Depending on the specific GPU you are running on, the combination of threadblock tile size + number of stages you've described may require more shared memory than available on a single SM. You can calculate the amount of shared memory needed for a GEMM as follows: ((threadblockM * threadblockK * sizeof(ElementA)) + (threadblockN * threadblockK * sizeof(ElementB))) * numStages.

Here's an example of a case in which the configuration you've described requires too much shared memory: using the calculation above, your kernel requires ((256 * 32 * 4) + (64 * 32 * 4)) * 4 = 163840 bytes of shared memory. If you were to try to run this on an A6000 GPU, which has only 102400 B of shared memory per SM, you would be unable to launch the kernel.

[umiswing]

Sorry for replying late. I found the error is caused by the shared memory limitation. It works fine after I set a proper tiling size. :)