import torch
import triton
import triton.language as tl

# The Triton kernel matmul_kernel_fp8dot_final remains the same as you provided.
# For brevity, I'm not repeating it here but it's assumed to be defined above this function.
@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,
    stride_c_m_dim, stride_c_n_dim,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    """
    Triton Kernel for FP8 Matrix Multiplication with Blockwise Scaling.

    This kernel computes C = A @ B.T where A and B are FP8 matrices and
    are scaled using FP32 scale factors.
    A: M x K (column-major, e4m3fnuz)
    B: N x K (column-major, e4m3fnuz) -> effectively K x N for dot product after loading
    A_SCALE: M x (K // 128) (column-major, fp32)
    B_SCALE: (N // 128) x (K // 128) (column-major, fp32)
    C: M x N (row-major, bf16)

    BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K: Tile dimensions for computation.
    It's assumed that BLOCK_SIZE_K is 128 due to scale factor granularity.
    It's assumed that BLOCK_SIZE_N is 128 due to b_scale indexing.
    """
    # Program IDs identify the current tile being processed
    pid_m_abs = tl.program_id(0)  # Absolute program ID for M dimension
    pid_n_abs = tl.program_id(1)  # Absolute program ID for N dimension

    # Offsets for the C tile this program instance will compute
    # arange creates a sequence like [0, 1, ..., BLOCK_SIZE - 1]
    offs_m_for_c_tile = pid_m_abs * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) # Rows in C tile
    offs_n_for_c_tile = pid_n_abs * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) # Cols in C tile

    # Boundary masks for M and N dimensions to handle non-perfectly tiled matrices
    mask_m_dim = offs_m_for_c_tile < M    # Mask for valid rows
    mask_n_dim = offs_n_for_c_tile < N    # Mask for valid columns

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

    # Base pointers for loading A scales. Each thread in M-dim of block loads its scale.
    # offs_m_for_c_tile has shape (BLOCK_SIZE_M,), so a_scale_load_ptrs_base also has this shape.
    a_scale_load_ptrs_base = A_SCALE_ptr + offs_m_for_c_tile * stride_a_scale_m_dim

    # Base pointers for loading A tiles.
    # offs_m_for_c_tile[:, None] makes it a column vector for broadcasting with k_indices.
    a_tile_ptrs_base = A_ptr + offs_m_for_c_tile[:, None] * stride_a_m_dim

    # Range of K indices for a K-block
    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):

        # Index for the current K-block in the scale tensors
        # Assumes scale block size in K is BLOCK_SIZE_K (typically 128)
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K # Integer division

        # --- Load A scales for the current M-tile and K-block ---
        # Pointers to load a vector of scales for the M-dimension of the C tile
        a_scale_load_ptrs = a_scale_load_ptrs_base + current_k_scale_idx * stride_a_scale_k_scaled_dim
        # Load a_scales for the current M-tile rows, for the current K-block. Shape: (BLOCK_SIZE_M,)
        # `other=1.0` provides a default value if mask is false (though ideally should not be needed if scales cover full M)
        a_scales_m_vector = tl.load(a_scale_load_ptrs, mask=mask_m_dim, other=1.0, eviction_policy="evict_first")

        # --- Load B scale for the current N-tile and K-block ---
        # `pid_n_abs` directly indexes the N-block dimension of b_scale.
        # This assumes BLOCK_SIZE_N corresponds to one block of b_scale's N-dimension (i.e., BLOCK_SIZE_N = 128).
        b_scale_n_idx_for_tile = pid_n_abs
        
        # Pointer to the scalar b_scale value
        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
        )
        # Load the scalar b_scale value.
        b_scale_val = tl.load(b_scale_scalar_ptr, eviction_policy="evict_first") # Assuming b_scale always exists for valid N,K blocks

        # Absolute K indices for the current block
        k_indices_absolute = k_block_start + k_indices_range # Shape: (BLOCK_SIZE_K,)

        # --- Load A tile (BLOCK_SIZE_M x BLOCK_SIZE_K) ---
        # Pointers for the A tile. Shape: (BLOCK_SIZE_M, BLOCK_SIZE_K)
        # a_tile_ptrs_base is (BLOCK_SIZE_M, 1)
        # k_indices_absolute[None, :] is (1, BLOCK_SIZE_K)
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        # Mask for loading A data. mask_m_dim[:, None] broadcasts over K.
        mask_a_data_load = mask_m_dim[:, None] # & (k_indices_absolute[None, :] < K) implicitly handled by loop range
        # Load A block. `other=0.0` for padding if K is not perfectly divisible by BLOCK_SIZE_K (though problem says K is div by 128).
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_data_load, other=0.0) # Shape: (BLOCK_SIZE_M, BLOCK_SIZE_K)

        # --- Load B tile (transposed, effectively BLOCK_SIZE_K x BLOCK_SIZE_N) ---
        # B is N x K (col-major). We load it as if it's K x N for the dot product.
        # offs_n_for_c_tile[None, :] gives shape (1, BLOCK_SIZE_N) - these are N indices (rows in B)
        # k_indices_absolute[:, None] gives shape (BLOCK_SIZE_K, 1) - these are K indices (cols in B)
        b_tile_ptrs = B_ptr + (offs_n_for_c_tile[None, :] * stride_b_n_dim +
                               k_indices_absolute[:, None] * stride_b_k_dim)
        # Mask for loading B data. mask_n_dim[None, :] broadcasts over K.
        mask_b_data_load = mask_n_dim[None, :] # & (k_indices_absolute[:, None] < K) implicitly handled
        # Load B block (transposed relative to how A is loaded for A@B.T thinking)
        # This loads a KxN shaped block from B.
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_b_data_load, other=0.0) # Shape: (BLOCK_SIZE_K, BLOCK_SIZE_N)

        # --- Perform dot product and scale ---
        # Compute dot product: (M_BLOCK x K_BLOCK) @ (K_BLOCK x N_BLOCK) -> (M_BLOCK x N_BLOCK)
        # Output is FP32 for accumulation. `allow_tf32=False` is important for true FP8/BF16 behavior.
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_transposed, out_dtype=tl.float32, allow_tf32=False)
        
        # Apply scaling factors.
        # a_scales_m_vector is (BLOCK_SIZE_M,), broadcast to (BLOCK_SIZE_M, 1)
        # b_scale_val is a scalar
        # Result is (BLOCK_SIZE_M, BLOCK_SIZE_N)
        scaled_dot_output = dot_output_f32 * a_scales_m_vector[:, None] * b_scale_val
        accumulator += scaled_dot_output

    # --- Store the accumulated C tile ---
    # Pointers for storing the C tile. C is row-major.
    # offs_m_for_c_tile[:, None] is (BLOCK_SIZE_M, 1)
    # offs_n_for_c_tile[None, :] is (1, BLOCK_SIZE_N)
    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 for storing C tile, ensuring writes are within bounds of C matrix.
    mask_c_store = mask_m_dim[:, None] & mask_n_dim[None, :]
    # Convert accumulator to bfloat16 and store.
    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_check should be same as K_orig
    K = K_orig # K dimension

    # Strides for input tensor a (M x K, column-major)
    # a[row, col] = ptr + col * M + row
    stride_a_m_dim = 1  # Stride for M dimension (rows)
    stride_a_k_dim = M  # Stride for K dimension (columns, leading dimension size)

    # Strides for input tensor b (N x K, column-major)
    # b[row, col] = ptr + col * N + row
    stride_b_n_dim = 1  # Stride for N dimension (rows)
    stride_b_k_dim = N  # Stride for K dimension (columns, leading dimension size)

    # Strides for a_scale tensor (M x (K // 128), column-major)
    # a_scale[row, col_scaled] = ptr + col_scaled * M + row
    stride_a_scale_m_dim = 1           # Stride for M dimension
    stride_a_scale_k_scaled_dim = M    # Stride for K_scaled dimension

    # Strides for b_scale tensor ( (N // 128) x (K // 128), column-major)
    # b_scale[row_scaled, col_scaled] = ptr + col_scaled * (N//128) + row_scaled
    # b_scale.shape[0] is (N // 128)
    stride_b_scale_n_scaled_dim = 1
    # Handle cases where b_scale might be empty if N or K are too small, though problem constraints suggest N, K are large enough.
    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: # If b_scale is 1D (e.g. K//128 = 1)
        stride_b_scale_k_scaled_dim = b_scale.shape[0] # This is N//128
    if b_scale.nelement() == 0 : # if N < 128 or K < 128
         stride_b_scale_k_scaled_dim = 1 # Avoid division by zero if shape[0] is 0


    # Strides for output tensor c (M x N, row-major)
    # c[row, col] = ptr + row * N + col
    stride_c_m_dim = N  # Stride for M dimension (rows)
    stride_c_n_dim = 1  # Stride for N dimension (columns)

    # Tile sizes for the kernel
    BLOCK_SIZE_M_val = 128
    BLOCK_SIZE_N_val = 128 # Must be 128 for current b_scale logic in kernel
    BLOCK_SIZE_K_val = 128 # Must be 128 for scale granularity

    # Grid dimensions for launching the kernel
    # triton.cdiv(X, Y) computes ceil(X / Y)
    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: num_warps and num_stages
    # Original: num_warps_val = 8, num_stages_val = 2
    # Tuned values:
    num_warps_val = 4   # Try reducing warps
    num_stages_val = 3  # Try increasing stages for better pipelining

    matmul_kernel_fp8dot_final[grid](
        a, b, a_scale, b_scale, c, # Data pointers
        M, N, K,                    # Matrix dimensions
        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, # Pass tile sizes as constexpr
        BLOCK_SIZE_N=BLOCK_SIZE_N_val,
        BLOCK_SIZE_K=BLOCK_SIZE_K_val,
        num_warps=num_warps_val,       # Number of warps per thread block
        num_stages=num_stages_val      # Number of stages for software pipelining
    )
    return c

python eval.py benchmark /tmp/tmpt7p_osm1