import torch
import triton
import triton.language as tl

# Triton kernel definition
@triton.jit
def matmul_kernel_fp8dot_final(
    A_ptr, B_ptr, A_SCALE_ptr, B_SCALE_ptr, C_ptr,
    M, N, K, # K is the total K dimension of the matrix
    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_TILE: tl.constexpr, # K-dimension of a single data tile load, e.g., 128
    B_SCALE_N_GRANULARITY: tl.constexpr,
    SCALE_K_GRANULARITY: tl.constexpr # Granularity of scales in K-dim, e.g., 128
):
    """
    Triton Kernel for FP8 Matrix Multiplication with Blockwise Scaling.
    Version 34 (K-Loop Unrolled by 2x):
    - BLOCK_SIZE_M=128, BLOCK_SIZE_N=64, BLOCK_SIZE_K_TILE=128
    - num_warps = 8, num_stages = 2
    - Scale loads eviction_policy: "evict_first".
    - Data (A, B tiles) loads eviction_policy: Triton default.
    - B_scale N-indexing uses B_SCALE_N_GRANULARITY.
    - Main K-loop unrolled once.

    A: M x K (column-major, e4m3fnuz)
    B: N x K (column-major, e4m3fnuz)
    A_SCALE: M x (K // SCALE_K_GRANULARITY) (column-major, fp32)
    B_SCALE: (N // B_SCALE_N_GRANULARITY) x (K // SCALE_K_GRANULARITY) (column-major, fp32)
    C: M x N (row-major, bf16)
    """
    pid_m_abs = tl.program_id(0) 
    pid_n_abs = tl.program_id(1) 

    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_tile = tl.arange(0, BLOCK_SIZE_K_TILE)

    # Loop over the K dimension, processing 2 * BLOCK_SIZE_K_TILE per iteration.
    for k_outer_loop_start in range(0, K, 2 * BLOCK_SIZE_K_TILE):
        
        # --- Process first K-block (k0) ---
        k_block_start_0 = k_outer_loop_start
        current_k_scale_idx_0 = k_block_start_0 // SCALE_K_GRANULARITY

        # Load A scales for k0
        a_scale_load_ptrs_0 = a_scale_load_ptrs_base + current_k_scale_idx_0 * stride_a_scale_k_scaled_dim
        a_scales_m_vector_0 = tl.load(a_scale_load_ptrs_0, mask=mask_m_dim, other=1.0, eviction_policy="evict_first")

        # Load B scale for k0
        b_scale_n_idx_for_tile_0 = (pid_n_abs * BLOCK_SIZE_N) // B_SCALE_N_GRANULARITY
        b_scale_scalar_ptr_0 = B_SCALE_ptr + (
            b_scale_n_idx_for_tile_0 * stride_b_scale_n_scaled_dim + 
            current_k_scale_idx_0 * stride_b_scale_k_scaled_dim    
        )
        b_scale_val_0 = tl.load(b_scale_scalar_ptr_0, eviction_policy="evict_first")

        # Load A data tile for k0
        k_indices_absolute_0 = k_block_start_0 + k_indices_range_tile
        a_tile_ptrs_0 = a_tile_ptrs_base + k_indices_absolute_0[None, :] * stride_a_k_dim
        a_block_fp8_0 = tl.load(a_tile_ptrs_0, mask=mask_m_dim[:, None], other=0.0)

        # Load B data tile for k0
        b_tile_ptrs_0 = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim + 
                                 k_indices_absolute_0[:, None] * stride_b_k_dim)
        b_block_fp8_transposed_0 = tl.load(b_tile_ptrs_0, mask=mask_n_dim[None, :], other=0.0)
        
        # Dot product and scale for k0
        dot_output_f32_0 = tl.dot(a_block_fp8_0, b_block_fp8_transposed_0, out_dtype=tl.float32, allow_tf32=False)
        scaled_dot_output_0 = dot_output_f32_0 * a_scales_m_vector_0[:, None] * b_scale_val_0
        accumulator += scaled_dot_output_0

        # --- Process second K-block (k1) ---
        k_block_start_1 = k_outer_loop_start + BLOCK_SIZE_K_TILE
        if k_block_start_1 < K: # Boundary check for the second block
            current_k_scale_idx_1 = k_block_start_1 // SCALE_K_GRANULARITY

            # Load A scales for k1
            a_scale_load_ptrs_1 = a_scale_load_ptrs_base + current_k_scale_idx_1 * stride_a_scale_k_scaled_dim
            a_scales_m_vector_1 = tl.load(a_scale_load_ptrs_1, mask=mask_m_dim, other=1.0, eviction_policy="evict_first")

            # Load B scale for k1
            # b_scale_n_idx_for_tile_1 is the same as for k0 because N-dim processing is independent of K-unrolling here
            b_scale_scalar_ptr_1 = B_SCALE_ptr + (
                b_scale_n_idx_for_tile_0 * stride_b_scale_n_scaled_dim + # Use same N-index
                current_k_scale_idx_1 * stride_b_scale_k_scaled_dim    
            )
            b_scale_val_1 = tl.load(b_scale_scalar_ptr_1, eviction_policy="evict_first")

            # Load A data tile for k1
            k_indices_absolute_1 = k_block_start_1 + k_indices_range_tile
            a_tile_ptrs_1 = a_tile_ptrs_base + k_indices_absolute_1[None, :] * stride_a_k_dim
            a_block_fp8_1 = tl.load(a_tile_ptrs_1, mask=mask_m_dim[:, None], other=0.0)

            # Load B data tile for k1
            b_tile_ptrs_1 = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim + 
                                     k_indices_absolute_1[:, None] * stride_b_k_dim)
            b_block_fp8_transposed_1 = tl.load(b_tile_ptrs_1, mask=mask_n_dim[None, :], other=0.0)
            
            # Dot product and scale for k1
            dot_output_f32_1 = tl.dot(a_block_fp8_1, b_block_fp8_transposed_1, out_dtype=tl.float32, allow_tf32=False)
            scaled_dot_output_1 = dot_output_f32_1 * a_scales_m_vector_1[:, None] * b_scale_val_1
            accumulator += scaled_dot_output_1

    # Store the final C tile
    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 and b_scale.dim() > 1 else 1
    if b_scale.nelement() > 0 and b_scale.dim() == 1: 
        stride_b_scale_k_scaled_dim = b_scale.shape[0] 
    if b_scale.nelement() == 0 : 
         stride_b_scale_k_scaled_dim = 1
    stride_c_m_dim = N
    stride_c_n_dim = 1

    # Tile sizes for the kernel (Based on v20/v33)
    BLOCK_SIZE_M_val = 128
    BLOCK_SIZE_N_val = 64   
    BLOCK_SIZE_K_TILE_val = 128  # K-dimension of individual A/B tiles loaded in the unrolled loop
    
    B_SCALE_N_GRANULARITY_val = 128 
    SCALE_K_GRANULARITY_val = 128   # Granularity of scales in K-dim

    grid_m = triton.cdiv(M, BLOCK_SIZE_M_val)
    grid_n = triton.cdiv(N, BLOCK_SIZE_N_val) 
    grid = (grid_m, grid_n)

    # Kernel launch parameters (Based on v20/v33)
    # Shared memory for M=128, N_tile=64, K_tile=128, num_stages=2:
    # 2 * (128*128*1_A + 64*128*1_B) = 2 * (16384 + 8192) = 49152 bytes. (Limit 65536)
    # The unrolling happens within one stage's worth of shared memory for A/B tiles.
    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_TILE=BLOCK_SIZE_K_TILE_val, 
        B_SCALE_N_GRANULARITY=B_SCALE_N_GRANULARITY_val,
        SCALE_K_GRANULARITY=SCALE_K_GRANULARITY_val,
        num_warps=num_warps_val, 
        num_stages=num_stages_val
    )
    return c

python eval.py benchmark /tmp/tmpag4w2eor