#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpu MI300

from task import input_t, output_t
import torch
import triton
import triton.language as tl

NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count


@triton.jit
def kernel(
    A_ptr,          # Pointer to matrix A (M x K) in FP8 format
    B_ptr,          # Pointer to matrix B (K x N) in FP8 format
    A_scale_ptr,    # Pointer to A's scaling factors
    B_scale_ptr,    # Pointer to B's scaling factors
    C_ptr,          # Pointer to output matrix C (M x N)
    M: tl.constexpr,        # Number of rows in A and C
    N: tl.constexpr,        # Number of columns in B and C
    K: tl.constexpr,        # Inner dimension
    BLOCK_M: tl.constexpr,  # Block size for M dimension (dynamically chosen)
    BLOCK_N: tl.constexpr,  # Block size for N dimension (dynamically chosen)
    BLOCK_K: tl.constexpr,  # Block size for K dimension processing (dynamically chosen)
    BLOCK_Q: tl.constexpr = 128,    # Quantization block size (fixed at 128)
    TRANSPOSE: tl.constexpr = False, # Whether matrices are in transposed layout
):
    
    program_id = tl.program_id(0)
    num_pid_across_n = tl.cdiv(N, BLOCK_N)

    program_id_m = program_id // num_pid_across_n
    program_id_n = program_id % num_pid_across_n

    # ===== ADAPTIVE MEMORY STRIDE CALCULATION =====
    # Support for both row-major (default) and column-major (transposed) layouts
    # This flexibility allows for better memory access patterns depending on how
    # the input matrices are stored in memory
    if not TRANSPOSE:
        # Row-major layout: A[i,j] = A_ptr[i + j*M]
        A_stride_m, A_stride_k = 1, M      # Standard row-major for A
        B_stride_n, B_stride_k = 1, N      # Standard row-major for B
    else:
        # Column-major layout: A[i,j] = A_ptr[i*K + j]
        A_stride_m, A_stride_k = K, 1      # Column-major for A
        B_stride_n, B_stride_k = K, 1      # Column-major for B
    
    # Output is always stored in row-major format
    C_stride_m, C_stride_n = N, 1
    
    # ===== SCALING FACTOR MEMORY LAYOUT =====
    # Scale matrices are stored in column-major order with fixed block structure
    # A scales: 1x128 blocks, B scales: 128x128 blocks
    # Note: BLOCK_Q is always 128 for FP8 quantization
    A_scale_stride_m, A_scale_stride_k = 1, M
    B_scale_stride_n, B_scale_stride_k = 1, tl.cdiv(N, BLOCK_Q)

    # ===== CALCULATE BLOCK BOUNDARIES =====
    # Determine which portion of the output this program instance will compute
    offset_m = program_id_m * BLOCK_M  # Starting row
    offset_n = program_id_n * BLOCK_N  # Starting column

    # ===== CREATE INDEX ARRAYS =====
    # Generate arrays of indices for efficient block-wise memory access
    block_offsets_m = offset_m + tl.arange(0, BLOCK_M)
    block_offsets_n = offset_n + tl.arange(0, BLOCK_N)
    block_offsets_k = tl.arange(0, BLOCK_K)

    # ===== INITIALIZE MEMORY POINTERS =====
    # Create pointer grids for accessing blocks of A and B matrices
    # These pointers will be updated as we iterate through the K dimension
    A_block_ptrs = A_ptr + (
        block_offsets_m[:, None] * A_stride_m + block_offsets_k[None, :] * A_stride_k
    )
    B_block_ptrs = B_ptr + (
        block_offsets_k[:, None] * B_stride_k + block_offsets_n[None, :] * B_stride_n
    )
    
    # ===== SCALING FACTOR POINTERS =====
    # Since A scales are organized as 1x128 blocks, we need pointers for BLOCK_M rows
    # Since B scales are 128x128 blocks and N,K <= BLOCK_Q, B scale is often scalar
    A_scale_block_ptrs = A_scale_ptr + (block_offsets_m[:, None] * A_scale_stride_m)
    B_scale_block_ptrs = B_scale_ptr + (offset_n // BLOCK_Q) * B_scale_stride_n

    # ===== MAIN COMPUTATION ACCUMULATOR =====
    # Master accumulator for this program's output block (BLOCK_M x BLOCK_N)
    master_accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # ===== QUANTIZATION-AWARE K-DIMENSION PROCESSING =====
    # Process K dimension in chunks of BLOCK_Q (128) elements
    # Each chunk uses the same scaling factors, so we can apply them once per chunk
    num_k_iters = K // BLOCK_Q
    for _ in range(0, num_k_iters):
        # Inner accumulator for current 128-element quantization block
        inner_accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        
        # ===== FINE-GRAINED K-DIMENSION PROCESSING =====
        # Within each quantization block, process in smaller BLOCK_K chunks
        # This provides better register usage and memory access patterns
        for _ in tl.range(0, BLOCK_Q // BLOCK_K):
            # Load blocks from A and B matrices
            A_block = tl.load(A_block_ptrs)  # Shape: (BLOCK_M, BLOCK_K)
            B_block = tl.load(B_block_ptrs)  # Shape: (BLOCK_K, BLOCK_N)
            
            # Perform blocked matrix multiplication
            # This is the core tensor contraction operation
            inner_accumulator = tl.dot(
                A_block, B_block, inner_accumulator
            )  # Result shape: (BLOCK_M, BLOCK_N)

            # ===== ADVANCE TO NEXT K-SLICE =====
            # Move pointers to the next BLOCK_K elements within current quantization block
            A_block_ptrs += BLOCK_K * A_stride_k
            B_block_ptrs += BLOCK_K * B_stride_k

        # ===== APPLY FP8 DEQUANTIZATION SCALING =====
        # Load and apply scaling factors for the completed quantization block
        A_scales = tl.load(A_scale_block_ptrs)  # Shape: (BLOCK_M, 1)
        B_scales = tl.load(B_scale_block_ptrs)  # Shape: scalar
        
        # Apply dequantization: inner_result * (A_scale * B_scale)
        # This converts FP8 quantized results back to proper floating-point values
        master_accumulator += inner_accumulator * (A_scales * B_scales)

        # ===== ADVANCE TO NEXT QUANTIZATION BLOCK =====
        # Move scale pointers to the next 128-element block
        A_scale_block_ptrs += A_scale_stride_k
        B_scale_block_ptrs += B_scale_stride_k

    # ===== STORE RESULTS TO GLOBAL MEMORY =====
    # Calculate exact output indices and handle boundary conditions
    block_offsets_m = (
        program_id_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
    )  # Shape: (BLOCK_M, 1)
    block_offsets_n = (
        program_id_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]
    )  # Shape: (1, BLOCK_N)
    
    # Create mask to handle cases where block extends beyond matrix boundaries
    mask = (block_offsets_m < M) & (block_offsets_n < N)  # Shape: (BLOCK_M, BLOCK_N)
    
    # Calculate output memory addresses
    C_block_ptrs = C_ptr + (block_offsets_m * C_stride_m + block_offsets_n * C_stride_n)
    
    # Store final accumulated results
    tl.store(C_block_ptrs, master_accumulator, mask=mask)


@torch.compile(dynamic=False, mode="max-autotune-no-cudagraphs")
def contiguous(x):
    return x.contiguous()


def get_config(M, N, K):
    num_blocks_ref = (M // 128) * (N // 128)
    
    TRANSPOSE = False  # Default to row-major layout
    matrix_instr_nonkdim = 16  # Matrix instruction size for non-K dimensions
    BLOCK_M, BLOCK_N, BLOCK_K = (128, 128, 64)
    
    if num_blocks_ref * 8 < NUM_SMS:  
        # ===== SMALL WORKLOAD CASE =====
        # When we have very few blocks relative to SMs, use smaller tiles
        # This increases parallelism and improves GPU utilization
        # Trade-off: More blocks but each does less work
        BLOCK_M, BLOCK_N, BLOCK_K = (32, 64, 128)
        matrix_instr_nonkdim = 16
        
    elif num_blocks_ref < NUM_SMS:
        # ===== MEDIUM WORKLOAD CASE =====  
        # Moderate number of blocks - use balanced tile sizes
        # This provides good balance between parallelism and per-block efficiency
        BLOCK_M, BLOCK_N, BLOCK_K = (64, 64, 64)
        
    # ===== LARGE WORKLOAD CASE =====
    # For large workloads (num_blocks_ref >= NUM_SMS), use default large tiles
    # This maximizes per-block work and computational intensity
    # Default values: BLOCK_M, BLOCK_N, BLOCK_K = (128, 128, 64)
    
    # ===== PACKAGE CONFIGURATION =====
    config = dict(
        BLOCK_M=BLOCK_M,                    # Row block size
        BLOCK_N=BLOCK_N,                    # Column block size  
        BLOCK_K=BLOCK_K,                    # K-dimension block size
        waves_per_eu=2,                     # Number of wavefronts per execution unit
        matrix_instr_nonkdim=matrix_instr_nonkdim,  # Matrix instruction sizing
        num_warps=4,                        # Number of warps per thread block
        num_stages=2,                       # Pipeline stages for memory/compute overlap
        TRANSPOSE=TRANSPOSE,                # Memory layout configuration
    )
    return config

def custom_kernel(data: input_t) -> output_t:
    A_tensor, B_tensor, A_scale_tensor, B_scale_tensor, C_tensor = data

    M, K = A_tensor.shape
    N, _ = B_tensor.shape

    # ===== INTELLIGENT CONFIGURATION SELECTION =====
    # Use auto-tuning heuristics to select optimal kernel parameters
    # This adapts the kernel to the specific problem size and GPU characteristics
    config = get_config(M, N, K)

    # ===== CALCULATE LAUNCH GRID =====
    # Determine number of thread blocks needed based on selected configuration
    # Each block processes (BLOCK_M x BLOCK_N) elements of the output
    num_blocks = triton.cdiv(M, config["BLOCK_M"]) * triton.cdiv(N, config["BLOCK_N"])
    
    # ===== LAUNCH OPTIMIZED KERNEL =====
    # Launch with dynamically selected parameters for optimal performance
    kernel[(num_blocks,)](
        A_tensor, B_tensor, A_scale_tensor, B_scale_tensor, C_tensor, 
        M, N, K, 
        **config  # Unpack all configuration parameters
    )

    return C_tensor

python eval.py benchmark /tmp/tmpwho132o8