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, # n_scaled_dim is stride for 1st dim of b_scale
    stride_c_m_dim, stride_c_n_dim,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    # Program IDs identify the C-tile this specific program instance will compute
    pid_m_abs = tl.program_id(0) # Absolute M-tile index
    pid_n_abs = tl.program_id(1) # Absolute N-tile index

    # Offsets for the C tile that this program instance will compute
    # 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)

    # Boundary masks for M and N dimensions
    # M, N are divisible by 64. BLOCK_SIZE_M/N are 128. Masking is needed if M/N is not multiple of 128.
    mask_m_dim = offs_m_for_c_tile < M    # Shape (BLOCK_SIZE_M)
    mask_n_dim = offs_n_for_c_tile < 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.
    # offs_m_for_c_tile needs to be masked for these loads.
    a_scale_load_ptrs_base = A_SCALE_ptr + offs_m_for_c_tile * stride_a_scale_m_dim
    # offs_m_for_c_tile[:, None] to allow broadcasting with K-dimension later for A tile
    a_tile_ptrs_base = A_ptr + offs_m_for_c_tile[:, None] * stride_a_m_dim

    # K-dimension loop setup
    k_indices_range = tl.arange(0, BLOCK_SIZE_K) # Shape (BLOCK_SIZE_K)
    # Loop over K dimension in blocks of BLOCK_SIZE_K
    # K is divisible by BLOCK_SIZE_K (128), so no boundary check for K-loop itself.
    for k_block_start in range(0, K, BLOCK_SIZE_K):
        # current_k_scale_idx: Index for the K-dimension of scale tensors.
        # Scale K-granularity is 128, BLOCK_SIZE_K is 128.
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K

        # Load A_SCALE for the current M-tile rows and K-block.
        # This is a vector of size BLOCK_SIZE_M. M-dimension is masked. K_scaled-dim is within bounds.
        a_scale_load_ptrs = a_scale_load_ptrs_base + current_k_scale_idx * stride_a_scale_k_scaled_dim
        # `other=1.0` for scales: if a_scale is 1.0 where a_fp8 is 0.0 (due to M-mask), product is 0.0, which is fine.
        a_scales_m_vector = tl.load(a_scale_load_ptrs, mask=mask_m_dim, other=1.0, eviction_policy="evict_first")

        # Load B_SCALE (scalar for this C tile and K-block).
        # b_scale N-dimension granularity is 128. BLOCK_SIZE_N is 128.
        # pid_n_abs directly indexes the N-scaled dimension of b_scale.
        # Grid for N ensures pid_n_abs is a valid index for b_scale.shape[0] (which is N_scaled).
        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 + # Stride for N_scaled dim
            current_k_scale_idx * stride_b_scale_k_scaled_dim    # Stride for K_scaled dim
        )
        # Scale indices (N_scaled, K_scaled) are assumed 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 # Shape (BLOCK_SIZE_K)

        # Load A tile (BLOCK_SIZE_M x BLOCK_SIZE_K). M-dimension masked. K-dimension no mask.
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim # (BM,BK)
        mask_a_data_load = mask_m_dim[:, None] # Shape (BM,1), broadcasts to (BM,BK). K-dim is always valid.
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_data_load, other=0.0)

        # Load B tile (transposed view for dot: BLOCK_SIZE_K x BLOCK_SIZE_N). K-dim no mask. N-dim masked.
        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_b_data_load = mask_n_dim[None, :] # Shape (1,BN), broadcasts to (BK,BN). K-dim is always valid.
        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 for FP8 inputs.
        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 and N-dimension are masked.
    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_c_store = mask_m_dim[:, None] & mask_n_dim[None, :] # Shape (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 
    
    # Input `a` is M x K in column-major order (from problem: a.T is K x M in generate_input)
    # A_ptr + k*M + m accesses A[m,k]
    stride_a_m_dim = 1 
    stride_a_k_dim = M

    # Input `b` is N x K in column-major order (from problem: b.T is K x N in generate_input)
    # B_ptr + k*N + n accesses B[n,k]
    stride_b_n_dim = 1
    stride_b_k_dim = N
    
    # Input `a_scale` is M x (K // 128) in column-major order (from problem: a_scale.T is K//128 x M)
    # A_SCALE_ptr + k_scaled*M + m accesses a_scale[m, k_scaled]
    stride_a_scale_m_dim = 1
    stride_a_scale_k_scaled_dim = M
    
    # Input `b_scale` is (N // 128) x (K // 128) in column-major order (from problem: b_scale.T is K//128 x N//128)
    # B_SCALE_ptr + k_scaled*(N_scaled_rows) + n_scaled accesses b_scale[n_scaled, k_scaled]
    # b_scale.shape[0] is N_scaled = N // 128 (or ceil(N/128) from ref generator's scale_n calculation)
    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 for K_scaled is num rows (N_scaled)

    # Output `c` is M x N in ROW-major order
    # C_ptr + m*N + n accesses C[m,n]
    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
    
    # Grid calculation
    # M, N are divisible by 64. They might not be divisible by BLOCK_SIZE_M/N (128).
    # So, cdiv is appropriate for grid calculation.
    grid_m = triton.cdiv(M, BLOCK_SIZE_M_val)
    grid_n = triton.cdiv(N, BLOCK_SIZE_N_val)
    grid = (grid_m, grid_n)
    
    # --- Tuning Parameters ---
    # Shared memory per stage for A and B tiles:
    # (128*128 for A_tile) + (128*128 for B_tile) = 16384 + 16384 = 32768 bytes.
    # For num_stages=2, total shared memory = 32768 * 2 = 65536 bytes.
    # This matches the MI300 hardware limit, maximizing pipelining under this constraint.
    num_warps_val = 8
    num_stages_val = 2

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