import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernel_fp8dot_no_masking(
    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,
    B_SCALE_N_GRANULARITY: tl.constexpr
):
    pid_m_abs = tl.program_id(0)
    pid_n_abs = tl.program_id(1)

    # Offsets for C tile. Since M, N are divisible by BLOCK_SIZE_M, BLOCK_SIZE_N,
    # these offsets are always within bounds of M and N respectively.
    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)

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

    # Base pointers for A and A_scale, adjusted for the M-dimension part of the tile.
    # (offs_m_for_c_tile < M) is effectively always true.
    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):
        # K is divisible by 128 (scale granularity for K).
        # current_k_scale_idx is an index into the K_scaled dimension.
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K # Assuming BLOCK_SIZE_K is the scale granularity (128)

        # Load A_SCALE: vector of size BLOCK_SIZE_M. No M or K_scaled boundary checks needed.
        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, eviction_policy="evict_first")

        # Calculate N-index for B_SCALE.
        # b_scale_n_tile_start_abs is the starting N-coordinate of the current C-tile.
        b_scale_n_tile_start_abs = pid_n_abs * BLOCK_SIZE_N
        b_scale_n_idx_for_tile = b_scale_n_tile_start_abs // B_SCALE_N_GRANULARITY
        
        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
        )
        # N_scaled and K_scaled indices for b_scale are implicitly valid due to grid and problem constraints.
        b_scale_val = tl.load(b_scale_scalar_ptr, eviction_policy="evict_first")

        # Absolute K indices for the current K-block.
        # (k_indices_absolute < K) is effectively always true as K is divisible by BLOCK_SIZE_K.
        k_indices_absolute = k_block_start + k_indices_range

        # Load A tile (BLOCK_SIZE_M x BLOCK_SIZE_K). No M or K boundary checks needed.
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        a_block_fp8 = tl.load(a_tile_ptrs) # No mask

        # Load B tile (transposed view for dot: BLOCK_SIZE_K x BLOCK_SIZE_N).
        # No K or N boundary checks needed.
        b_tile_ptrs = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim + 
                               k_indices_absolute[:, None] * stride_b_k_dim) 
        b_block_fp8_transposed = tl.load(b_tile_ptrs) # No mask
        
        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

    # Store the computed C tile. No M or N boundary checks needed for store.
    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)
    tl.store(c_tile_ptrs, accumulator.to(tl.bfloat16)) # No mask

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 

    # Strides (assuming inputs are already in specified column-major/row-major)
    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 # a_scale is M x K_scaled (col-major view)
    
    stride_b_scale_n_scaled_dim = 1 
    stride_b_scale_k_scaled_dim = b_scale.shape[0] # b_scale is N_scaled x K_scaled (col-major view), shape[0] is N_scaled

    stride_c_m_dim = N 
    stride_c_n_dim = 1

    # --- Kernel Parameters ---
    # Block sizes chosen to match divisibility constraints, enabling mask removal.
    BLOCK_SIZE_M_val = 64   # M is divisible by 64
    BLOCK_SIZE_N_val = 64   # N is divisible by 64
    BLOCK_SIZE_K_val = 128  # K is divisible by 128 (also scale K-granularity)
    
    # Granularity for N-dimension of b_scale (as per problem description/ref_gen)
    B_SCALE_N_GRANULARITY_val = 128

    # Grid calculation using exact division due to problem constraints and block size choices.
    # This ensures pid_m_abs * BLOCK_SIZE_M_val < M and pid_n_abs * BLOCK_SIZE_N_val < N.
    grid_m = M // BLOCK_SIZE_M_val
    grid_n = N // BLOCK_SIZE_N_val
    grid = (grid_m, grid_n)
    
    # Tuning parameters: number of warps and pipeline stages
    num_warps_val = 4
    num_stages_val = 4  # Shared mem: (64*128 + 64*128)*1B * 4 stages = 16KB * 4 = 64KB. Fits easily.

    matmul_kernel_fp8dot_no_masking[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,
        B_SCALE_N_GRANULARITY=B_SCALE_N_GRANULARITY_val,
        num_warps=num_warps_val,
        num_stages=num_stages_val
    )
    return c
python eval.py benchmark /tmp/tmph4n417ev