import torch
import triton
import triton.language as tl

# The Triton kernel for FP8 blockwise matrix multiplication
@triton.jit
def matmul_kernel(
    A_ptr, B_ptr, A_SCALE_ptr, B_SCALE_ptr, C_ptr,
    M, N, K, # K is the original full K dimension
    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_dot: tl.constexpr,     # K dimension processed by each tl.dot call
    SCALE_K_GRANULARITY: tl.constexpr, # K dimension for scale factor granularity (usually 128)
    SCALE_N_CONST: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m_for_c_tile = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_for_c_tile = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

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

    # Outer loop: iterates over K in chunks of SCALE_K_GRANULARITY (128)
    # This loop determines which scale factors to use.
    for k_scale_block_outer_start in range(0, K, SCALE_K_GRANULARITY):
        # Load scales for this 128-block of K once
        current_scale_k_block_idx = k_scale_block_outer_start // SCALE_K_GRANULARITY
        
        a_scale_load_ptrs = A_SCALE_ptr + (offs_m_for_c_tile * stride_a_scale_m_dim +
                                           current_scale_k_block_idx * stride_a_scale_k_scaled_dim)
        mask_a_scale_load = (offs_m_for_c_tile < M)
        # tl.static_print(f"Loading a_scales_vector for k_scale_block_idx: {current_scale_k_block_idx}")
        a_scales_vector = tl.load(a_scale_load_ptrs, mask=mask_a_scale_load, other=1.0) # Shape: (BLOCK_SIZE_M)

        n_scaled_indices_for_b_block = offs_n_for_c_tile // SCALE_N_CONST
        b_scale_load_ptrs = B_SCALE_ptr + (n_scaled_indices_for_b_block * stride_b_scale_n_scaled_dim +
                                           current_scale_k_block_idx * stride_b_scale_k_scaled_dim)
        N_s_ceil_for_mask = (N + SCALE_N_CONST - 1) // SCALE_N_CONST
        mask_b_scale_load = (n_scaled_indices_for_b_block < N_s_ceil_for_mask)
        # tl.static_print(f"Loading b_scales_vector for k_scale_block_idx: {current_scale_k_block_idx}")
        b_scales_vector = tl.load(b_scale_load_ptrs, mask=mask_b_scale_load, other=1.0) # Shape: (BLOCK_SIZE_N)

        # Inner loop: iterates within the current 128-K scale block, in chunks of BLOCK_SIZE_K_dot
        # k_dot_loop_end needs to be capped by K and also by the end of the current scale block
        k_dot_loop_end_for_this_scale_block = tl.minimum(k_scale_block_outer_start + SCALE_K_GRANULARITY, K)
        
        for k_dot_loop_inner_start in range(k_scale_block_outer_start, k_dot_loop_end_for_this_scale_block, BLOCK_SIZE_K_dot):
            # --- Load A_dot block ---
            a_k_indices_range_dot = tl.arange(0, BLOCK_SIZE_K_dot)
            a_k_indices_absolute_dot = k_dot_loop_inner_start + a_k_indices_range_dot
            # Mask for K dimension of data load, relative to full K
            k_mask_for_data_dot_block = a_k_indices_absolute_dot < K 

            a_dot_tile_ptrs = A_ptr + (offs_m_for_c_tile[:, None] * stride_a_m_dim +
                                   a_k_indices_absolute_dot[None, :] * stride_a_k_dim)
            final_a_dot_mask = (offs_m_for_c_tile[:, None] < M) & k_mask_for_data_dot_block[None, :]
            a_block_fp8_dot = tl.load(a_dot_tile_ptrs, mask=final_a_dot_mask, other=0.0) # Shape: (BLOCK_SIZE_M, BLOCK_SIZE_K_dot)

            # --- Load B_dot block (transposed) ---
            b_dot_tile_ptrs = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim +
                                   a_k_indices_absolute_dot[:, None] * stride_b_k_dim)
            final_b_dot_mask = k_mask_for_data_dot_block[:, None] & (offs_n_for_c_tile[None, :] < N)
            b_block_fp8_dot = tl.load(b_dot_tile_ptrs, mask=final_b_dot_mask, other=0.0) # Shape: (BLOCK_SIZE_K_dot, BLOCK_SIZE_N)
            
            # --- Dequantization (using the scales loaded in outer loop) ---
            a_block_f32_dot = a_block_fp8_dot.to(tl.float32)
            b_block_f32_dot = b_block_fp8_dot.to(tl.float32)
            
            a_dequant_dot = a_block_f32_dot * a_scales_vector[:, None] # Broadcast scale over BLOCK_SIZE_K_dot
            b_dequant_dot = b_block_f32_dot * b_scales_vector[None, :] # Broadcast scale over BLOCK_SIZE_K_dot
            
            accumulator += tl.dot(a_dequant_dot, b_dequant_dot, allow_tf32=False)

    # --- Store C block ---
    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 = (offs_m_for_c_tile[:, None] < M) & (offs_n_for_c_tile[None, :] < N)
    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 # Full K dimension
    
    assert K == K_check, f"K dimension mismatch: a.shape[1]={K_orig}, b.shape[1]={K_check}"
    assert a.dtype == torch.float8_e4m3fnuz
    assert b.dtype == torch.float8_e4m3fnuz
    assert a_scale.dtype == torch.float32
    assert b_scale.dtype == torch.float32
    assert c.dtype == torch.bfloat16
    assert c.shape == (M, N)
    
    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 0
    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_dot_val = 64        # K dimension for each tl.dot call
    SCALE_K_GRANULARITY_val = 128    # Scales are defined for blocks of 128 in K
    
    SCALE_N_CONST_val = 128          # For indexing b_scale

    grid = (triton.cdiv(M, BLOCK_SIZE_M_val), triton.cdiv(N, BLOCK_SIZE_N_val))
    
    num_warps_val = 4 
    num_stages_val = 3 # Potentially more stages due to smaller K_dot blocks

    matmul_kernel[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_dot=BLOCK_SIZE_K_dot_val,
        SCALE_K_GRANULARITY=SCALE_K_GRANULARITY_val,
        SCALE_N_CONST=SCALE_N_CONST_val,
        num_warps=num_warps_val,
        num_stages=num_stages_val
    )
    return c
python eval.py benchmark /tmp/tmpkcruvyed