import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernel_fp8dot_optimized(
    A_ptr, B_ptr, A_SCALE_ptr, B_SCALE_ptr, C_ptr,
    M, N, K,
    stride_a_m_dim, stride_a_k_dim,          # Strides for A (col-major)
    stride_b_n_dim, stride_b_k_dim,          # Strides for B (col-major)
    stride_a_scale_m_dim, stride_a_scale_k_scaled_dim, # Strides for A_SCALE (col-major view of M x K_scaled)
    stride_b_scale_n_scaled_dim, stride_b_scale_k_scaled_dim, # Strides for B_SCALE (col-major view of N_scaled x K_scaled)
    stride_c_m_dim, stride_c_n_dim,          # Strides for C (row-major)
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    B_SCALE_N_GRANULARITY: tl.constexpr  # Granularity for N-dim of b_scale (typically 128)
):
    # Program IDs identify the C-tile this specific program instance will compute
    pid_m_abs = tl.program_id(0) # Program ID along M dimension
    pid_n_abs = tl.program_id(1) # Program ID along N dimension

    # Offsets for the C tile:
    # offs_m_for_c_tile: [pid_m*BM, pid_m*BM + 1, ..., pid_m*BM + BM-1]
    # offs_n_for_c_tile: [pid_n*BN, pid_n*BN + 1, ..., pid_n*BN + BN-1]
    offs_m_for_c_tile = pid_m_abs * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) # Shape (BLOCK_SIZE_M)
    offs_n_for_c_tile = pid_n_abs * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) # Shape (BLOCK_SIZE_N)

    # Accumulator for the C tile, initialized to zeros. Operates in float32.
    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.
    # M is divisible by BLOCK_SIZE_M (64). Grid ensures pid_m_abs is valid.
    # So, (offs_m_for_c_tile < M) is always true. No M-dim boundary masking needed for these.
    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 # offs_m expanded for broadcasting with K-dim later

    # K-dimension loop setup
    k_indices_range = tl.arange(0, BLOCK_SIZE_K)
    # Loop over K dimension in blocks of BLOCK_SIZE_K
    for k_block_start in range(0, K, BLOCK_SIZE_K):
        # current_k_scale_idx: Index for the K-dimension of scale tensors.
        # Problem states K is divisible by 128. If BLOCK_SIZE_K=128, this is the scale granularity.
        # Indexing for scales along K will be valid.
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K # K-scale index for a_scale and b_scale
        
        # Load A_SCALE for the current M-tile rows and K-block.
        # This is a vector of size BLOCK_SIZE_M. No boundary checks needed for M or K_scaled dimensions.
        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")

        # Load B_SCALE (scalar for this C tile and K-block).
        # b_scale_n_idx_for_tile: Index for the N-dimension of B_SCALE.
        # Based on problem spec and reference generator, b_scale dimensions are ceil(N/GRANULARITY), ceil(K/GRANULARITY).
        # The indexing pid_n_abs * BLOCK_SIZE_N should map to a valid N_scaled index.
        b_scale_n_idx_for_tile = (pid_n_abs * BLOCK_SIZE_N) // B_SCALE_N_GRANULARITY
        
        b_scale_scalar_ptr = B_SCALE_ptr + (
            b_scale_n_idx_for_tile * stride_b_scale_n_scaled_dim + # Stride for N_scaled dim
            current_k_scale_idx * stride_b_scale_k_scaled_dim    # Stride for K_scaled dim
        )
        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_block_start + k_indices_range # Shape (BLOCK_SIZE_K)

        # Load A tile (BLOCK_SIZE_M x BLOCK_SIZE_K)
        # No boundary check needed for M or K dimensions due to divisibility constraints
        # and BLOCK_SIZE_M, BLOCK_SIZE_K matching these.
        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, no other=0.0

        # Load B tile (BLOCK_SIZE_K x BLOCK_SIZE_N, effectively B.T for the dot product)
        # K-dimension boundary check not needed.
        # N-dimension boundary check IS required, as N may not be a multiple of BLOCK_SIZE_N (e.g. N=64, BN=128).
        b_tile_ptrs = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim +  # (1,BN) for N-offsets
                               k_indices_absolute[:, None] * stride_b_k_dim) # (BK,1) for K-offsets
        
        # Mask for N dimension of B
        mask_b_data_load = (offs_n_for_c_tile[None, :] < N) # Shape (1, BN), broadcasts to (BK, BN)
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_b_data_load, other=0.0)
        
        # Perform dot product: (BLOCK_SIZE_M x BLOCK_SIZE_K) @ (BLOCK_SIZE_K x BLOCK_SIZE_N)
        # Output is float32. allow_tf32=False is good practice for non-TF32 GEMMs.
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_transposed, out_dtype=tl.float32, allow_tf32=False)
        
        # Apply scales: C_sub_tile = (A_sub_tile @ B_sub_tile.T) * a_scale_vec * b_scale_scalar
        # a_scales_m_vector[:, None] broadcasts M-vector of scales across N columns of dot_output_f32
        scaled_dot_output = dot_output_f32 * a_scales_m_vector[:, None] * b_scale_val
        accumulator += scaled_dot_output

    # Store the computed C tile to global memory (M x N, row-major)
    # M-dimension boundary check not needed.
    # N-dimension boundary check IS required for C.
    c_tile_ptrs = C_ptr + (offs_m_for_c_tile[:, None] * stride_c_m_dim + # (BM,1) for M-offsets
                           offs_n_for_c_tile[None, :] * stride_c_n_dim) # (1,BN) for N-offsets
    
    # Mask for N dimension of C
    mask_c_store = (offs_n_for_c_tile[None, :] < N) # Shape (1, BN), broadcasts to (BM, BN)
    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 
    
    # Asserts for basic correctness (can be removed in final submission if overhead matters)
    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)
    assert M % 64 == 0 and N % 64 == 0 and K % 128 == 0 # From problem statement

    # Strides for column-major A (M x K) -> A[m, k] is at A_ptr + k*M + m
    stride_a_m_dim = 1 
    stride_a_k_dim = M
    
    # Strides for column-major B (N x K) -> B[n, k] is at B_ptr + k*N + n
    stride_b_n_dim = 1
    stride_b_k_dim = N 
    
    # Strides for A_SCALE
    # generate_input: a_scale = torch.randn([scale_k, m],...).T
    # So a_scale is (M, K_scaled) with strides (1, M) if contiguous after transpose.
    # K_scaled = ceil(K/128)
    stride_a_scale_m_dim = 1
    stride_a_scale_k_scaled_dim = M # a_scale.stride(1)
    
    # Strides for B_SCALE
    # generate_input: b_scale = torch.randn([scale_k, scale_n],...).T
    # So b_scale is (N_scaled, K_scaled) with strides (1, N_scaled) if contiguous.
    # N_scaled = ceil(N/128), K_scaled = ceil(K/128)
    stride_b_scale_n_scaled_dim = 1 # b_scale.stride(0)
    # b_scale.shape[0] is N_scaled. This is the stride for the K_scaled dimension.
    stride_b_scale_k_scaled_dim = b_scale.shape[0] if b_scale.nelement() > 0 else 0 

    # Strides for C (M x N, row-major) -> C[m, n] is at C_ptr + m*N + n
    stride_c_m_dim = N 
    stride_c_n_dim = 1

    # Kernel launch parameters
    BLOCK_SIZE_M_val = 64   # M is divisible by 64.
    BLOCK_SIZE_N_val = 128  # N is divisible by 64. N-masking handles BN > N or N % BN != 0.
    BLOCK_SIZE_K_val = 128  # K is divisible by 128. Scale granularity is 128.
    
    B_SCALE_N_GRANULARITY_val = 128 # b_scale N-dimension is per 128 elements of N.

    grid = (triton.cdiv(M, BLOCK_SIZE_M_val), triton.cdiv(N, BLOCK_SIZE_N_val))
    
    # --- Key Tuning Parameters ---
    num_warps_val = 8   # Typically 4 or 8 for such block sizes.
    num_stages_val = 3  # Increased from 2. Could try 4 if beneficial and fits LDS.
                        # 2 stages: (64*128 + 128*128) * 1B * 2 = 48KB
                        # 3 stages: (64*128 + 128*128) * 1B * 3 = 72KB
                        # 4 stages: (64*128 + 128*128) * 1B * 4 = 96KB
                        # All fit comfortably in MI300's 256KB LDS per CU.

    matmul_kernel_fp8dot_optimized[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/tmp_z933yqh