import torch
import triton
import triton.language as tl

# The Triton kernel for FP8 blockwise matrix multiplication using FP8 tl.dot
@triton.jit
def matmul_kernel_fp8dot_v2( # Name kept, but logic updated from previous version
    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
    stride_b_n_dim, stride_b_k_dim,          # Strides for B
    stride_a_scale_m_dim, stride_a_scale_k_scaled_dim, # Strides for A_SCALE
    stride_b_scale_n_dim, stride_b_scale_k_scaled_dim, # Strides for B_SCALE (col major)
    stride_c_m_dim, stride_c_n_dim,          # Strides for C
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr     # K dimension processed by each tl.dot call, also scale granularity for K
):
    pid_m = tl.program_id(0) # Block index for M dimension
    pid_n = tl.program_id(1) # Block index for N dimension, also N-scaled index for b_scale

    # Offsets for the C tile this program is responsible for
    offs_m_for_c_tile = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_for_c_tile = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    # Accumulator for the C tile, in FP32
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # Outer loop: iterates over K in chunks of BLOCK_SIZE_K
    for k_block_start in range(0, K, BLOCK_SIZE_K):
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K
        
        # Load A_scales for the current M-tile and K-scale-block
        # a_scales_m_vector shape: (BLOCK_SIZE_M)
        a_scale_load_ptrs = A_SCALE_ptr + (offs_m_for_c_tile * stride_a_scale_m_dim +
                                           current_k_scale_idx * stride_a_scale_k_scaled_dim)
        mask_a_scale_load = (offs_m_for_c_tile < M)
        a_scales_m_vector = tl.load(a_scale_load_ptrs, mask=mask_a_scale_load, other=1.0)

        # Load B_scale for the current N-program-block (pid_n) and K-scale-block
        # Assumes b_scale is dimensioned N_scaled x K_scaled,
        # where N_scaled corresponds to N // BLOCK_SIZE_N (pid_n indexes this directly)
        b_scale_scalar_ptr = B_SCALE_ptr + (
            pid_n * stride_b_scale_n_dim + # pid_n is the N-scaled block index
            current_k_scale_idx * stride_b_scale_k_scaled_dim
        )
        # Given problem constraints (M,N div by 64, K div by 128) and block sizes (128),
        # pid_n and current_k_scale_idx should be valid indices into b_scale's dimensions.
        b_scale_val = tl.load(b_scale_scalar_ptr) # Load scalar for this (N_block, K_block)

        # K indices for the current dot product block
        k_indices_range = tl.arange(0, BLOCK_SIZE_K)
        k_indices_absolute = k_block_start + k_indices_range

        # Load A data tile: (BLOCK_SIZE_M, BLOCK_SIZE_K)
        a_tile_ptrs = A_ptr + (offs_m_for_c_tile[:, None] * stride_a_m_dim +
                               k_indices_absolute[None, :] * stride_a_k_dim)
        mask_a_data_load = (offs_m_for_c_tile[:, None] < M) & (k_indices_absolute[None, :] < K)
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_data_load, other=0.0)

        # Load B data tile (for B^T): (BLOCK_SIZE_K, BLOCK_SIZE_N)
        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_data_load = (k_indices_absolute[:, None] < K) & (offs_n_for_c_tile[None, :] < N)
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_b_data_load, other=0.0)
        
        # Perform FP8 dot product, accumulate in FP32
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_transposed, out_dtype=tl.float32, allow_tf32=False)
        
        # Apply scales: C_tile = (A_dot * B_dot^T) * A_scale_m_vec[:,None] * B_scale_scalar_val
        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 = (offs_m_for_c_tile[:, None] < M) & (offs_n_for_c_tile[None, :] < N)
    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 
    
    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)

    # Strides for column-major inputs
    stride_a_m_dim = 1 
    stride_a_k_dim = M
    stride_b_n_dim = 1
    stride_b_k_dim = N 
    
    # Strides for column-major a_scale (M x K_scaled)
    stride_a_scale_m_dim = 1
    stride_a_scale_k_scaled_dim = M
    
    # Strides for column-major b_scale (N_scaled x K_scaled)
    stride_b_scale_n_dim = 1 
    stride_b_scale_k_scaled_dim = b_scale.shape[0] if b_scale.nelement() > 0 else 0

    # Strides for row-major output C
    stride_c_m_dim = N 
    stride_c_n_dim = 1

    # Tuned kernel launch parameters to reduce shared memory
    BLOCK_SIZE_M_val = 128
    BLOCK_SIZE_N_val = 128  # Reduced from 256
    BLOCK_SIZE_K_val = 128  # K block for dot product, matches K-scale granularity
    
    grid = (triton.cdiv(M, BLOCK_SIZE_M_val), triton.cdiv(N, BLOCK_SIZE_N_val))
    
    num_warps_val = 8    # Kept from previous OOM attempt; can be reduced if other issues.
    num_stages_val = 2   # Reduced from 4 to manage shared memory.

    matmul_kernel_fp8dot_v2[grid]( # Use the updated kernel version
        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_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/tmppwcvkh57