mlx-native 0.4.2

//! RMS Normalization GPU dispatch.
//!
//! Computes: `x * rsqrt(mean(x^2) + eps) * weight`
//!
//! The mean is computed over the last dimension.  eps=1e-6 is the standard
//! value for Gemma 4.

use metal::MTLSize;

use crate::buffer::MlxBuffer;
use crate::dtypes::DType;
use crate::encoder::{CapturedOpKind, CommandEncoder};
use crate::error::{MlxError, Result};
use crate::kernel_registry::KernelRegistry;

/// MSL source for the RMS norm kernels (embedded at compile time).
pub static RMS_NORM_SHADER_SOURCE: &str = include_str!("../shaders/rms_norm.metal");

/// Register RMS norm shader sources with the given kernel registry.
pub fn register(registry: &mut KernelRegistry) {
    registry.register_source("rms_norm_f32", RMS_NORM_SHADER_SOURCE);
    registry.register_source("rms_norm_f16", RMS_NORM_SHADER_SOURCE);
    registry.register_source("rms_norm_bf16", RMS_NORM_SHADER_SOURCE);
    registry.register_source("rms_norm_no_scale_bf16", RMS_NORM_SHADER_SOURCE);
    registry.register_source("rms_norm_no_scale_f32", RMS_NORM_SHADER_SOURCE);
    // Fused RMS norm + elementwise multiply (Phase 4e.2)
    registry.register_source("rms_norm_mul_f32", RMS_NORM_SHADER_SOURCE);
    registry.register_source("rms_norm_mul_f16", RMS_NORM_SHADER_SOURCE);
    registry.register_source("rms_norm_mul_bf16", RMS_NORM_SHADER_SOURCE);
}

/// Select the fused RMS norm + multiply kernel name based on input dtype.
fn fused_rms_norm_mul_kernel_name(dtype: DType) -> Result<&'static str> {
    match dtype {
        DType::F32 => Ok("rms_norm_mul_f32"),
        DType::F16 => Ok("rms_norm_mul_f16"),
        DType::BF16 => Ok("rms_norm_mul_bf16"),
        _ => Err(MlxError::InvalidArgument(format!(
            "Fused RMS norm+mul unsupported dtype: {}",
            dtype
        ))),
    }
}

/// Dispatch an RMS normalization operation on the GPU.
///
/// # Arguments
///
/// * `encoder`    - Command encoder to record the dispatch into.
/// * `registry`   - Kernel registry (must have rms_norm sources registered).
/// * `device`     - Metal device for pipeline compilation.
/// * `input`      - Input buffer of shape `[rows, dim]` (f32, f16, or bf16).
/// * `weight`     - Weight buffer of shape `[dim]` (same dtype as input; f32 for f32/f16 kernels, bf16 for bf16).
/// * `output`     - Output buffer (same dtype and shape as input).
/// * `params_buf` - Params buffer containing `[eps, dim]` as f32.
/// * `rows`       - Number of rows to normalize.
/// * `dim`        - Dimension of the last axis.
///
/// # Errors
///
/// Returns `MlxError::InvalidArgument` if:
/// - Input dtype is not f32, f16, or bf16.
/// - Input element count does not match rows * dim.
pub fn dispatch_rms_norm(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    input: &MlxBuffer,
    weight: &MlxBuffer,
    output: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    dim: u32,
) -> Result<()> {
    if rows == 0 || dim == 0 {
        return Err(MlxError::InvalidArgument(
            "RMS norm rows and dim must be > 0".into(),
        ));
    }

    let expected = (rows as usize) * (dim as usize);
    if input.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "RMS norm input element count {} != rows({}) * dim({})",
            input.element_count(),
            rows,
            dim
        )));
    }
    if output.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "RMS norm output element count {} != rows({}) * dim({})",
            output.element_count(),
            rows,
            dim
        )));
    }

    let kernel_name = match input.dtype() {
        DType::F32 => "rms_norm_f32",
        DType::F16 => "rms_norm_f16",
        DType::BF16 => "rms_norm_bf16",
        _ => {
            return Err(MlxError::InvalidArgument(format!(
                "RMS norm unsupported dtype: {}",
                input.dtype()
            )));
        }
    };

    let pipeline = registry.get_pipeline(kernel_name, device)?;

    // One threadgroup per row.  Threadgroup size must be a power of 2
    // for the tree reduction to work correctly.
    let tg_size = std::cmp::min(256, dim.next_power_of_two()) as u64;

    // Threadgroup shared memory: tg_size floats for the reduction.
    let shared_mem_bytes = tg_size * 4; // sizeof(float) = 4

    // Tag for the fusion pass (Phase 4e.2): RMS norm can fuse with a
    // subsequent elementwise multiply.
    encoder.set_op_kind(CapturedOpKind::RmsNorm);

    encoder.encode_threadgroups_with_shared(
        pipeline,
        &[
            (0, input),
            (1, weight),
            (2, output),
            (3, params_buf),
        ],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}

/// Dispatch a fused 3-output RMS normalization (f32 only).
///
/// Computes RMS(input) ONCE, then applies three different per-element
/// weight vectors to produce three outputs:
///   output_a = input * rsqrt(mean(input^2) + eps) * weight_a
///   output_b = input * rsqrt(mean(input^2) + eps) * weight_b
///   output_c = input * rsqrt(mean(input^2) + eps) * weight_c
///
/// Replaces three separate `dispatch_rms_norm` calls when the input is
/// the same.  Saves 2 dispatches and reads `input` once instead of
/// three times.  Wave P4.9.
///
/// # Arguments
///
/// * `encoder`    - Command encoder to record the dispatch into.
/// * `registry`   - Kernel registry (must have rms_norm_f32_triple registered).
/// * `device`     - Metal device for pipeline compilation.
/// * `input`      - Shared input buffer of shape `[rows, dim]` (f32).
/// * `weight_a`   - First weight vector `[dim]` (f32).
/// * `weight_b`   - Second weight vector `[dim]` (f32).
/// * `weight_c`   - Third weight vector `[dim]` (f32).
/// * `output_a`   - First output buffer `[rows, dim]` (f32).
/// * `output_b`   - Second output buffer `[rows, dim]` (f32).
/// * `output_c`   - Third output buffer `[rows, dim]` (f32).
/// * `params_buf` - Params buffer containing `[eps, dim]` as f32.
/// * `rows`       - Number of rows to normalize.
/// * `dim`        - Dimension of the last axis.
///
/// # Errors
///
/// Returns `MlxError::InvalidArgument` if parameters are invalid.
pub fn dispatch_rms_norm_f32_triple(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    input: &MlxBuffer,
    weight_a: &MlxBuffer,
    weight_b: &MlxBuffer,
    weight_c: &MlxBuffer,
    output_a: &MlxBuffer,
    output_b: &MlxBuffer,
    output_c: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    dim: u32,
) -> Result<()> {
    if rows == 0 || dim == 0 {
        return Err(MlxError::InvalidArgument(
            "RMS norm triple: rows and dim must be > 0".into(),
        ));
    }

    let expected = (rows as usize) * (dim as usize);
    for (name, buf) in [
        ("input", input),
        ("output_a", output_a),
        ("output_b", output_b),
        ("output_c", output_c),
    ] {
        if buf.element_count() != expected {
            return Err(MlxError::InvalidArgument(format!(
                "RMS norm triple: {} element count {} != rows({}) * dim({})",
                name, buf.element_count(), rows, dim
            )));
        }
        if buf.dtype() != DType::F32 {
            return Err(MlxError::InvalidArgument(format!(
                "RMS norm triple: {} must be f32, got {}",
                name, buf.dtype()
            )));
        }
    }
    for (name, buf) in [
        ("weight_a", weight_a),
        ("weight_b", weight_b),
        ("weight_c", weight_c),
    ] {
        if buf.element_count() != dim as usize {
            return Err(MlxError::InvalidArgument(format!(
                "RMS norm triple: {} element count {} != dim({})",
                name, buf.element_count(), dim
            )));
        }
        if buf.dtype() != DType::F32 {
            return Err(MlxError::InvalidArgument(format!(
                "RMS norm triple: {} must be f32, got {}",
                name, buf.dtype()
            )));
        }
    }

    let pipeline = registry.get_pipeline("rms_norm_f32_triple", device)?;

    let tg_size = std::cmp::min(256, dim.next_power_of_two()) as u64;
    let shared_mem_bytes = tg_size * 4;

    encoder.encode_threadgroups_with_shared(
        pipeline,
        &[
            (0, input),
            (1, weight_a),
            (2, weight_b),
            (3, weight_c),
            (4, output_a),
            (5, output_b),
            (6, output_c),
            (7, params_buf),
        ],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}

/// Dispatch the Wave P4.18 fused post-attention norm+add + triple pre-FF norm.
///
/// Replaces the two dispatches
///   (a) `fused_norm_add_f32(hidden, attn_out, post_attn_w -> pf_residual)`
///   (b) `rms_norm_f32_triple(pf_residual, w_a/b/c -> out_a/b/c)`
/// with one kernel that:
///   1. computes rms_inv over attn_out
///   2. writes `residual_out = hidden + attn_out * rms_inv * post_attn_w` AND
///      accumulates sum(residual_out^2) in the same pass
///   3. computes rms_inv over the residual stream
///   4. reads `residual_out` back, applies three weight vectors, writes 3 outputs
///
/// Eliminates one write+read of the residual buffer (~60 MB / layer @ pp2455
/// x 30 layers = ~1.8 GB of memory traffic) and one dispatch per layer.
///
/// # Arguments
/// * `encoder`      - Command encoder to record the dispatch into.
/// * `registry`     - Kernel registry (must have fused_post_attn_triple_norm_f32).
/// * `device`       - Metal device for pipeline compilation.
/// * `hidden`       - Pre-attention residual stream, f32 `[rows, dim]`.
/// * `attn_out`     - Attention O-proj output, f32 `[rows, dim]`.
/// * `post_attn_w`  - Post-attention layernorm weight, f32 `[dim]`.
/// * `weight_a`     - First pre-FF weight (e.g. MLP norm), f32 `[dim]`.
/// * `weight_b`     - Second pre-FF weight (e.g. MoE input norm), f32 `[dim]`.
/// * `weight_c`     - Third pre-FF weight (e.g. router norm), f32 `[dim]`.
/// * `residual_out` - Output: residual stream (hidden + normed_attn), f32 `[rows, dim]`.
/// * `output_a/b/c` - Triple-normed outputs, f32 `[rows, dim]`.
/// * `eps`          - RMS epsilon.
/// * `rows`         - Number of rows (= seq_len).
/// * `dim`          - Model hidden size.
#[allow(clippy::too_many_arguments)]
pub fn dispatch_fused_post_attn_triple_norm_f32(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    hidden:       &MlxBuffer,
    attn_out:     &MlxBuffer,
    post_attn_w:  &MlxBuffer,
    weight_a:     &MlxBuffer,
    weight_b:     &MlxBuffer,
    weight_c:     &MlxBuffer,
    residual_out: &MlxBuffer,
    output_a:     &MlxBuffer,
    output_b:     &MlxBuffer,
    output_c:     &MlxBuffer,
    eps:          f32,
    rows:         u32,
    dim:          u32,
) -> Result<()> {
    if rows == 0 || dim == 0 {
        return Err(MlxError::InvalidArgument(
            "fused_post_attn_triple_norm_f32: rows and dim must be > 0".into(),
        ));
    }
    let expected = (rows as usize) * (dim as usize);
    for (name, buf) in [
        ("hidden",       hidden),
        ("attn_out",     attn_out),
        ("residual_out", residual_out),
        ("output_a",     output_a),
        ("output_b",     output_b),
        ("output_c",     output_c),
    ] {
        if buf.element_count() < expected {
            return Err(MlxError::InvalidArgument(format!(
                "fused_post_attn_triple_norm_f32: {} size {} < rows({}) * dim({})",
                name, buf.element_count(), rows, dim
            )));
        }
        if buf.dtype() != DType::F32 {
            return Err(MlxError::InvalidArgument(format!(
                "fused_post_attn_triple_norm_f32: {} must be f32, got {}",
                name, buf.dtype()
            )));
        }
    }
    for (name, buf) in [
        ("post_attn_w", post_attn_w),
        ("weight_a",    weight_a),
        ("weight_b",    weight_b),
        ("weight_c",    weight_c),
    ] {
        if buf.element_count() != dim as usize {
            return Err(MlxError::InvalidArgument(format!(
                "fused_post_attn_triple_norm_f32: {} size {} != dim({})",
                name, buf.element_count(), dim
            )));
        }
    }

    let pipeline = registry.get_pipeline("fused_post_attn_triple_norm_f32", device)?;

    let tg_size = std::cmp::min(256u32, dim.next_power_of_two()) as u64;
    let shared_mem_bytes = tg_size * 4;

    #[repr(C)]
    #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
    struct Params { eps: f32, dim: u32 }
    let params = Params { eps, dim };

    use super::encode_helpers::{as_bytes, KernelArg};
    encoder.set_op_kind(CapturedOpKind::Other);
    encoder.encode_threadgroups_with_args_and_shared(
        pipeline,
        &[
            (0,  KernelArg::Buffer(hidden)),
            (1,  KernelArg::Buffer(attn_out)),
            (2,  KernelArg::Buffer(post_attn_w)),
            (3,  KernelArg::Buffer(weight_a)),
            (4,  KernelArg::Buffer(weight_b)),
            (5,  KernelArg::Buffer(weight_c)),
            (6,  KernelArg::Buffer(residual_out)),
            (7,  KernelArg::Buffer(output_a)),
            (8,  KernelArg::Buffer(output_b)),
            (9,  KernelArg::Buffer(output_c)),
            (10, KernelArg::Bytes(as_bytes(&params))),
        ],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}

/// Dispatch an RMS normalization without learned scale (bf16 only).
///
/// Computes: `output = x * rsqrt(mean(x^2) + eps)` — no weight multiplication.
/// Used for per-head V normalization in Gemma 4.
///
/// # Arguments
///
/// * `encoder`    - Command encoder to record the dispatch into.
/// * `registry`   - Kernel registry (must have rms_norm_no_scale_bf16 registered).
/// * `device`     - Metal device for pipeline compilation.
/// * `input`      - Input buffer of shape `[rows, dim]` (bf16).
/// * `output`     - Output buffer (same dtype and shape as input).
/// * `params_buf` - Params buffer containing `[eps, dim]` as f32.
/// * `rows`       - Number of rows to normalize.
/// * `dim`        - Dimension of the last axis.
///
/// # Errors
///
/// Returns `MlxError::InvalidArgument` if parameters are invalid.
pub fn dispatch_rms_norm_no_scale_bf16(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    input: &MlxBuffer,
    output: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    dim: u32,
) -> Result<()> {
    if rows == 0 || dim == 0 {
        return Err(MlxError::InvalidArgument(
            "RMS norm no_scale: rows and dim must be > 0".into(),
        ));
    }

    let expected = (rows as usize) * (dim as usize);
    if input.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "RMS norm no_scale: input element count {} != rows({}) * dim({})",
            input.element_count(),
            rows,
            dim
        )));
    }
    if output.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "RMS norm no_scale: output element count {} != rows({}) * dim({})",
            output.element_count(),
            rows,
            dim
        )));
    }

    let pipeline = registry.get_pipeline("rms_norm_no_scale_bf16", device)?;

    // One threadgroup per row.  Threadgroup size must be a power of 2
    // for the tree reduction to work correctly.
    let tg_size = std::cmp::min(256, dim.next_power_of_two()) as u64;

    // Threadgroup shared memory: tg_size floats for the reduction.
    let shared_mem_bytes = tg_size * 4; // sizeof(float) = 4

    encoder.encode_threadgroups_with_shared(
        pipeline,
        &[
            (0, input),
            (1, output),
            (2, params_buf),
        ],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}

/// Dispatch an RMS normalization without learned scale (f32).
///
/// Computes: `output = x * rsqrt(mean(x^2) + eps)` -- no weight multiplication.
/// Used for per-head V normalization in Gemma 4 when activations are f32.
///
/// # Arguments
///
/// * `encoder`    - Command encoder to record the dispatch into.
/// * `registry`   - Kernel registry (must have rms_norm_no_scale_f32 registered).
/// * `device`     - Metal device for pipeline compilation.
/// * `input`      - Input buffer of shape `[rows, dim]` (f32).
/// * `output`     - Output buffer (same dtype and shape as input).
/// * `params_buf` - Params buffer containing `[eps, dim]` as f32.
/// * `rows`       - Number of rows to normalize.
/// * `dim`        - Dimension of the last axis.
///
/// # Errors
///
/// Returns `MlxError::InvalidArgument` if parameters are invalid.
pub fn dispatch_rms_norm_no_scale_f32(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    input: &MlxBuffer,
    output: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    dim: u32,
) -> Result<()> {
    if rows == 0 || dim == 0 {
        return Err(MlxError::InvalidArgument(
            "RMS norm no_scale f32: rows and dim must be > 0".into(),
        ));
    }

    let expected = (rows as usize) * (dim as usize);
    if input.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "RMS norm no_scale f32: input element count {} != rows({}) * dim({})",
            input.element_count(),
            rows,
            dim
        )));
    }
    if output.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "RMS norm no_scale f32: output element count {} != rows({}) * dim({})",
            output.element_count(),
            rows,
            dim
        )));
    }

    let pipeline = registry.get_pipeline("rms_norm_no_scale_f32", device)?;

    let tg_size = std::cmp::min(256, dim.next_power_of_two()) as u64;
    let shared_mem_bytes = tg_size * 4;

    encoder.encode_threadgroups_with_shared(
        pipeline,
        &[
            (0, input),
            (1, output),
            (2, params_buf),
        ],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}

/// Dispatch a fused RMS normalization + elementwise multiply.
///
/// Computes: `output = (input * rsqrt(mean(input^2) + eps) * weight) * scale`
///
/// This replaces the two-dispatch pattern:
///   1. `rms_norm(input, weight) -> tmp`
///   2. `elementwise_mul(tmp, scale) -> output`
///
/// with a single kernel pass, eliminating one barrier and one global memory
/// round-trip for the intermediate `tmp` buffer.
///
/// # Arguments
///
/// * `encoder`      - Command encoder to record the dispatch into.
/// * `registry`     - Kernel registry.
/// * `device`       - Metal device for pipeline compilation.
/// * `input`        - Input buffer of shape `[rows, dim]`.
/// * `norm_weight`  - Norm weight buffer of shape `[dim]`.
/// * `scale_weight` - Scale (MUL operand) buffer of shape `[rows, dim]`.
/// * `output`       - Output buffer of shape `[rows, dim]`.
/// * `params_buf`   - Params buffer containing `[eps, dim]` as f32.
/// * `rows`         - Number of rows.
/// * `dim`          - Dimension of the last axis.
#[allow(clippy::too_many_arguments)]
pub fn dispatch_rms_norm_mul(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    input: &MlxBuffer,
    norm_weight: &MlxBuffer,
    scale_weight: &MlxBuffer,
    output: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    dim: u32,
) -> Result<()> {
    if rows == 0 || dim == 0 {
        return Err(MlxError::InvalidArgument(
            "Fused RMS norm+mul: rows and dim must be > 0".into(),
        ));
    }

    let expected = (rows as usize) * (dim as usize);
    if input.element_count() != expected {
        return Err(MlxError::InvalidArgument(format!(
            "Fused RMS norm+mul: input element count {} != rows({}) * dim({})",
            input.element_count(),
            rows,
            dim
        )));
    }

    let kernel_name = fused_rms_norm_mul_kernel_name(input.dtype())?;
    let pipeline = registry.get_pipeline(kernel_name, device)?;

    let tg_size = std::cmp::min(256, dim.next_power_of_two()) as u64;
    let shared_mem_bytes = tg_size * 4; // sizeof(float) = 4

    encoder.encode_threadgroups_with_shared(
        pipeline,
        &[
            (0, input),
            (1, norm_weight),
            (2, scale_weight),
            (3, output),
            (4, params_buf),
        ],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}