[QST] What can go wrong without cute::warpgroup_fence_operand(accum) in GEMM kernels?

[hyhieu]

I see warpgroup_fence_operand(accum) typically being included in GEMM kernels.

For the sake of discussion, let us consider this one:

cutlass/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp

Lines 424 to 445 in a8f2c80

	warpgroup_fence_operand(accum);
	// Prologue MMAs
	CUTLASS_PRAGMA_UNROLL
	for (int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
	prologue_mma_count > 0; --prologue_mma_count)
	{
	// WAIT on smem_pipe_read until it's data is available
	pipeline.consumer_wait(smem_pipe_read);
	warpgroup_arrive();
	// Unroll the K mode manually to set scale D to 1
	CUTLASS_PRAGMA_UNROLL
	for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
	// (V,M,K) x (V,N,K) => (V,M,N)
	cute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum);
	tiled_mma.accumulate_ = GMMA::ScaleOut::One;
	}
	warpgroup_commit_batch();
	++smem_pipe_read;
	--k_tile_count;
	}
	warpgroup_fence_operand(accum);

I noticed that locally, I could disable the first warpgroup_fence_operand(accum) (at line 424 above), and unit tests still pass.

Hence, I wonder what the fence is for?

I did a little bit of digging, and I found that the asm implementation of warpgroup_fence_operand seems to be:

cutlass/include/cute/arch/mma_sm90_gmma.hpp

Lines 92 to 98 in a8f2c80

	CUTE_HOST_DEVICE
	void
	warpgroup_fence_operand(float& reg) {
	#if defined(__CUDA_ARCH__)
	asm volatile("" : "+f"(reg) :: "memory");
	#endif
	}

So, unlike warpgroup_arrive, warpgroup_wait, and warpgroup_commit_batch, warpgroup_fence_operand does not correspond to a wgmma.xxx instruction. Wondering what that magical ASM is doing:

asm volatile("" : "+f"(reg) :: "memory");

Thank you!

[thakkarV]

this does not have any correctness implication for the kernel. it simply is an NVVM code motion fence to ensure the registers of the WGMMA instruction do not get touched by anything else in the middle of the WGMMA batch of instructions

[hyhieu]

this does not have any correctness implication for the kernel. it simply is an NVVM code motion fence to ensure the registers of the WGMMA instruction do not get touched by anything else in the middle of the WGMMA batch of instructions

A reply in 3 mins... Wow! Thank you, @thakkarV!

Follow up on your reply, can we come up with an example where something would go wrong without warpgroup_fence_operand?

For instance, (I am not sure if it is the case, but) does this CUTLASS_PRAGMA_UNROLL loop touch the accum registers multiple times, and hence require the fence?

   // Unroll the K mode manually to set scale D to 1 
   CUTLASS_PRAGMA_UNROLL 
   for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { 
     // (V,M,K) x (V,N,K) => (V,M,N) 
     cute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum); 
     tiled_mma.accumulate_ = GMMA::ScaleOut::One; 
   }

[thakkarV]

in this case, no, it does not touch the accumulator registers. accumulate_ is a predicate specifying whether the input accumulators should be scaled by zero (cleared) before accumulating the result into them or not. (D = A @ B as opposed to D += A @ B)