import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernel_fp8dot_final(
    A_ptr, B_ptr, A_SCALE_ptr, B_SCALE_ptr, C_ptr,
    M, N, K,
    stride_a_m_dim, stride_a_k_dim,
    stride_b_n_dim, stride_b_k_dim,
    stride_a_scale_m_dim, stride_a_scale_k_scaled_dim,
    stride_b_scale_n_scaled_dim, stride_b_scale_k_scaled_dim,
    stride_c_m_dim, stride_c_n_dim,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, # Common K-dimension for the dot product
    GROUP_M_SIZE: tl.constexpr
):
    # Program IDs
    pid_m_local = tl.program_id(0)  # M-tile index within the group
    pid_n_abs = tl.program_id(1)    # Absolute N-tile index
    pid_group_m = tl.program_id(2)  # Group index along M dimension

    # Calculate total number of M-tiles and absolute M-tile index
    grid_m_total_tiles = tl.cdiv(M, BLOCK_SIZE_M)
    first_m_tile_in_group = pid_group_m * GROUP_M_SIZE
    
    # Number of M-tiles this specific group will handle (can be less than GROUP_M_SIZE for the last group)
    num_m_tiles_in_this_group = tl.minimum(grid_m_total_tiles - first_m_tile_in_group, GROUP_M_SIZE)

    # Early exit for threads in the grid that are padding for the last group
    if pid_m_local >= num_m_tiles_in_this_group:
        return

    pid_m_abs = first_m_tile_in_group + pid_m_local

    # --- Kernel logic mostly from v7 (128x128x128 base) ---
    offs_m_for_c_tile = pid_m_abs * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_for_c_tile = pid_n_abs * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    mask_m_dim = offs_m_for_c_tile < M
    mask_n_dim = offs_n_for_c_tile < N

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    a_scale_load_ptrs_base = A_SCALE_ptr + offs_m_for_c_tile * stride_a_scale_m_dim
    a_tile_ptrs_base = A_ptr + offs_m_for_c_tile[:, None] * stride_a_m_dim

    k_indices_range = tl.arange(0, BLOCK_SIZE_K)
    for k_block_start in range(0, K, BLOCK_SIZE_K):
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K

        a_scale_load_ptrs = a_scale_load_ptrs_base + current_k_scale_idx * stride_a_scale_k_scaled_dim
        a_scales_m_vector = tl.load(a_scale_load_ptrs, mask=mask_m_dim, other=1.0, eviction_policy="evict_first")

        b_scale_n_idx_for_tile = pid_n_abs
        b_scale_scalar_ptr = B_SCALE_ptr + (
            b_scale_n_idx_for_tile * stride_b_scale_n_scaled_dim +
            current_k_scale_idx * stride_b_scale_k_scaled_dim
        )
        b_scale_val = tl.load(b_scale_scalar_ptr, eviction_policy="evict_first")

        k_indices_absolute = k_block_start + k_indices_range

        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        mask_a_data_load = mask_m_dim[:, None]
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_data_load, other=0.0)

        b_tile_ptrs = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim +
                               k_indices_absolute[:, None] * stride_b_k_dim)
        mask_b_data_load = mask_n_dim[None, :]
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_b_data_load, other=0.0)
        
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_transposed, out_dtype=tl.float32, allow_tf32=False)
        
        scaled_dot_output = dot_output_f32 * a_scales_m_vector[:, None] * b_scale_val
        accumulator += scaled_dot_output

    c_tile_ptrs = C_ptr + (offs_m_for_c_tile[:, None] * stride_c_m_dim +
                           offs_n_for_c_tile[None, :] * stride_c_n_dim)
    mask_c_store = mask_m_dim[:, None] & mask_n_dim[None, :]
    tl.store(c_tile_ptrs, accumulator.to(tl.bfloat16), mask=mask_c_store)


def custom_kernel(kernel_input_data: tuple) -> torch.Tensor:
    a, b, a_scale, b_scale, c = kernel_input_data

    M, K_orig = a.shape 
    N, K_check = b.shape
    K = K_orig 
    
    stride_a_m_dim = 1 
    stride_a_k_dim = M
    stride_b_n_dim = 1 
    stride_b_k_dim = N 
    stride_a_scale_m_dim = 1
    stride_a_scale_k_scaled_dim = M
    stride_b_scale_n_scaled_dim = 1 
    stride_b_scale_k_scaled_dim = b_scale.shape[0] if b_scale.nelement() > 0 else 1
    stride_c_m_dim = N 
    stride_c_n_dim = 1

    # --- Kernel Launch Parameters ---
    BLOCK_SIZE_M_val = 128
    BLOCK_SIZE_N_val = 128
    BLOCK_SIZE_K_val = 128 # Common K-dimension for dot product
    
    GROUP_M_SIZE_val = 8 # Experiment with GROUP_M

    # Grid calculation for GROUP_M
    grid_m_total_tiles = triton.cdiv(M, BLOCK_SIZE_M_val)
    grid_n_total_tiles = triton.cdiv(N, BLOCK_SIZE_N_val)
    
    num_m_groups = triton.cdiv(grid_m_total_tiles, GROUP_M_SIZE_val)
    
    # Grid: (num_m_tiles_per_group, total_n_tiles, num_m_groups)
    # The first dimension of the grid (program_id(0)) will run up to GROUP_M_SIZE_val.
    # The kernel handles the case where the last group might have fewer tiles.
    grid = (GROUP_M_SIZE_val, grid_n_total_tiles, num_m_groups)
    
    # --- Tuning Parameters ---
    num_warps_val = 8
    num_stages_val = 2 # Max shared memory with 128x128x128 blocks

    matmul_kernel_fp8dot_final[grid](
        a, b, a_scale, b_scale, c,
        M, N, K, 
        stride_a_m_dim, stride_a_k_dim,
        stride_b_n_dim, stride_b_k_dim,
        stride_a_scale_m_dim, stride_a_scale_k_scaled_dim,
        stride_b_scale_n_scaled_dim, stride_b_scale_k_scaled_dim, 
        stride_c_m_dim, stride_c_n_dim,
        BLOCK_SIZE_M=BLOCK_SIZE_M_val,
        BLOCK_SIZE_N=BLOCK_SIZE_N_val,
        BLOCK_SIZE_K=BLOCK_SIZE_K_val,
        GROUP_M_SIZE=GROUP_M_SIZE_val,
        num_warps=num_warps_val,
        num_stages=num_stages_val
    )
    return c
python eval.py benchmark /tmp/tmp0k287nvj