submission 77598

"""
v147c: Fix MLIR Context Issue

v147b error: "MLIR function requires a Context"
Solution: Create tiled_mma inside a function, not at module level

Try creating it inside the kernel where MLIR context exists.
"""

import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
from cutlass.cute.nvgpu import tcgen05
from cutlass.cute.nvgpu.tcgen05 import OperandSource, CtaGroup
import cutlass.utils.blockscaled_layout as blockscaled_utils

# Data types
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
acc_dtype = cutlass.Float32
c_dtype = cutlass.Float16
sf_vec_size = 16

# Configuration
mma_tiler_mnk = (128, 1, 256)
threads_per_cta = 128

# MMA descriptor (this part is safe at module level)
mma_op = tcgen05.MmaMXF4NVF4Op(
    sf_dtype=sf_dtype,
    instruction_shape=(128, 16, 64),
    cta_group=CtaGroup.ONE,
    a_src=OperandSource.SMEM,
)


@cute.kernel
def tcgen05_kernel(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
):
    """
    v147c: Create tiled_mma inside kernel where context exists
    """
    # Create mma_atom and tiled_mma HERE (inside kernel with context)
    mma_atom = cute.make_mma_atom(mma_op)
    tiled_mma = cute.make_tiled_mma(mma_atom)

    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, _, _ = cute.arch.thread_idx()

    warp_id = tidx // 32
    lane_id = tidx % 32

    # Tile the input tensors
    gA_mkl = cute.local_tile(
        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    gSFA_mkl = cute.local_tile(
        mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    gB_nkl = cute.local_tile(
        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    gSFB_nkl = cute.local_tile(
        mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    gC_mnl = cute.local_tile(
        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
    )

    m_idx = warp_id * 32 + lane_id
    valid_idx = m_idx if m_idx < 128 else 0

    # Output accumulator
    tCgC = gC_mnl[valid_idx, None, bidx, bidy, bidz]
    tCgC = cute.make_tensor(tCgC.iterator, 1)
    res = cute.zeros_like(tCgC, acc_dtype)

    # K-loop
    k_tile_cnt = gA_mkl.layout[3].shape

    for k_tile in range(k_tile_cnt):
        # Load current K-tile
        tAgA = gA_mkl[valid_idx, None, bidx, k_tile, bidz]
        tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
        tAgSFA = gSFA_mkl[valid_idx, None, bidx, k_tile, bidz]
        tBgSFB = gSFB_nkl[0, None, bidy, k_tile, bidz]

        # Load and convert (v146m pattern)
        a_val = tAgA.load().to(acc_dtype)
        b_val = tBgB.load().to(acc_dtype)
        sfa_val = tAgSFA.load().to(acc_dtype)
        sfb_val = tBgSFB.load().to(acc_dtype)

        # Store to register memory
        tArA = cute.make_rmem_tensor_like(tAgA, acc_dtype)
        tBrB = cute.make_rmem_tensor_like(tBgB, acc_dtype)
        tArSFA = cute.make_rmem_tensor_like(tAgSFA, acc_dtype)
        tBrSFB = cute.make_rmem_tensor_like(tBgSFB, acc_dtype)

        tArA.store(a_val)
        tBrB.store(b_val)
        tArSFA.store(sfa_val)
        tBrSFB.store(sfb_val)

        # Manual SIMT (working pattern from v146m)
        num_blocks = mma_tiler_mnk[2] // 16

        for i in cutlass.range_constexpr(num_blocks):
            idx = i * 16
            scale = tArSFA[idx] * tBrSFB[idx]
            sum_ab = acc_dtype(0.0)

            # Unrolled sum
            sum_ab += tArA[idx + 0] * tBrB[idx + 0]
            sum_ab += tArA[idx + 1] * tBrB[idx + 1]
            sum_ab += tArA[idx + 2] * tBrB[idx + 2]
            sum_ab += tArA[idx + 3] * tBrB[idx + 3]
            sum_ab += tArA[idx + 4] * tBrB[idx + 4]
            sum_ab += tArA[idx + 5] * tBrB[idx + 5]
            sum_ab += tArA[idx + 6] * tBrB[idx + 6]
            sum_ab += tArA[idx + 7] * tBrB[idx + 7]
            sum_ab += tArA[idx + 8] * tBrB[idx + 8]
            sum_ab += tArA[idx + 9] * tBrB[idx + 9]
            sum_ab += tArA[idx + 10] * tBrB[idx + 10]
            sum_ab += tArA[idx + 11] * tBrB[idx + 11]
            sum_ab += tArA[idx + 12] * tBrB[idx + 12]
            sum_ab += tArA[idx + 13] * tBrB[idx + 13]
            sum_ab += tArA[idx + 14] * tBrB[idx + 14]
            sum_ab += tArA[idx + 15] * tBrB[idx + 15]

            res += sum_ab * scale

    # Store result
    if m_idx < 128:
        tCgC.store(res.to(c_dtype))


@cute.jit
def my_kernel(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: tuple,
):
    m, _, k, l = problem_size

    a_tensor = cute.make_tensor(
        a_ptr,
        cute.make_layout(
            (m, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
        ),
    )

    n_padded_128 = 128
    b_tensor = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (n_padded_128, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * k, 32)),
        ),
    )

    c_tensor = cute.make_tensor(
        c_ptr,
        cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
    )

    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)

    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
    sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

    grid = (
        cute.ceil_div(c_tensor.shape[0], 128),
        1,
        c_tensor.shape[2],
    )

    tcgen05_kernel(
        a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor
    ).launch(
        grid=grid,
        block=[threads_per_cta, 1, 1],
        cluster=(1, 1, 1),
    )


_compiled_kernel_cache = None


def compile_kernel():
    global _compiled_kernel_cache

    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache

    a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
    sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)

    _compiled_kernel_cache = cute.compile(
        my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
    )

    return _compiled_kernel_cache


def custom_kernel(data):
    """
    v147c: tiled_mma created inside kernel (where context exists)
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    compiled_func = compile_kernel()

    m, k, l = a.shape
    k = k * 2
    n = 1

    a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(
        sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
    )
    sfb_ptr = make_ptr(
        sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
    )

    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))

    return c


if __name__ == "__main__":
    print("v147c: tiled_mma created inside kernel to fix MLIR context issue")