import torch
import triton
import triton.language as tl

# Triton kernel definition with autotuning
@triton.autotune(
    configs=[
        # Based on v20/v33/v35 (best so far: M=128, N=64, K_tile=128, 8 warps, 2 stages)
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 128, 'num_stages': 2}, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 128, 'num_stages': 1}, num_warps=8),
        
        # Based on v15 (another good one: M=128, N=128, K_tile=128, 8 warps, 2 stages)
        # This one is at the shared memory limit (65536 bytes for 2 stages)
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K_DATA_TILE': 128, 'num_stages': 2}, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K_DATA_TILE': 128, 'num_stages': 1}, num_warps=8),

        # Configurations with larger M and smaller K_DATA_TILE
        # M=256, N=64, K_tile=64, 2 stages -> shmem = 2*(256*64 + 64*64) = 2*(16384+4096) = 40960 (OK)
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 64,  'num_stages': 2}, num_warps=8),
        # M=256, N=64, K_tile=64, 3 stages -> shmem = 3*20480 = 61440 (OK)
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 64,  'num_stages': 3}, num_warps=8),

        # Trying 4 warps for some promising block sizes
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 128, 'num_stages': 2}, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 64,  'num_stages': 2}, num_warps=4),
        
        # A configuration with smaller K_DATA_TILE for M=128, N=64 block
        # M=128, N=64, K_tile=64, 2 stages -> shmem = 2*(128*64 + 64*64) = 2*(8192+4096) = 24576 (OK)
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K_DATA_TILE': 64,  'num_stages': 2}, num_warps=8),
    ],
    key=['M', 'N', 'K'], # Autotuning will be specific to these problem dimensions
)
@triton.jit
def _matmul_kernel_fp8dot_autotuned(
    A_ptr, B_ptr, A_SCALE_ptr, B_SCALE_ptr, C_ptr,
    M, N, K, # Matrix dimensions (used as autotune keys)
    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,
    # Fixed granularities based on problem spec, passed as constexpr
    B_SCALE_N_GRANULARITY: tl.constexpr, 
    SCALE_K_GRANULARITY: tl.constexpr,
    # Autotuned parameters (must be tl.constexpr in kernel signature)
    BLOCK_SIZE_M: tl.constexpr, 
    BLOCK_SIZE_N: tl.constexpr, 
    BLOCK_SIZE_K_DATA_TILE: tl.constexpr # K-dimension of a single data tile loaded
    # num_warps and num_stages are handled by autotuner from Config, not direct kernel args here
):
    """
    Autotuned Triton Kernel for FP8 Matrix Multiplication with Blockwise Scaling.
    Version 38.
    """
    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_data_tile = tl.arange(0, BLOCK_SIZE_K_DATA_TILE)

    # Loop over the K dimension in blocks of BLOCK_SIZE_K_DATA_TILE.
    for k_block_start in range(0, K, BLOCK_SIZE_K_DATA_TILE):
        # Index for the current K-block in the scale tensors (granularity is SCALE_K_GRANULARITY).
        current_k_scale_idx = k_block_start // SCALE_K_GRANULARITY 

        # --- Load A scales ---
        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, mask=mask_m_dim, other=1.0, eviction_policy="evict_first")

        # --- Load B scale ---
        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 + 
            current_k_scale_idx * stride_b_scale_k_scaled_dim    
        )
        b_scale_val = tl.load(b_scale_scalar_ptr, eviction_policy="evict_first")

        # Absolute K indices for the current data block.
        k_indices_absolute = k_block_start + k_indices_range_data_tile

        # --- Load A data tile ---
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        # Boundary check for K dimension within the tile load itself is implicitly handled if K is not a multiple of BLOCK_SIZE_K_DATA_TILE
        # However, the problem statement says K is divisible by 128, and K_DATA_TILE will be 64 or 128.
        # So, k_indices_absolute should stay within K for most of the loop.
        # The mask for K is more critical if K is not perfectly divisible by BLOCK_SIZE_K_DATA_TILE.
        # For simplicity and given constraints, we assume K is large enough / divisible.
        mask_a_k_tile = k_indices_absolute[None, :] < K 
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_m_dim[:, None] & mask_a_k_tile, other=0.0) # Default eviction policy

        # --- Load B data tile (transposed) ---
        b_tile_ptrs = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim + 
                               k_indices_absolute[:, None] * stride_b_k_dim)
        mask_b_k_tile = k_indices_absolute[:, None] < K
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_n_dim[None, :] & mask_b_k_tile, other=0.0) # Default eviction policy
        
        # --- Perform dot product and apply scaling ---
        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 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 

    # Strides (same as before)
    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
    
    # Fixed granularities for scale indexing
    B_SCALE_N_GRANULARITY_val = 128 
    SCALE_K_GRANULARITY_val = 128   

    # Grid is a lambda function that the autotuner will use with the chosen META config
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']),
        triton.cdiv(N, META['BLOCK_SIZE_N']),
    )

    # Call the autotuned kernel
    _matmul_kernel_fp8dot_autotuned[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, 
        B_SCALE_N_GRANULARITY_val, # Pass as a normal argument, kernel expects tl.constexpr
        SCALE_K_GRANULARITY_val,   # Pass as a normal argument, kernel expects tl.constexpr
        # BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K_DATA_TILE are set by autotuner
    )
    return c

python eval.py benchmark /tmp/tmp4mit5svz