//! Flash Attention Metal Shaders
//!
//! Metal Shading Language implementation of Flash Attention for Apple Silicon
//! Based on the Flash Attention algorithm by Dao et al. (2022)
//!
//! Implements tiled attention with online softmax for O(N) memory complexity.

#include <metal_stdlib>
using namespace metal;

// Block sizes for tiling - tuned for Apple Silicon's 32KB threadgroup memory limit
// Forward kernel: 32x32 blocks → Q(8KB) + K(8KB) + V(8KB) + S(4KB) = 28KB ✓
// Backward kernels: 16x16 blocks → Q+K+V+dO+O (5 tiles * 4KB) = 20KB ✓
constant uint BLOCK_Q_FWD = 32;
constant uint BLOCK_KV_FWD = 32;
constant uint BLOCK_Q_BWD = 16;
constant uint BLOCK_KV_BWD = 16;
constant uint MAX_HEAD_DIM = 128;

/// Flash Attention kernel parameters
struct FlashAttentionParams {
    uint batch_size;
    uint num_heads;
    uint seq_len_q;
    uint seq_len_kv;
    uint head_dim;
    float softmax_scale;
    uint causal;  // 0 = false, 1 = true
    uint block_size_q;
    uint block_size_kv;
};

/// Helper: Simdgroup reduction for max
inline float simdgroup_reduce_max(float val) {
    for (uint offset = 16; offset > 0; offset >>= 1) {
        val = max(val, simd_shuffle_down(val, offset));
    }
    return simd_broadcast_first(val);
}

/// Helper: Simdgroup reduction for sum
inline float simdgroup_reduce_sum(float val) {
    for (uint offset = 16; offset > 0; offset >>= 1) {
        val += simd_shuffle_down(val, offset);
    }
    return simd_broadcast_first(val);
}

/// Flash Attention forward kernel
///
/// Implements the Flash Attention algorithm with:
/// - Tiled loading of Q, K, V into threadgroup memory
/// - Online softmax with running max/sum
/// - Causal masking support
/// - Optimized memory access patterns for Apple Silicon
kernel void flash_attention_forward(
    device const half* Q [[buffer(0)]],
    device const half* K [[buffer(1)]],
    device const half* V [[buffer(2)]],
    device half* O [[buffer(3)]],
    device float* LSE [[buffer(4)]],
    constant FlashAttentionParams& params [[buffer(5)]],
    uint3 tgid [[threadgroup_position_in_grid]],
    uint3 tid [[thread_position_in_threadgroup]],
    uint3 tg_size [[threads_per_threadgroup]],
    uint thread_idx [[thread_index_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_group [[simdgroup_index_in_threadgroup]]
) {
    // Threadgroup memory for Q, K, V tiles (32KB limit on Apple Silicon)
    // Forward: 8KB + 8KB + 8KB + 4KB = 28KB ✓
    threadgroup half Q_tile[BLOCK_Q_FWD][MAX_HEAD_DIM];
    threadgroup half K_tile[BLOCK_KV_FWD][MAX_HEAD_DIM];
    threadgroup half V_tile[BLOCK_KV_FWD][MAX_HEAD_DIM];
    threadgroup float S_tile[BLOCK_Q_FWD][BLOCK_KV_FWD];  // Attention scores

    // Identify which batch, head, and Q block we're processing
    uint batch_idx = tgid.z;
    uint head_idx = tgid.y;
    uint q_block_idx = tgid.x;

    // Base offsets for this batch and head
    uint qkv_stride = params.seq_len_q * params.head_dim;
    uint head_offset = (batch_idx * params.num_heads + head_idx) * qkv_stride;

    // Q block start position
    uint q_start = q_block_idx * BLOCK_Q_FWD;

    // Thread's row within the Q block
    uint q_row = thread_idx;
    uint q_idx = q_start + q_row;

    // Online softmax state (per-thread for its Q row)
    float m_i = -INFINITY;  // Running max
    float l_i = 0.0f;       // Running sum of exp(scores - max)

    // Output accumulator (per-thread for its Q row)
    float o_acc[MAX_HEAD_DIM];
    for (uint d = 0; d < params.head_dim; d++) {
        o_acc[d] = 0.0f;
    }

    // Load Q tile into threadgroup memory (each thread loads one row)
    if (q_idx < params.seq_len_q && q_row < BLOCK_Q_FWD) {
        uint q_offset = head_offset + q_idx * params.head_dim;
        for (uint d = 0; d < params.head_dim; d++) {
            Q_tile[q_row][d] = Q[q_offset + d];
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Number of K/V blocks to iterate over
    uint num_kv_blocks = (params.seq_len_kv + BLOCK_KV_FWD - 1) / BLOCK_KV_FWD;

    // For causal attention, limit K/V blocks based on Q position
    uint max_kv_block = num_kv_blocks;
    if (params.causal) {
        // Only process K/V positions up to q_start + BLOCK_Q_FWD - 1
        max_kv_block = min(num_kv_blocks, (q_start + BLOCK_Q_FWD + BLOCK_KV_FWD - 1) / BLOCK_KV_FWD);
    }

    // Iterate over K/V blocks
    for (uint kv_block = 0; kv_block < max_kv_block; kv_block++) {
        uint kv_start = kv_block * BLOCK_KV_FWD;

        // Load K and V tiles (collaborative loading)
        if (q_row < BLOCK_KV_FWD) {
            uint kv_idx = kv_start + q_row;
            if (kv_idx < params.seq_len_kv) {
                uint kv_offset = head_offset + kv_idx * params.head_dim;
                for (uint d = 0; d < params.head_dim; d++) {
                    K_tile[q_row][d] = K[kv_offset + d];
                    V_tile[q_row][d] = V[kv_offset + d];
                }
            } else {
                // Padding for out-of-bounds
                for (uint d = 0; d < params.head_dim; d++) {
                    K_tile[q_row][d] = half(0.0);
                    V_tile[q_row][d] = half(0.0);
                }
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Skip if this thread's Q row is out of bounds
        if (q_idx >= params.seq_len_q || q_row >= BLOCK_Q_FWD) {
            threadgroup_barrier(mem_flags::mem_threadgroup);
            continue;
        }

        // Compute attention scores for this Q row against all K rows in the block
        // S[i,j] = Q[i] @ K[j]^T * scale
        float row_max = -INFINITY;

        for (uint j = 0; j < BLOCK_KV_FWD; j++) {
            uint kv_idx = kv_start + j;

            // Skip if K/V position is out of bounds
            if (kv_idx >= params.seq_len_kv) {
                S_tile[q_row][j] = -INFINITY;
                continue;
            }

            // Causal mask: future positions get -infinity
            if (params.causal && kv_idx > q_idx) {
                S_tile[q_row][j] = -INFINITY;
                continue;
            }

            // Dot product Q[i] @ K[j]
            float score = 0.0f;
            for (uint d = 0; d < params.head_dim; d++) {
                score += float(Q_tile[q_row][d]) * float(K_tile[j][d]);
            }
            score *= params.softmax_scale;
            S_tile[q_row][j] = score;
            row_max = max(row_max, score);
        }

        // Online softmax update
        // m_new = max(m_i, row_max)
        // l_new = l_i * exp(m_i - m_new) + sum(exp(S - m_new))
        // O_new = O_i * exp(m_i - m_new) + sum(exp(S - m_new) * V)

        float m_new = max(m_i, row_max);
        float exp_diff = exp(m_i - m_new);

        // Rescale existing accumulator
        l_i *= exp_diff;
        for (uint d = 0; d < params.head_dim; d++) {
            o_acc[d] *= exp_diff;
        }

        // Accumulate new contributions
        for (uint j = 0; j < BLOCK_KV_FWD; j++) {
            float score = S_tile[q_row][j];
            if (score > -INFINITY * 0.5f) {  // Skip masked positions
                float exp_score = exp(score - m_new);
                l_i += exp_score;

                // Accumulate weighted V
                for (uint d = 0; d < params.head_dim; d++) {
                    o_acc[d] += exp_score * float(V_tile[j][d]);
                }
            }
        }

        m_i = m_new;
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write final output: O = acc / l_i
    if (q_idx < params.seq_len_q && q_row < BLOCK_Q_FWD) {
        float inv_l = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f;
        uint o_offset = head_offset + q_idx * params.head_dim;

        for (uint d = 0; d < params.head_dim; d++) {
            O[o_offset + d] = half(o_acc[d] * inv_l);
        }

        // Store log-sum-exp for backward pass: LSE = m + log(l)
        uint lse_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q + q_idx;
        LSE[lse_offset] = m_i + log(max(l_i, 1e-10f));
    }
}

/// Flash Attention backward kernel for dQ
///
/// Computes gradient with respect to Q:
/// dQ[i] = sum_j( softmax_grad[i,j] * K[j] )
/// where softmax_grad[i,j] = P[i,j] * (dO[i] @ V[j]^T - D[i])
/// and D[i] = sum_j(P[i,j] * dO[i] @ V[j]^T) = O[i] @ dO[i]^T (element-wise then sum)
kernel void flash_attention_backward_dq(
    device const half* dO [[buffer(0)]],
    device const half* Q [[buffer(1)]],
    device const half* K [[buffer(2)]],
    device const half* V [[buffer(3)]],
    device const half* O [[buffer(4)]],
    device const float* LSE [[buffer(5)]],
    device half* dQ [[buffer(6)]],
    constant FlashAttentionParams& params [[buffer(7)]],
    uint3 tgid [[threadgroup_position_in_grid]],
    uint3 tid [[thread_position_in_threadgroup]],
    uint thread_idx [[thread_index_in_threadgroup]]
) {
    // Threadgroup memory for tiles (32KB limit on Apple Silicon)
    // Backward: 5 tiles * 16 * 128 * 2 = 20KB ✓
    threadgroup half Q_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
    threadgroup half K_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
    threadgroup half V_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
    threadgroup half dO_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
    threadgroup half O_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
    threadgroup float D_tile[BLOCK_Q_BWD];  // D[i] = dO[i] @ O[i]

    uint batch_idx = tgid.z;
    uint head_idx = tgid.y;
    uint q_block_idx = tgid.x;

    uint head_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q * params.head_dim;
    uint lse_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q;

    uint q_start = q_block_idx * BLOCK_Q_BWD;
    uint q_row = thread_idx;
    uint q_idx = q_start + q_row;

    // Output gradient accumulator
    float dq_acc[MAX_HEAD_DIM];
    for (uint d = 0; d < params.head_dim; d++) {
        dq_acc[d] = 0.0f;
    }

    // Load Q, dO, O tiles and compute D = dO @ O (element-wise sum)
    if (q_idx < params.seq_len_q && q_row < BLOCK_Q_BWD) {
        uint offset = head_offset + q_idx * params.head_dim;
        float d_val = 0.0f;
        for (uint d = 0; d < params.head_dim; d++) {
            Q_tile[q_row][d] = Q[offset + d];
            dO_tile[q_row][d] = dO[offset + d];
            O_tile[q_row][d] = O[offset + d];
            d_val += float(dO_tile[q_row][d]) * float(O_tile[q_row][d]);
        }
        D_tile[q_row] = d_val;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Iterate over K/V blocks
    uint num_kv_blocks = (params.seq_len_kv + BLOCK_KV_BWD - 1) / BLOCK_KV_BWD;
    uint max_kv_block = num_kv_blocks;
    if (params.causal) {
        max_kv_block = min(num_kv_blocks, (q_start + BLOCK_Q_BWD + BLOCK_KV_BWD - 1) / BLOCK_KV_BWD);
    }

    for (uint kv_block = 0; kv_block < max_kv_block; kv_block++) {
        uint kv_start = kv_block * BLOCK_KV_BWD;

        // Load K, V tiles
        if (q_row < BLOCK_KV_BWD) {
            uint kv_idx = kv_start + q_row;
            if (kv_idx < params.seq_len_kv) {
                uint offset = head_offset + kv_idx * params.head_dim;
                for (uint d = 0; d < params.head_dim; d++) {
                    K_tile[q_row][d] = K[offset + d];
                    V_tile[q_row][d] = V[offset + d];
                }
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        if (q_idx >= params.seq_len_q || q_row >= BLOCK_Q_BWD) {
            threadgroup_barrier(mem_flags::mem_threadgroup);
            continue;
        }

        float lse_val = LSE[lse_offset + q_idx];

        for (uint j = 0; j < BLOCK_KV_BWD; j++) {
            uint kv_idx = kv_start + j;
            if (kv_idx >= params.seq_len_kv) continue;
            if (params.causal && kv_idx > q_idx) continue;

            // Recompute attention score
            float score = 0.0f;
            for (uint d = 0; d < params.head_dim; d++) {
                score += float(Q_tile[q_row][d]) * float(K_tile[j][d]);
            }
            score *= params.softmax_scale;

            // Recompute P[i,j] = exp(score - LSE[i])
            float p_ij = exp(score - lse_val);

            // Compute dO[i] @ V[j]
            float dov = 0.0f;
            for (uint d = 0; d < params.head_dim; d++) {
                dov += float(dO_tile[q_row][d]) * float(V_tile[j][d]);
            }

            // softmax_grad = P * (dO @ V - D)
            float softmax_grad = p_ij * (dov - D_tile[q_row]);

            // dQ[i] += softmax_grad * K[j] * scale
            for (uint d = 0; d < params.head_dim; d++) {
                dq_acc[d] += softmax_grad * float(K_tile[j][d]) * params.softmax_scale;
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write dQ output
    if (q_idx < params.seq_len_q && q_row < BLOCK_Q_BWD) {
        uint offset = head_offset + q_idx * params.head_dim;
        for (uint d = 0; d < params.head_dim; d++) {
            dQ[offset + d] = half(dq_acc[d]);
        }
    }
}

/// Flash Attention backward kernel for dK and dV
///
/// Computes gradients with respect to K and V:
/// dK[j] = sum_i( softmax_grad[i,j] * Q[i] )
/// dV[j] = sum_i( P[i,j] * dO[i] )
kernel void flash_attention_backward_dkv(
    device const half* dO [[buffer(0)]],
    device const half* Q [[buffer(1)]],
    device const half* K [[buffer(2)]],
    device const half* V [[buffer(3)]],
    device const half* O [[buffer(4)]],
    device const float* LSE [[buffer(5)]],
    device half* dK [[buffer(6)]],
    device half* dV [[buffer(7)]],
    constant FlashAttentionParams& params [[buffer(8)]],
    uint3 tgid [[threadgroup_position_in_grid]],
    uint3 tid [[thread_position_in_threadgroup]],
    uint thread_idx [[thread_index_in_threadgroup]]
) {
    // Threadgroup memory for tiles (32KB limit on Apple Silicon)
    // Backward: 5 tiles * 16 * 128 * 2 = 20KB ✓
    threadgroup half Q_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
    threadgroup half K_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
    threadgroup half V_tile[BLOCK_KV_BWD][MAX_HEAD_DIM];
    threadgroup half dO_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
    threadgroup half O_tile[BLOCK_Q_BWD][MAX_HEAD_DIM];
    threadgroup float D_tile[BLOCK_Q_BWD];

    uint batch_idx = tgid.z;
    uint head_idx = tgid.y;
    uint kv_block_idx = tgid.x;

    uint head_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q * params.head_dim;
    uint lse_offset = (batch_idx * params.num_heads + head_idx) * params.seq_len_q;

    uint kv_start = kv_block_idx * BLOCK_KV_BWD;
    uint kv_row = thread_idx;
    uint kv_idx = kv_start + kv_row;

    // Output gradient accumulators
    float dk_acc[MAX_HEAD_DIM];
    float dv_acc[MAX_HEAD_DIM];
    for (uint d = 0; d < params.head_dim; d++) {
        dk_acc[d] = 0.0f;
        dv_acc[d] = 0.0f;
    }

    // Load K, V tiles
    if (kv_idx < params.seq_len_kv && kv_row < BLOCK_KV_BWD) {
        uint offset = head_offset + kv_idx * params.head_dim;
        for (uint d = 0; d < params.head_dim; d++) {
            K_tile[kv_row][d] = K[offset + d];
            V_tile[kv_row][d] = V[offset + d];
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Iterate over Q blocks
    uint num_q_blocks = (params.seq_len_q + BLOCK_Q_BWD - 1) / BLOCK_Q_BWD;

    for (uint q_block = 0; q_block < num_q_blocks; q_block++) {
        uint q_start = q_block * BLOCK_Q_BWD;

        // For causal: skip Q blocks that come before this K/V position
        if (params.causal && q_start + BLOCK_Q_BWD - 1 < kv_start) {
            continue;
        }

        // Load Q, dO, O tiles and compute D
        if (kv_row < BLOCK_Q_BWD) {
            uint q_idx = q_start + kv_row;
            if (q_idx < params.seq_len_q) {
                uint offset = head_offset + q_idx * params.head_dim;
                float d_val = 0.0f;
                for (uint d = 0; d < params.head_dim; d++) {
                    Q_tile[kv_row][d] = Q[offset + d];
                    dO_tile[kv_row][d] = dO[offset + d];
                    O_tile[kv_row][d] = O[offset + d];
                    d_val += float(dO_tile[kv_row][d]) * float(O_tile[kv_row][d]);
                }
                D_tile[kv_row] = d_val;
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        if (kv_idx >= params.seq_len_kv || kv_row >= BLOCK_KV_BWD) {
            threadgroup_barrier(mem_flags::mem_threadgroup);
            continue;
        }

        for (uint i = 0; i < BLOCK_Q_BWD; i++) {
            uint q_idx = q_start + i;
            if (q_idx >= params.seq_len_q) continue;
            if (params.causal && kv_idx > q_idx) continue;

            float lse_val = LSE[lse_offset + q_idx];

            // Recompute attention score
            float score = 0.0f;
            for (uint d = 0; d < params.head_dim; d++) {
                score += float(Q_tile[i][d]) * float(K_tile[kv_row][d]);
            }
            score *= params.softmax_scale;

            // Recompute P[i,j] = exp(score - LSE[i])
            float p_ij = exp(score - lse_val);

            // dV[j] += P[i,j] * dO[i]
            for (uint d = 0; d < params.head_dim; d++) {
                dv_acc[d] += p_ij * float(dO_tile[i][d]);
            }

            // Compute dO[i] @ V[j]
            float dov = 0.0f;
            for (uint d = 0; d < params.head_dim; d++) {
                dov += float(dO_tile[i][d]) * float(V_tile[kv_row][d]);
            }

            // softmax_grad = P * (dO @ V - D)
            float softmax_grad = p_ij * (dov - D_tile[i]);

            // dK[j] += softmax_grad * Q[i] * scale
            for (uint d = 0; d < params.head_dim; d++) {
                dk_acc[d] += softmax_grad * float(Q_tile[i][d]) * params.softmax_scale;
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write dK, dV outputs
    if (kv_idx < params.seq_len_kv && kv_row < BLOCK_KV_BWD) {
        uint offset = head_offset + kv_idx * params.head_dim;
        for (uint d = 0; d < params.head_dim; d++) {
            dK[offset + d] = half(dk_acc[d]);
            dV[offset + d] = half(dv_acc[d]);
        }
    }
}