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, # K is the common K-dimension for the matmul operation
    stride_a_m_dim, stride_a_k_dim,
    stride_b_n_dim, stride_b_k_dim, # Strides for the original B matrix (N x K, column-major)
    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, # These define the C tile (e.g., 128x32)
    BLOCK_SIZE_K_COMMON: tl.constexpr, # Common K-dim for A_tile and B_transpose_tile in dot product
    A_B_SCALE_K_GRANULARITY: tl.constexpr,
    B_SCALE_N_GRANULARITY: tl.constexpr
):
    pid_m_abs = tl.program_id(0) # M-tile index
    pid_n_abs = tl.program_id(1) # N-tile index (for C-tiles of size BLOCK_SIZE_N)

    # Offsets for the C tile (e.g., 128x32)
    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)

    # Boundary masks
    mask_m_for_ac = offs_m_for_c_tile < M    # Mask for A loads, A_scale loads, C stores
    mask_n_for_bc = offs_n_for_c_tile < N    # Mask for B loads, C stores

    # Accumulator for the C tile (e.g., 128x32)
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # Base pointers for A and A_scale
    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 # Shape (BM, 1) for broadcasting with K

    # K-dimension loop setup. K (problem K) is divisible by BLOCK_SIZE_K_COMMON (128).
    k_indices_common_tile = tl.arange(0, BLOCK_SIZE_K_COMMON)
    
    for k_block_start in range(0, K, BLOCK_SIZE_K_COMMON):
        current_k_scale_idx = k_block_start // A_B_SCALE_K_GRANULARITY # Scale K-granularity is 128

        # Load A_SCALE (vector of size BLOCK_SIZE_M=128). M-dim masked.
        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_for_ac, other=1.0, eviction_policy="evict_first")

        # B_SCALE indexing. B_SCALE_N_GRANULARITY is 128. BLOCK_SIZE_N for data C-tile is 32.
        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")

        k_indices_absolute = k_block_start + k_indices_common_tile # Shape (BLOCK_SIZE_K_COMMON=128)

        # Load A tile (BLOCK_SIZE_M x BLOCK_SIZE_K_COMMON = 128x128). M-dim masked. K-dim no mask.
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        mask_a_load = mask_m_for_ac[:, None] # K-dim is full, M-dim is masked
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_load, other=0.0) # Shape (128, 128)

        # Load B tile for tl.dot(A_tile (128,128), B_tile_transposed (128,32)).
        # B is N x K (col-major). We need to load K_COMMON rows from B, and BLOCK_SIZE_N columns from B.
        # Pointer structure for B: B_ptr + k_abs * Stride_K_B (which is N) + n_abs * Stride_N_B (which is 1)
        b_tile_ptrs = B_ptr + (
            k_indices_absolute[:, None] * stride_b_k_dim +      # Vary K along rows (stride is N)
            offs_n_for_c_tile[None, :] * stride_b_n_dim         # Vary N along columns (stride is 1)
        )
        # Mask for B: K-dim is full, N-dim needs masking (mask_n_for_bc)
        mask_b_load = mask_n_for_bc[None, :] # Shape (1, BN=32), broadcasts to (BK_COMMON, BN)
        b_block_fp8_for_dot = tl.load(b_tile_ptrs, mask=mask_b_load, other=0.0) # Shape (128, 32)
        
        # dot(A_tile[128,128], B_tile_transposed[128,32]) -> C_tile_accumulator[128,32]
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_for_dot, 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 C tile (128x32). M and N dims masked.
    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_for_ac[:, None] & mask_n_for_bc[None, :] # Mask for the 128x32 C-tile
    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_common = K_orig # K is the common dimension for the matmul
    
    # Strides based on input tensor definitions
    # a: M x K_common, col-major. Access A[m, k] with ptr + k*M + m
    stride_a_m_dim = 1 
    stride_a_k_dim = M
    # b: N x K_common, col-major. Access B[n, k] with ptr + k*N + n
    stride_b_n_dim = 1 
    stride_b_k_dim = N 
    
    # a_scale: M x (K_common // 128), col-major
    stride_a_scale_m_dim = 1
    stride_a_scale_k_scaled_dim = M
    
    # b_scale: (N // 128) x (K_common // 128), col-major
    stride_b_scale_n_scaled_dim = 1 
    stride_b_scale_k_scaled_dim = b_scale.shape[0] if b_scale.nelement() > 0 else 1
    
    # c: M x N, row-major
    stride_c_m_dim = N 
    stride_c_n_dim = 1

    # --- Kernel Launch Parameters ---
    BLOCK_SIZE_M_val = 128          # C-tile M-dimension
    BLOCK_SIZE_N_val = 32           # C-tile N-dimension (small N)
    BLOCK_SIZE_K_COMMON_val = 128   # Common K-dimension for matmul A[BMxBKc] @ B_T[BKcxBN]
    
    A_B_SCALE_K_GRANULARITY_val = 128 # K granularity for scales
    B_SCALE_N_GRANULARITY_val = 128 # N granularity for b_scale

    # Grid calculation
    grid_m = triton.cdiv(M, BLOCK_SIZE_M_val)
    grid_n = triton.cdiv(N, BLOCK_SIZE_N_val) # N may not be div by BLOCK_SIZE_N_val (32)
    grid = (grid_m, grid_n)
    
    # --- Tuning Parameters ---
    # Shared memory per stage: A(128x128) + B_T(128x32) = 16384 + 4096 = 20480 bytes.
    # Max stages = floor(65536 / 20480) = 3.
    num_warps_val = 4 # For 128x32 C-tile.
    num_stages_val = 3 # Maximize pipelining. Total shared: 20480 * 3 = 61440 bytes.

    matmul_kernel_fp8dot_final[grid](
        a, b, a_scale, b_scale, c,
        M, N, K_common, 
        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_COMMON=BLOCK_SIZE_K_COMMON_val,
        A_B_SCALE_K_GRANULARITY=A_B_SCALE_K_GRANULARITY_val,
        B_SCALE_N_GRANULARITY=B_SCALE_N_GRANULARITY_val,
        num_warps=num_warps_val,
        num_stages=num_stages_val
    )
    return c
python eval.py benchmark /tmp/tmp6pkwpj2r