import torch
import triton
import triton.language as tl

from task import input_t, output_t

# Triton kernel re-architected with MI300X characteristics as primary design targets.
# This version emphasizes:
# 1. Minimized Local Data Share (LDS) footprint: Through smaller K-tiles and num_stages=2.
# 2. Leveraging high HBM3 bandwidth: More frequent, smaller data loads due to reduced K-tile depth.
# 3. Utilizing ample registers (VGPRs): Smaller tiles and efficient tl.dot operations are intended
#    to allow accumulators and current operands to make better use of registers.
# 4. Targeting good occupancy: Via waves_per_eu hint and appropriate num_warps.
@triton.jit
def _gemm_a8w8_blockscale_kernel_mi300x(  # Renamed to signify MI300X focus
    a_ptr,
    b_ptr,
    c_ptr,
    a_scale_ptr,
    b_scale_ptr,
    M,
    N,
    K,
    stride_am,  # Stride for A's M dimension
    stride_ak,  # Stride for A's K dimension
    stride_bk,  # Stride for B's K dimension (effectively W's K-dim stride)
    stride_bn,  # Stride for B's N dimension (effectively W's N-dim stride)
    stride_cm,
    stride_cn,
    stride_ascale_m,
    stride_ascale_k,
    stride_bscale_k,  # K-dim stride for w_scale (transposed b_scale)
    stride_bscale_n,  # N-dim stride for w_scale
    GROUP_K: tl.constexpr,  # K-dimension granularity for scales (e.g., 128)
    GROUP_N: tl.constexpr,  # N-dimension granularity for B scales (e.g., 128 from b_scale shape)
    # MI300X Specific Tuning Parameters as Kernel Defaults:
    BLOCK_SIZE_M: tl.constexpr = 128,
    BLOCK_SIZE_N: tl.constexpr = 128,
    BLOCK_SIZE_K: tl.constexpr = 64,  # Reduced K-tile depth.
    # A_tile (128x64 fp8) = 8KB. B_tile (64x128 fp8) = 8KB. Total = 16KB/stage.
    # Smaller K-tiles -> more K-loop iterations -> "streaming" from HBM.
    num_stages: tl.constexpr = 2,  # Double-buffering for A/B tiles in LDS.
    # LDS for A/B tiles: 2 * 16KB = 32KB. Fits well in 64KiB LDS.
    # AMD guidance: num_stages <= 2 is often optimal.
    num_warps: tl.constexpr = 4,  # 4 warps * 64 lanes/warp (MI300X wavefront) = 256 threads/block.
    # A common, often efficient, workgroup size.
    waves_per_eu: tl.constexpr = 2,  # Target occupancy hint for AMD scheduler (2-4 typical).
    # Helps hide memory and instruction latencies.
    matrix_instr_nonkdim: tl.constexpr = 32,  # Hint for non-K dimension of hardware MFMAs (MI300X has 32x32).
):
    pid = tl.program_id(0)  # Kernel uses a 1D launch grid.
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    # Determine current M and N block this program instance is responsible for.
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    # --- Pointer Setup for Matrix Tiles ---
    # Offsets for the M, N dimensions for the current thread block's output tile.
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    # Offsets for iterating within a K-block tile [0 ... BLOCK_SIZE_K-1].
    offs_k_loop = tl.arange(0, BLOCK_SIZE_K)

    # Base pointers for A matrix. `a_ptr` is start of A (M,K).
    # `a_ptrs_base` will hold pointers to the start of each row in the current M-block of A.
    a_ptrs_base = a_ptr + offs_m[:, None] * stride_am

    # Base pointers for B matrix (which is W = B_transposed, K,N, row-major). `b_ptr` is start of W.
    # `b_ptrs_base` will hold pointers to the start of each relevant column slice in W for the current N-block.
    # (Effectively, for W_kj, offs_n[None,:]*stride_bn gives offsets to W_k0, W_k1 .. for the N-block columns)
    b_ptrs_base = b_ptr + offs_n[None, :] * stride_bn

    # Accumulator tile, initialized to zeros. Uses fp32 for precision during accumulation.
    # Size: (BLOCK_SIZE_M, BLOCK_SIZE_N). For 128x128, this is 64KB if stored as fp32.
    # Triton aims to keep this in registers, spilling to LDS if necessary.
    # Reduced LDS usage by A/B tiles (due to smaller BLOCK_SIZE_K) helps manage overall LDS pressure.
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # --- Pointer Setup for Scale Factors ---
    # a_scale is (M, K/GROUP_K). `a_scale_ptr` points to it.
    # `a_scale_row_ptrs` points to the start of rows in `a_scale` for the current M-block.
    a_scale_row_ptrs = a_scale_ptr + offs_m * stride_ascale_m

    # b_scale (w_scale from host) is (K/GROUP_K, N/GROUP_N), row-major. `b_scale_ptr` points to it.
    # `n_group_idx` is the N-group index for b_scale corresponding to current N-block (pid_n).
    # This assumes BLOCK_SIZE_N is aligned with GROUP_N (both 128 in this setup).
    n_group_idx = pid_n
    # `b_scale_col_base_ptr` points to the start of the N-group slice in `w_scale` for the current N-block.
    # (i.e., &w_scale[0, n_group_idx])
    b_scale_col_base_ptr = b_scale_ptr + n_group_idx * stride_bscale_n

    # --- Main Loop over K Dimension ---
    # Iterates over K in chunks of `BLOCK_SIZE_K`.
    # Smaller `BLOCK_SIZE_K` increases iterations, enhancing "streaming" from HBM.
    for k_block_iter_idx in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        k_base_offset_in_k_dim = k_block_iter_idx * BLOCK_SIZE_K

        # --- Load A and B Tiles ---
        # Advance A pointers for the current K-block slice.
        cur_a_ptrs = a_ptrs_base + (k_base_offset_in_k_dim + offs_k_loop[None, :]) * stride_ak
        # Advance B (W) pointers for the current K-block slice.
        cur_b_ptrs = b_ptrs_base + (k_base_offset_in_k_dim + offs_k_loop[:, None]) * stride_bk

        # Boundary masks for A and B tiles (for M, N, K not perfectly divisible by block sizes).
        a_mask = (offs_m[:, None] < M) & ((k_base_offset_in_k_dim + offs_k_loop[None, :]) < K)
        b_mask = ((k_base_offset_in_k_dim + offs_k_loop[:, None]) < K) & (offs_n[None, :] < N)

        # Load A and B tiles (fp8 data). `other=0.0` pads with zeros for out-of-bounds accesses.
        a_tile = tl.load(cur_a_ptrs, mask=a_mask, other=0.0)
        b_tile = tl.load(cur_b_ptrs, mask=b_mask, other=0.0)

        # --- Load Scale Factors for Current K-Group ---
        # `k_group_for_scales_idx` is the index into the K-dimension of scale matrices (e.g., K_dim / GROUP_K).
        # Scales are reused if BLOCK_SIZE_K < GROUP_K.
        k_group_for_scales_idx = k_base_offset_in_k_dim // GROUP_K

        # Load a_scale: vector of size BLOCK_SIZE_M for the current M-block and K-group.
        a_scale_ptrs_for_k_group = a_scale_row_ptrs + k_group_for_scales_idx * stride_ascale_k
        a_scale_load_mask = offs_m < M  # Mask for rows of a_scale within M dimension.
        a_sc = tl.load(a_scale_ptrs_for_k_group, mask=a_scale_load_mask, other=0.0)

        # Load b_scale: scalar for the current N-block and K-group.
        b_scale_ptr_for_k_group = b_scale_col_base_ptr + k_group_for_scales_idx * stride_bscale_k
        # Assuming N-group and K-group indexing for b_scale is always valid or scales are padded.
        b_sc = tl.load(
            b_scale_ptr_for_k_group
        )  # Add mask/other if strict boundary check needed for b_scale itself.

        # --- Compute and Accumulate ---
        # `tl.dot` for fp8 inputs accumulates into fp32.
        # Scales (fp32) are applied after the dot product.
        # a_sc (BLOCK_SIZE_M,) is broadcast to (BLOCK_SIZE_M, 1). b_sc (scalar) broadcasts automatically.
        accumulator += tl.dot(a_tile, b_tile) * a_sc[:, None] * b_sc

    # --- Store Output Block ---
    # Convert accumulator from fp32 to bfloat16 for output.
    c_block = accumulator.to(tl.bfloat16)

    # Pointers to write the output C_block.
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    # Mask for writing C_block, to handle M, N not divisible by BLOCK_SIZE_M, BLOCK_SIZE_N.
    output_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c_block, mask=output_mask)


def custom_kernel(data):  # Type hint: data: input_t
    a_fp8, b_fp8, a_scale, b_scale, c = data

    # Prepare inputs as expected by the kernel:
    # A (x): M x K, column-major (kernel handles its strides directly)
    # B (w): K x N, row-major (host transposes original B_colmajor (N,K))
    # a_scale (x_scale): M x K/GROUP_K, column-major
    # b_scale (w_scale): K/GROUP_K x N/GROUP_N, row-major (host transposes original b_scale_colmajor)
    x = a_fp8
    w = b_fp8.T.contiguous()
    x_scale = a_scale
    w_scale = b_scale.T.contiguous()

    M_val, K_val_x = x.shape
    K_val_w, N_val = w.shape
    assert K_val_x == K_val_w, f"K dimension mismatch: A has {K_val_x}, B (as W) has {K_val_w}"
    K_val = K_val_x

    # Scale granularities are fixed by the problem's scale tensor shapes.
    GROUP_K_const = 128
    GROUP_N_const = 128  # This N-grouping is for b_scale, consistent with its (N/128, K/128) shape.

    # Grid calculation using the MI300X kernel's default block sizes.
    # These are BLOCK_SIZE_M=128, BLOCK_SIZE_N=128 in the kernel definition.
    kernel_default_block_m = 128
    kernel_default_block_n = 128
    grid_m = triton.cdiv(M_val, kernel_default_block_m)
    grid_n = triton.cdiv(N_val, kernel_default_block_n)
    grid = (grid_m * grid_n,)  # 1D grid of thread blocks.

    _gemm_a8w8_blockscale_kernel_mi300x[grid](
        x,  # a_ptr
        w,  # b_ptr (points to W, the transposed B)
        c,  # c_ptr
        x_scale,  # a_scale_ptr
        w_scale,  # b_scale_ptr (points to w_scale, transposed b_scale)
        M_val,
        N_val,
        K_val,
        x.stride(0),  # stride_am
        x.stride(1),  # stride_ak
        w.stride(0),  # stride_bk (K-dim stride of W, which is W.stride(0))
        w.stride(1),  # stride_bn (N-dim stride of W, which is W.stride(1))
        c.stride(0),  # stride_cm
        c.stride(1),  # stride_cn
        x_scale.stride(0),  # stride_ascale_m
        x_scale.stride(1),  # stride_ascale_k
        w_scale.stride(0),  # stride_bscale_k (K-dim stride of w_scale)
        w_scale.stride(1),  # stride_bscale_n (N-dim stride of w_scale)
        GROUP_K_const,
        GROUP_N_const,
        # BLOCK_SIZE_M, N, K, num_stages, etc., will use defaults from kernel definition.
    )
    return c

python eval.py benchmark /tmp/tmpel354385