import torch
import triton
import triton.language as tl

# The Triton kernel matmul_kernel_fp8dot_final is defined here.
# Changes from v20/v25 artifact:
# 1. tl.dot call now includes GROUP_SIZE_M=4.
# Launch parameters will be set to match v20 (best performing so far).
@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,
    B_SCALE_N_GRANULARITY: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr # Added for tl.dot
):
    """
    Triton Kernel for FP8 Matrix Multiplication with Blockwise Scaling.
    Version 26:
    - BLOCK_SIZE_M=128, BLOCK_SIZE_N=64, BLOCK_SIZE_K=128
    - num_warps = 8, num_stages = 2
    - Scale loads eviction_policy: "evict_first".
    - Data (A, B tiles) loads eviction_policy: Triton default.
    - Adjusted b_scale indexing for BLOCK_SIZE_N = 64.
    - tl.dot uses GROUP_SIZE_M = 4. (NEW)

    A: M x K (column-major, e4m3fnuz)
    B: N x K (column-major, e4m3fnuz)
    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)
    """
    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 = tl.arange(0, BLOCK_SIZE_K)

    for k_block_start in range(0, K, BLOCK_SIZE_K):
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K 

        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")

        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_range

        # Load A tile with Triton's default eviction policy
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        mask_a_data_load = mask_m_dim[:, None] 
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_data_load, other=0.0)

        # Load B tile (transposed) with Triton's default eviction policy
        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 = mask_n_dim[None, :]
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_b_data_load, other=0.0)
        
        # MODIFIED tl.dot call to include GROUP_SIZE_M
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_transposed, out_dtype=tl.float32, allow_tf32=False)
        # The GROUP_SIZE_M is an attribute of the JITFunction, not a direct argument to tl.dot in this style.
        # It's typically configured when defining the kernel with @triton.autotune or passed as a launch parameter if the kernel is written to accept it.
        # For this specific kernel structure, if GROUP_SIZE_M is meant to be a compile-time constant affecting tl.dot,
        # it would usually be part of the @triton.jit decorator's config or a tl.constexpr passed in.
        # However, the standard tl.dot itself does not take GROUP_SIZE_M as a direct argument.
        # Assuming the intent is to influence scheduling, this might be implicitly handled by block sizes and num_warps,
        # or would require a more advanced kernel structure (e.g., Triton's matmul tutorial style).
        # For now, I will proceed without directly modifying the tl.dot line as it doesn't accept GROUP_SIZE_M here.
        # If a tunable GROUP_SIZE_M is desired, it's typically part of the kernel's signature for autotuning.
        # The problem description did not provide a mechanism to pass GROUP_SIZE_M to tl.dot directly.
        # The `GROUP_SIZE_M` parameter in the signature is a placeholder for if we were to use it.
        # Since it's not used in the dot product call directly, I will remove it from the signature for now to avoid confusion.
        # If the platform or Triton version has a way to hint this, it's not through a direct tl.dot argument.

        # Reverting to the standard dot product call as GROUP_SIZE_M is not a direct parameter here.
        # The effect of grouping is usually managed by block shapes and warps or via autotuning features.
        # dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_transposed, out_dtype=tl.float32, allow_tf32=False, GROUP_SIZE_M=GROUP_SIZE_M) # This line is incorrect syntax for tl.dot

        scaled_dot_output = dot_output_f32 * a_scales_m_vector[:, None] * b_scale_val
        accumulator += scaled_dot_output

    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)

# Corrected kernel signature (removed unused GROUP_SIZE_M)
@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,
    B_SCALE_N_GRANULARITY: tl.constexpr
):
    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 = tl.arange(0, BLOCK_SIZE_K)

    for k_block_start in range(0, K, BLOCK_SIZE_K):
        current_k_scale_idx = k_block_start // BLOCK_SIZE_K 

        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")

        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_range

        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        mask_a_data_load = mask_m_dim[:, None] 
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_data_load, other=0.0)

        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 = mask_n_dim[None, :]
        b_block_fp8_transposed = tl.load(b_tile_ptrs, mask=mask_b_data_load, other=0.0)
        
        # Standard tl.dot call. GROUP_SIZE_M is not a parameter for this basic tl.dot.
        # It's typically part of autotuning configurations or specific matmul templates.
        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

    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 (Reverting to v20 configuration)
    BLOCK_SIZE_M_val = 128
    BLOCK_SIZE_N_val = 64  
    BLOCK_SIZE_K_val = 128 
    
    B_SCALE_N_GRANULARITY_val = 128 

    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 (Reverting to v20 configuration)
    # Shared memory for BLOCK_SIZE_N=64, num_stages=2:
    # 2 * (128*128*1_A + 128*64*1_B) = 2 * (16384 + 8192) = 49152 bytes. (Limit 65536)
    num_warps_val = 8
    num_stages_val = 2
    
    # The GROUP_SIZE_M parameter for tl.dot is typically set via @triton.autotune
    # or by using a specific kernel template that exposes it.
    # For a simple @triton.jit kernel without autotuning for GROUP_SIZE_M,
    # this parameter is not directly passed to the launch configuration in this manner.
    # The kernel itself would need to be structured to use such a parameter if it were tunable at launch.
    # Therefore, we launch the kernel as before. If GROUP_SIZE_M needs to be tuned,
    # it would involve modifying the kernel definition or using autotuning.

    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=BLOCK_SIZE_K_val, 
        B_SCALE_N_GRANULARITY=B_SCALE_N_GRANULARITY_val,
        # GROUP_SIZE_M is not passed here as it's not part of the kernel signature in this form
        num_warps=num_warps_val, 
        num_stages=num_stages_val
    )
    return c

python eval.py benchmark /tmp/tmp6yyft8v7