Multi-stage pipeline: Why wait for two groups to complete at the beginning?

[masahi]

In

cutlass/include/cutlass/gemm/threadblock/mma_multistage.h

Line 433 in 12f4108

cutlass::arch::cp_async_wait<Base::kStages - 2>();

, we wait for two oldest stages of async copy groups to finish. Why do we need to wait for two, instead of one?

If my read of the innermost warp_mma_k loop is correct, we only consume one-stage worth of async copy result in the inner most loop before we commit another group and wait for the next oldest group to be ready at

cutlass/include/cutlass/gemm/threadblock/mma_multistage.h

Lines 555 to 563 in 12f4108

	copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A,
	group_start_iteration_B);
	// Inserts a memory fence between stages of cp.async instructions.
	cutlass::arch::cp_async_fence();
	// Waits until kStages-2 stages have committed.
	arch::cp_async_wait<Base::kStages - 2>();
	__syncthreads();

.

So I want to know why we keep two stages of copy results always ready in the innermost loop.

[masahi]

oh maybe I got it. At the last warp_mma_k iteration of a given stage, we want to start loading from the next stage copy results:

cutlass/include/cutlass/gemm/threadblock/mma_multistage.h

Lines 500 to 501 in 12f4108

	this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]);
	this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]);

so yeah, I can see why we need two stages of copy results always ready.

[hwu36]

cutlass::arch::cp_async_wait<Base::kStages - 2>(); means 2 stages are still on the fly and the rest are committed. Check this comment.

we only consume one-stage worth of async copy result in the inner most loop before we commit another group and wait for the next oldest group to be ready

Correct.

At the last warp_mma_k iteration of a given stage, we want to start loading from the next stage copy results:

Yes, we do this.

[masahi]

means 2 stages are still on the fly and the rest are committed

Isn't it the opposite? I thought it means 2 stages are completed and Base::kStages - 2 stages are still in flight. I'm confused with the meaning of "committed" in the comment:

https://github.com/NVIDIA/cutlass/blob/master/include/cutlass/arch/memory_sm80.h#L401

If it said "Blocks until all but N previous cp.async.commit_group operations have completed", that agrees with my understanding of the instruction, according to the PTX doc https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-wait-group

[hwu36]

sn't it the opposite? I thought it means 2 stages are completed and Base::kStages - 2 stages are still in flight.

You are correct. I misspoke.

[TigerYang414]

Is this multi-stage discussed here related to stages argument used to init tile description?
If it is, the stages here means the number of stages used to copy memory to shared memory?
But for op code of simt, I only see 2 stages used, is there any limit for this?

[hwu36]

Is this multi-stage discussed here related to stages argument used to init tile description?

yes

If it is, the stages here means the number of stages used to copy memory to shared memory?

yes

But for op code of simt, I only see 2 stages used, is there any limit for this?

sgemm on sm80 can use multiple stages.

[TigerYang414]

So what determines stages we can use for different sms? Is there any doc or best practice guidelines?

[hwu36]

multistage is only useful for >= sm80. you can profile to choose the best stage for you. you can start from those used by cutlass profiler.