import torch
import triton
import triton.language as tl

# The Triton kernel code (matmul_kernel_fp8dot_final) remains unchanged from your provided code.
# It is assumed to be defined above this function in the actual execution environment.
@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
):
    """
    Triton Kernel for FP8 Matrix Multiplication with Blockwise Scaling.
    This kernel computes C = (A * A_scale) @ (B * B_scale)^T.
    A is M x K (col-major), B is N x K (col-major).
    A_scale is M x K//128 (col-major), B_scale is N//128 x K//128 (col-major).
    C is M x N (row-major).
    """
    # Program ID determines the C-tile this program is responsible for
    pid_m_abs = tl.program_id(0)  # M-tile index
    pid_n_abs = tl.program_id(1)  # N-tile index

    # Offsets for the C tile this program will compute
    # e.g., if BLOCK_SIZE_M=128, BLOCK_SIZE_N=128:
    # offs_m_for_c_tile will be [0, 1, ..., 127] for the first block in M-dim
    # offs_n_for_c_tile will be [0, 1, ..., 127] for the first block in N-dim
    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 to prevent out-of-bounds access for tiles at the edges of M and N dimensions
    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, B_scale loads (indirectly), C stores

    # Accumulator for the C tile, initialized to zeros. Computation is done in float32.
    # Shape: (BLOCK_SIZE_M, BLOCK_SIZE_N), e.g., (128, 128)
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # Base pointers for A and A_scale, adjusted by the M-tile offset
    # offs_m_for_c_tile has shape (BLOCK_SIZE_M)
    # stride_a_scale_m_dim is 1 (A_scale is col-major, M is the contiguous dim for a fixed scale_k)
    a_scale_load_ptrs_base = A_SCALE_ptr + offs_m_for_c_tile * stride_a_scale_m_dim
    
    # stride_a_m_dim is 1 (A is col-major, M is the contiguous dim for a fixed K)
    # a_tile_ptrs_base has shape (BLOCK_SIZE_M, 1) to allow broadcasting with K-offsets later
    a_tile_ptrs_base = A_ptr + offs_m_for_c_tile[:, None] * stride_a_m_dim

    # K-dimension loop setup.
    # The problem states K is divisible by BLOCK_SIZE_K_COMMON (e.g., 128).
    k_indices_common_tile = tl.arange(0, BLOCK_SIZE_K_COMMON) # Shape (BLOCK_SIZE_K_COMMON), e.g., (128)

    # Loop over K dimension in blocks of BLOCK_SIZE_K_COMMON
    for k_block_start in range(0, K, BLOCK_SIZE_K_COMMON):
        # Determine the K-index for scale matrices (A_scale, B_scale)
        # A_B_SCALE_K_GRANULARITY is typically 128
        current_k_scale_idx = k_block_start // A_B_SCALE_K_GRANULARITY

        # --- Load A_SCALE ---
        # A_scale is M x (K // 128), col-major.
        # We need a vector of scales for the current M-tile rows and current K-scale-block.
        # stride_a_scale_k_scaled_dim is M (stride to get to the next K_scaled column in A_scale)
        a_scale_load_ptrs = a_scale_load_ptrs_base + current_k_scale_idx * stride_a_scale_k_scaled_dim
        # Load a vector of scales of size BLOCK_SIZE_M. Masked for M-dimension.
        # eviction_policy="evict_first" suggests this data might be reused less if L1 is pressured.
        a_scales_m_vector = tl.load(a_scale_load_ptrs, mask=mask_m_for_ac, other=1.0, eviction_policy="evict_first") # Shape (BLOCK_SIZE_M)

        # --- Load B_SCALE ---
        # B_scale is (N // 128) x (K // 128), col-major.
        # We need a scalar B_scale corresponding to the current N-tile and K-scale-block.
        # B_SCALE_N_GRANULARITY is typically 128.
        # pid_n_abs * BLOCK_SIZE_N gives the starting N-index of the C-tile.
        b_scale_n_idx_for_tile = (pid_n_abs * BLOCK_SIZE_N) // B_SCALE_N_GRANULARITY
        
        # stride_b_scale_n_scaled_dim is 1 (N_scaled is the contiguous dim for fixed K_scaled)
        # stride_b_scale_k_scaled_dim is (N // 128) (stride to get to next K_scaled col in B_scale)
        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") # Scalar

        # Absolute K indices for the current K-block
        # Shape (BLOCK_SIZE_K_COMMON), e.g., (128)
        k_indices_absolute = k_block_start + k_indices_common_tile

        # --- Load A tile ---
        # A is M x K, col-major. We load a tile of size (BLOCK_SIZE_M, BLOCK_SIZE_K_COMMON).
        # stride_a_k_dim is M (stride to get to the next K column in A)
        # k_indices_absolute[None, :] broadcasts K-offsets across M-dimension pointers.
        a_tile_ptrs = a_tile_ptrs_base + k_indices_absolute[None, :] * stride_a_k_dim
        # Mask for A load: M-dimension is masked (mask_m_for_ac), K-dimension is full (within block).
        mask_a_load = mask_m_for_ac[:, None] # Shape (BLOCK_SIZE_M, 1)
        # Load A tile. Shape (BLOCK_SIZE_M, BLOCK_SIZE_K_COMMON), e.g., (128, 128)
        a_block_fp8 = tl.load(a_tile_ptrs, mask=mask_a_load, other=0.0)

        # --- Load B tile ---
        # B is N x K, col-major. For tl.dot(A_tile, B_tile_transposed_effectively),
        # B_tile_transposed_effectively should be (BLOCK_SIZE_K_COMMON, BLOCK_SIZE_N).
        # We load B data that corresponds to B[k_abs, n_abs].
        # stride_b_k_dim is N (stride to get to the next K row in B, as it's col-major)
        # stride_b_n_dim is 1 (N is the contiguous dim for a fixed K in B)
        # k_indices_absolute[:, None] provides K-offsets for rows.
        # offs_n_for_c_tile[None, :] provides N-offsets for columns.
        b_tile_ptrs = B_ptr + (
            k_indices_absolute[:, None] * stride_b_k_dim +      # Vary K along rows of loaded tile (stride N for B)
            offs_n_for_c_tile[None, :] * stride_b_n_dim         # Vary N along columns of loaded tile (stride 1 for B)
        )
        # Mask for B load: K-dimension is full (within block), N-dimension is masked (mask_n_for_bc).
        mask_b_load = mask_n_for_bc[None, :] # Shape (1, BLOCK_SIZE_N)
        # Load B tile. Shape (BLOCK_SIZE_K_COMMON, BLOCK_SIZE_N), e.g., (128, 128) or (128, 32)
        b_block_fp8_for_dot = tl.load(b_tile_ptrs, mask=mask_b_load, other=0.0)
        
        # --- Perform Dot Product ---
        # tl.dot(A_tile[BM, BKc], B_tile_for_dot[BKc, BN]) -> C_tile_accumulator[BM, BN]
        # A_tile shape: (BLOCK_SIZE_M, BLOCK_SIZE_K_COMMON)
        # B_tile_for_dot shape: (BLOCK_SIZE_K_COMMON, BLOCK_SIZE_N)
        # Output shape: (BLOCK_SIZE_M, BLOCK_SIZE_N)
        dot_output_f32 = tl.dot(a_block_fp8, b_block_fp8_for_dot, out_dtype=tl.float32, allow_tf32=False)
        
        # Apply scaling factors
        # a_scales_m_vector is (BLOCK_SIZE_M), needs to be (BLOCK_SIZE_M, 1) for broadcasting
        # b_scale_val is a scalar
        scaled_dot_output = dot_output_f32 * a_scales_m_vector[:, None] * b_scale_val
        accumulator += scaled_dot_output

    # --- Store C tile ---
    # C is M x N, row-major.
    # stride_c_m_dim is N (stride to get to the next M row in C)
    # stride_c_n_dim is 1 (N is the contiguous dim for a fixed M in C)
    # offs_m_for_c_tile[:, None] provides M-offsets for rows.
    # offs_n_for_c_tile[None, :] provides N-offsets for columns.
    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 C store: Both M and N dimensions are masked.
    mask_c_store = mask_m_for_ac[:, None] & mask_n_for_bc[None, :] # Shape (BLOCK_SIZE_M, BLOCK_SIZE_N)
    # Store the accumulated result, converted to bfloat16.
    tl.store(c_tile_ptrs, accumulator.to(tl.bfloat16), mask=mask_c_store)


def custom_kernel(kernel_input_data: tuple) -> torch.Tensor:
    """
    Host wrapper for launching the Triton FP8 GEMM kernel.
    It sets up strides, grid dimensions, and tuning parameters.
    """
    a, b, a_scale, b_scale, c = kernel_input_data

    M, K_orig = a.shape  # a is M x K
    N, K_check = b.shape # b is N x K (problem states b is N x K col-major, so shape is K,N if .T in generate_input)
                         # generate_input returns (a.T, b.T, ...). So a is KxM, b is KxN originally.
                         # After .T, a is M x K, b is N x K. This matches the kernel's expectation for shapes.

    # Validate K dimensions
    if K_orig != K_check:
        raise ValueError(f"K dimension mismatch between a ({K_orig}) and b ({K_check})")
    K_common = K_orig

    # Strides based on input tensor definitions (as received by this function)
    # a: M x K_common, e4m3fnuz, column-major. Access A[m, k] with ptr + k*M + m
    stride_a_m_dim = 1        # Stride for M dimension (fastest changing for col-major if accessing A[m,k])
    stride_a_k_dim = M        # Stride for K dimension (to move to next element in K direction)

    # b: N x K_common, e4m3fnuz, column-major. Access B[n, k] with ptr + k*N + n
    stride_b_n_dim = 1        # Stride for N dimension
    stride_b_k_dim = N        # Stride for K dimension

    # a_scale: M x (K_common // 128), fp32, column-major
    # Access A_scale[m, k_scaled] with ptr + k_scaled*M + m
    stride_a_scale_m_dim = 1
    stride_a_scale_k_scaled_dim = M # Stride to move to the next k_scaled block

    # b_scale: (N // 128) x (K_common // 128), fp32, column-major
    # Access B_scale[n_scaled, k_scaled] with ptr + k_scaled*(N//128) + n_scaled
    # b_scale.shape[0] is (N // 128)
    stride_b_scale_n_scaled_dim = 1
    stride_b_scale_k_scaled_dim = b_scale.shape[0] if b_scale.nelement() > 0 else 1 # Stride to move to next k_scaled

    # c: M x N, bf16, ROW-major (output)
    # Access C[m, n] with ptr + m*N + n
    stride_c_m_dim = N        # Stride for M dimension (to move to next row)
    stride_c_n_dim = 1        # Stride for N dimension (to move to next element in a row)

    # --- Kernel Launch Parameters ---
    # These are constexpr in the kernel, so they are passed as Python literals.
    
    # Optimized parameters:
    BLOCK_SIZE_M_val = 128          # C-tile M-dimension. Kept at 128.
    BLOCK_SIZE_N_val = 128          # C-tile N-dimension. Increased from 32 to 128.
    BLOCK_SIZE_K_COMMON_val = 128   # Common K-dimension for matmul A[BMxBKc] @ B_T[BKcxBN]. Kept at 128.

    A_B_SCALE_K_GRANULARITY_val = 128 # K granularity for scales. Matches problem.
    B_SCALE_N_GRANULARITY_val = 128   # N granularity for b_scale. Matches problem.

    # Grid calculation: How many thread blocks to launch
    # cdiv(X, BLOCK_X) ensures enough blocks are launched to cover the entire dimension.
    grid_m = triton.cdiv(M, BLOCK_SIZE_M_val)
    grid_n = triton.cdiv(N, BLOCK_SIZE_N_val)
    grid = (grid_m, grid_n)

    # --- Tuning Parameters for the Kernel Launch ---
    # num_warps: Number of warps (groups of 32 threads) per thread block.
    # num_stages: How many stages of K-loop to pipeline in shared memory.
    
    # Optimized parameters:
    num_warps_val = 8               # Increased from 4 to 8 for the larger 128x128 C-tile.
                                    # (128*128 C-tile elements) / (8 warps * 32 threads/warp) = 64 elements per thread.
    num_stages_val = 2              # Adjusted from 3 to 2.
                                    # Shared memory per stage for 128x128 C-tile:
                                    # A_tile (128x128xFP8) = 16KB
                                    # B_tile (128x128xFP8) = 16KB
                                    # Total per stage = 32KB.
                                    # 2 stages * 32KB = 64KB, fitting well into typical L1/shared memory.
    
    # Print launch parameters for debugging (optional)
    # print(f"Grid: ({grid_m}, {grid_n}), BlockM: {BLOCK_SIZE_M_val}, BlockN: {BLOCK_SIZE_N_val}, BlockK: {BLOCK_SIZE_K_COMMON_val}")
    # print(f"NumWarps: {num_warps_val}, NumStages: {num_stages_val}")
    # print(f"M={M}, N={N}, K={K_common}")
    # print(f"a_scale shape: {a_scale.shape}, b_scale shape: {b_scale.shape}")
    # print(f"stride_a_scale_k_scaled_dim: {stride_a_scale_k_scaled_dim}, stride_b_scale_k_scaled_dim: {stride_b_scale_k_scaled_dim}")


    matmul_kernel_fp8dot_final[grid](
        a, b, a_scale, b_scale, c, # Data pointers
        M, N, K_common,             # 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,
        # Triton 'constexpr' arguments, passed by Python literal value
        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,
        # Kernel launch tuning parameters
        num_warps=num_warps_val,
        num_stages=num_stages_val
    )
    return c

python eval.py benchmark /tmp/tmpcsnmpudz