[BUG] mma_atom_call for mma.sync is eliminated if C is RZ #3689

import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack

WARP_SIZE = 32
A_VALUES_PER_LANE = 8
B_VALUES_PER_LANE = 4
C_VALUES_PER_LANE = 4

@cute.kernel
def mma_sync_kernel(
    gA: cute.Tensor,
    gB: cute.Tensor,
    gC: cute.Tensor,
    mma_atom: cute.MmaAtom,
):
    lane, _, _ = cute.arch.thread_idx()
    rA = cute.make_rmem_tensor((A_VALUES_PER_LANE,), cutlass.BFloat16)
    rB = cute.make_rmem_tensor((B_VALUES_PER_LANE,), cutlass.BFloat16)
    rC = cute.make_rmem_tensor((C_VALUES_PER_LANE,), cutlass.Float32)
    rD = cute.make_rmem_tensor((C_VALUES_PER_LANE,), cutlass.Float32)
    cute.autovec_copy(gA[lane, None], rA)
    cute.autovec_copy(gB[lane, None], rB)
    rC.fill(0.0)

    # IMPORTANT :The current warp-MMA lowering incorrectly eliminates that separate-D/C form.
    rD = rC

    cute.mma_atom_call(mma_atom, rD, rA, rB, rC)
    cute.autovec_copy(rD, gC[lane, None])


@cute.jit
def run_mma_sync(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
    op = cute.nvgpu.warp.MmaF16BF16Op(
        cutlass.BFloat16,
        cutlass.Float32,
        (16, 8, 16),
    )
    mma_atom = cute.make_mma_atom(op)

    mma_sync_kernel(gA, gB, gC, mma_atom).launch(
        grid=(1, 1, 1),
        block=(WARP_SIZE, 1, 1),
    )


def main() -> None:
    if not torch.cuda.is_available():
        raise RuntimeError("An NVIDIA GPU with mma.sync support is required")

    # These shapes describe native [lane, lane_value] fragments, not ordinary
    # [M, K], [K, N], and [M, N] matrix layouts.
    a = torch.ones((WARP_SIZE, A_VALUES_PER_LANE), device="cuda", dtype=torch.bfloat16)
    b = torch.ones((WARP_SIZE, B_VALUES_PER_LANE), device="cuda", dtype=torch.bfloat16)
    c = torch.empty((WARP_SIZE, C_VALUES_PER_LANE), device="cuda", dtype=torch.float32)

    run_mma_sync(
        from_dlpack(a, assumed_align=16),
        from_dlpack(b, assumed_align=16),
        from_dlpack(c, assumed_align=16),
    )

    # With A=1, B=1, and K=16, all 128 values in the native C fragment are 16.
    torch.testing.assert_close(c, torch.full_like(c, 16.0), rtol=0.0, atol=0.0)
    print("PASS: one BF16 m16n8k16 mma.sync, every native C-fragment value is 16")
    print(c)


if __name__ == "__main__":
    main()