//! Runtime-generated Compiled Metal expressions from the pinned
//! metal/compiled.cpp emitter. The shared functor/header source is unchanged.
use super::array::{Array, Dtype};
use super::{compiled, ops::Operation};
use sha2::{Digest, Sha256};
use std::fmt::Write;

#[derive(Debug)]
struct Input {
    dtype: Dtype,
    scalar: bool,
    constant: Option<String>,
}

#[derive(Debug)]
pub(super) struct Statement {
    pub(super) dtype: Dtype,
    /// Reference primitive name; Broadcast/AsType emit a static cast.
    pub(super) operator: &'static str,
    pub(super) inputs: Vec<usize>,
}

pub(super) struct Expression {
    inputs: Vec<Input>,
    statements: Vec<Statement>,
    outputs: Vec<usize>,
    prefix: String,
}

pub(super) fn operator(op: &Operation) -> Option<&'static str> {
    use super::{binary::Binary as B, unary::Unary as U};
    Some(match op {
        Operation::Broadcast => "Broadcast",
        Operation::AsType => "AsType",
        Operation::Select => "Select",
        Operation::Unary(op) => match op {
            U::LogicalNot => "LogicalNot",
            U::Abs => "Abs",
            U::Sign => "Sign",
            U::Negative => "Negative",
            U::Square => "Square",
            U::Exp => "Exp",
            U::Log => "Log",
            U::Sigmoid => "Sigmoid",
            U::Tanh => "Tanh",
            U::Sqrt => "Sqrt",
            U::Rsqrt => "Rsqrt",
            U::Sin => "Sin",
            U::Cos => "Cos",
            U::Erf => "Erf",
        },
        Operation::Binary(op) => match op {
            B::Equal => "Equal",
            B::NotEqual => "NotEqual",
            B::GreaterEqual => "GreaterEqual",
            B::Greater => "Greater",
            B::Less => "Less",
            B::LessEqual => "LessEqual",
            B::LogicalAnd => "LogicalAnd",
            B::LogicalOr => "LogicalOr",
            B::BitwiseAnd => "BitwiseAnd",
            B::BitwiseOr => "BitwiseOr",
            B::BitwiseXor => "BitwiseXor",
            B::Remainder => "Remainder",
            B::Add => "Add",
            B::Subtract => "Subtract",
            B::Multiply => "Multiply",
            B::Divide => "Divide",
            B::Maximum => "Maximum",
            B::Minimum => "Minimum",
            B::LogAddExp => "LogAddExp",
            B::Power => "Power",
        },
        _ => return None,
    })
}

fn dtype(d: Dtype) -> &'static str {
    match d {
        Dtype::Bool => "bool",
        Dtype::U8 => "uint8_t",
        Dtype::U16 => "uint16_t",
        Dtype::U32 => "uint32_t",
        Dtype::U64 => "uint64_t",
        Dtype::I8 => "int8_t",
        Dtype::I16 => "int16_t",
        Dtype::I32 => "int32_t",
        Dtype::I64 => "int64_t",
        Dtype::F16 => "float16_t",
        Dtype::BF16 => "bfloat16_t",
        Dtype::F32 => "float",
        Dtype::F64 => "double",
        Dtype::C64 => "complex64_t",
    }
}

// ostream defaultfloat with digits10+1 (7 for F32/F16/BF16, 16 for F64).
fn float_constant(value: f64, precision: usize) -> String {
    if value.is_nan() {
        return "NAN".into();
    }
    if value.is_infinite() {
        return if value.is_sign_negative() {
            "-INFINITY"
        } else {
            "INFINITY"
        }
        .into();
    }
    let scientific = format!("{:.*e}", precision - 1, value);
    let (mantissa, exponent) = scientific.split_once('e').unwrap();
    let exponent: i32 = exponent.parse().unwrap();
    let trim = |s: String| {
        if s.contains('.') {
            s.trim_end_matches('0').trim_end_matches('.').to_owned()
        } else {
            s
        }
    };
    if exponent < -4 || exponent >= precision as i32 {
        format!("{}e{:+03}", trim(mantissa.to_owned()), exponent)
    } else {
        trim(format!(
            "{:.*}",
            (precision as i32 - 1 - exponent).max(0) as usize,
            value
        ))
    }
}

fn constant(a: &Array) -> Result<String, String> {
    a.wait()?;
    if !a.has_data() {
        return Err("compiled constant is unavailable".into());
    }
    let mut bits = [0u8; 8];
    a.buffer().read(a.offset(), &mut bits[..a.itemsize()])?;
    let u16v = u16::from_le_bytes(bits[..2].try_into().unwrap());
    let u32v = u32::from_le_bytes(bits[..4].try_into().unwrap());
    Ok(match a.layout().dtype() {
        Dtype::Bool => (bits[0] != 0).to_string(),
        Dtype::U8 => bits[0].to_string(),
        Dtype::I8 => (bits[0] as i8).to_string(),
        Dtype::U16 => u16v.to_string(),
        Dtype::I16 => (u16v as i16).to_string(),
        Dtype::U32 => u32v.to_string(),
        Dtype::I32 => (u32v as i32).to_string(),
        Dtype::U64 => u64::from_le_bytes(bits).to_string(),
        Dtype::I64 => i64::from_le_bytes(bits).to_string(),
        Dtype::BF16 => float_constant(f32::from_bits(u32::from(u16v) << 16) as f64, 7),
        Dtype::F16 => float_constant(super::super::half_to_f32(u16v) as f64, 7),
        Dtype::F32 => float_constant(f32::from_bits(u32v) as f64, 7),
        Dtype::F64 => float_constant(f64::from_le_bytes(bits), 16),
        Dtype::C64 => format!(
            "complex64_t({}, {})",
            float_constant(f32::from_bits(u32v) as f64, 7),
            float_constant(f32::from_le_bytes(bits[4..].try_into().unwrap()) as f64, 7)
        ),
    })
}

struct Names {
    slots: Vec<Option<String>>,
    count: usize,
}
impl Names {
    fn get(&mut self, slot: usize) -> String {
        self.slots[slot]
            .get_or_insert_with(|| {
                self.count += 1;
                let mut n = self.count;
                let mut letters = Vec::new();
                while n > 0 {
                    letters.push(b'A' + ((n - 1) % 26) as u8);
                    n = (n - 1) / 26;
                }
                letters.reverse();
                String::from_utf8(letters).unwrap()
            })
            .clone()
    }
}

impl Expression {
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) fn name(&self) -> String {
        std::iter::once("Compiled")
            .chain(self.statements.iter().map(|s| s.operator))
            .collect()
    }

    pub(super) fn new(
        inputs: &[Array],
        statements: Vec<Statement>,
        outputs: Vec<usize>,
        constants: &[bool],
    ) -> Result<Self, String> {
        if inputs.len() != constants.len() || statements.is_empty() || outputs.is_empty() {
            return Err("invalid compiled expression arity".into());
        }
        for (i, s) in statements.iter().enumerate() {
            if s.inputs.is_empty() || s.inputs.iter().any(|&j| j >= inputs.len() + i) {
                return Err("compiled expression is not topological".into());
            }
        }
        if outputs
            .iter()
            .any(|&o| o < inputs.len() || o >= inputs.len() + statements.len())
        {
            return Err("invalid compiled expression output".into());
        }
        let inputs = inputs
            .iter()
            .zip(constants)
            .map(|(a, &embed)| {
                Ok(Input {
                    dtype: a.layout().dtype(),
                    scalar: a.layout().shape().is_empty(),
                    constant: if embed { Some(constant(a)?) } else { None },
                })
            })
            .collect::<Result<Vec<_>, String>>()?;
        // Symbol renaming only; expression spelling and argument ordering
        // follow the original emitter. Hash includes embedded constant text.
        let hash = Sha256::digest(format!("{inputs:?}{statements:?}{outputs:?}"));
        let hash = hash.iter().map(|b| format!("{b:02x}")).collect::<String>();
        let prefix = format!("kernel_qwen_mtplx_compiled_{hash}");
        Ok(Self {
            inputs,
            statements,
            outputs,
            prefix,
        })
    }

    pub(super) fn evaluate(&self, inputs: &[Array], outputs: &[Array]) -> Result<(), String> {
        super::compile_custom(&format!("{}_contiguous", self.prefix), || self.source())?;
        compiled::evaluate_many(inputs, outputs, &self.prefix, |i| {
            self.inputs[i].constant.is_some()
        })
    }

    fn value_dtype(&self, slot: usize) -> Dtype {
        if slot < self.inputs.len() {
            self.inputs[slot].dtype
        } else {
            self.statements[slot - self.inputs.len()].dtype
        }
    }

    pub(super) fn source(&self) -> Result<String, String> {
        let mut source =
            include_str!("../../../../metal/mtplx-qsa-compiled-header.metal").to_owned();
        let work = (8 / self.value_dtype(self.outputs[0]).itemsize()).max(1);
        source.push_str(&self.kernel(true, 0, false, false, 1)?);
        if work > 1 {
            source.push_str(&self.kernel(true, 0, false, false, work)?);
        }
        source.push_str(&self.kernel(true, 0, false, true, work)?);
        for rank in 1..8 {
            source.push_str(&self.kernel(
                false,
                rank,
                false,
                false,
                if rank > 3 { 2 } else { 1 },
            )?);
            source.push_str(&self.kernel(
                false,
                rank,
                false,
                true,
                if rank > 3 { 4 } else { 1 },
            )?);
        }
        source.push_str(&self.kernel(false, 0, true, false, 2)?);
        source.push_str(&self.kernel(false, 0, true, true, 4)?);
        Ok(source)
    }

    fn kernel(
        &self,
        contiguous: bool,
        rank: usize,
        dynamic: bool,
        large: bool,
        work: usize,
    ) -> Result<String, String> {
        let suffix = if contiguous {
            format!("contiguous{}", if work > 1 && !large { "_n" } else { "" })
        } else {
            format!(
                "strided_{}",
                if dynamic {
                    "dynamic".to_owned()
                } else {
                    rank.to_string()
                }
            )
        };
        let name = format!(
            "{}_{}{}",
            self.prefix,
            suffix,
            if large { "_large" } else { "" }
        );
        let mut os = format!("[[host_name(\"{name}\")]]\n[[kernel]] void {name}(\n");
        let mut names = Names {
            slots: vec![None; self.inputs.len() + self.statements.len()],
            count: 0,
        };
        let mut count = 0;
        let mut indices = false;
        for (i, input) in self.inputs.iter().enumerate() {
            if input.constant.is_some() {
                continue;
            }
            if !input.scalar && !contiguous {
                indices = true;
            }
            writeln!(
                os,
                "    device const {}* {} [[buffer({count})]],",
                dtype(input.dtype),
                names.get(i)
            )
            .unwrap();
            count += 1;
        }
        let index_type = if large { "int64_t" } else { "uint" };
        if indices {
            writeln!(
                os,
                "    constant const int64_t* in_strides [[buffer({count})]],"
            )
            .unwrap();
            count += 1;
        }
        for &o in &self.outputs {
            writeln!(
                os,
                "    device {}* {} [[buffer({count})]],",
                dtype(self.value_dtype(o)),
                names.get(o)
            )
            .unwrap();
            count += 1;
        }
        if contiguous {
            writeln!(
                os,
                "    constant const {index_type}& size [[buffer({count})]],"
            )
            .unwrap();
        } else {
            writeln!(
                os,
                "    constant const int* output_shape [[buffer({count})]],"
            )
            .unwrap();
        }
        count += 1;
        if dynamic {
            writeln!(os, "    constant const int& ndim [[buffer({count})]],").unwrap();
            count += 1;
        }
        if count > 31 {
            return Err("compiled expression exceeds Metal's 31 argument buffers".into());
        }
        os.push_str(
            "    uint3 pos [[thread_position_in_grid]],\n    uint3 grid [[threads_per_grid]]) {\n",
        );
        writeln!(os, "  constexpr int N_ = {work};").unwrap();
        if contiguous && large {
            os.push_str("  int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));\n");
        } else if contiguous {
            os.push_str("  uint index = N_ * pos.x;\n");
        } else if work > 1 {
            writeln!(
                os,
                "  int xshape = output_shape[{}];",
                if dynamic {
                    "ndim - 1".to_owned()
                } else {
                    (rank - 1).to_string()
                }
            )
            .unwrap();
            writeln!(os, "  {index_type} index = N_ * pos.x + xshape * (pos.y + {index_type}(grid.y) * pos.z);").unwrap();
        } else {
            writeln!(
                os,
                "  {index_type} index = pos.x + grid.x * (pos.y + {index_type}(grid.y) * pos.z);"
            )
            .unwrap();
        }
        if work > 1 && contiguous {
            os.push_str("  for (int i = 0; i < N_ && index < size; ++i) {\n");
        }
        let mut nc = Vec::new();
        for (i, input) in self.inputs.iter().enumerate() {
            let n = names.get(i);
            if let Some(c) = &input.constant {
                writeln!(
                    os,
                    "  auto tmp_{n} = static_cast<{}>({c});",
                    dtype(input.dtype)
                )
                .unwrap();
            } else if input.scalar {
                writeln!(os, "  {} tmp_{n} = {n}[0];", dtype(input.dtype)).unwrap();
            } else if contiguous {
                writeln!(os, "  {} tmp_{n} = {n}[index];", dtype(input.dtype)).unwrap();
            } else {
                nc.push(i);
            }
        }
        for (i, &slot) in nc.iter().enumerate() {
            let n = names.get(slot);
            write!(os, "  {index_type} index_{n} = ").unwrap();
            match rank {
                1 => writeln!(os, "elem_to_loc_1<{index_type}>(pos.x, in_strides[{}]);", i * rank).unwrap(),
                2 => writeln!(os, "elem_to_loc_2<{index_type}>({{pos.x, pos.y}}, in_strides + {});", i * rank).unwrap(),
                3 => writeln!(os, "elem_to_loc_3<{index_type}>(pos, in_strides + {});", i * rank).unwrap(),
                _ if !dynamic => writeln!(os, "N_ * pos.x * {index_type}(in_strides[{}]) + pos.y * {index_type}(in_strides[{}]);", (i+1)*rank-1, (i+1)*rank-2).unwrap(),
                _ => writeln!(os, "N_ * pos.x * {index_type}(in_strides[ndim * {i} + ndim - 1]) + pos.y * {index_type}(in_strides[ndim * {i} + ndim - 2]);").unwrap(),
            }
        }
        if !nc.is_empty() && (rank > 3 || dynamic) {
            os.push_str("  uint zpos = pos.z;\n");
            writeln!(
                os,
                "  for (int d = {}; d >= 0; --d) {{",
                if dynamic {
                    "ndim - 3".to_owned()
                } else {
                    (rank - 3).to_string()
                }
            )
            .unwrap();
            os.push_str("    uint l = zpos % output_shape[d];\n");
            for (i, &slot) in nc.iter().enumerate() {
                let n = names.get(slot);
                if dynamic {
                    writeln!(
                        os,
                        "    index_{n} += l * {index_type}(in_strides[{i} * ndim + d]);"
                    )
                    .unwrap();
                } else {
                    writeln!(
                        os,
                        "    index_{n} += l * {index_type}(in_strides[{} + d]);",
                        i * rank
                    )
                    .unwrap();
                }
            }
            os.push_str("    zpos /= output_shape[d];\n  }\n");
        }
        if work > 1 && !contiguous {
            os.push_str("  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {\n");
        }
        for &slot in &nc {
            let n = names.get(slot);
            writeln!(
                os,
                "  {} tmp_{n} = {n}[index_{n}];",
                dtype(self.inputs[slot].dtype)
            )
            .unwrap();
        }
        for (i, s) in self.statements.iter().enumerate() {
            let n = names.get(self.inputs.len() + i);
            write!(os, "  {} tmp_{n} = ", dtype(s.dtype)).unwrap();
            if matches!(s.operator, "Broadcast" | "AsType") {
                writeln!(
                    os,
                    "cast_to<{}>(tmp_{});",
                    dtype(s.dtype),
                    names.get(s.inputs[0])
                )
                .unwrap();
            } else {
                let args = s
                    .inputs
                    .iter()
                    .map(|&s| format!("tmp_{}", names.get(s)))
                    .collect::<Vec<_>>()
                    .join(", ");
                writeln!(os, "{}()({args});", s.operator).unwrap();
            }
        }
        for &slot in &self.outputs {
            let n = names.get(slot);
            writeln!(os, "  {n}[index] = tmp_{n};").unwrap();
        }
        if work > 1 {
            for (i, &slot) in nc.iter().enumerate() {
                let n = names.get(slot);
                if dynamic {
                    writeln!(os, "  index_{n} += in_strides[{i} * ndim + ndim - 1];").unwrap();
                } else {
                    writeln!(os, "  index_{n} += in_strides[{}];", i * rank + rank - 1).unwrap();
                }
            }
            os.push_str("  index++;\n  }\n");
        }
        os.push_str("}\n");
        Ok(os)
    }
}

#[test]
fn mtplx_compiled_constant_spelling() {
    for (value, text) in [
        (0., "0"),
        (-0., "-0"),
        (0.1, "0.1"),
        (1e-4, "0.0001"),
        (1e-5, "1e-05"),
        (1e7, "1e+07"),
        (1234567., "1234567"),
        (12345678., "1.234568e+07"),
        (f64::NAN, "NAN"),
        (f64::INFINITY, "INFINITY"),
        (f64::NEG_INFINITY, "-INFINITY"),
    ] {
        assert_eq!(float_constant(value, 7), text);
    }
    assert_eq!(float_constant(1.2345678901234567, 16), "1.234567890123457");
}

#[test]
fn mtplx_compiled_emitter_matches_all_pinned_silu_bodies() {
    // The pinned shapeless SiLU has two explicit Broadcast primitives;
    // compare the emitter, independently of shape-fixed trace selection.
    let input = Array::unallocated(&[8], Dtype::BF16).unwrap();
    let s = |operator, inputs| Statement {
        dtype: Dtype::BF16,
        operator,
        inputs,
    };
    let expression = Expression::new(
        &[input],
        vec![
            s("Sigmoid", vec![0]),
            s("Broadcast", vec![0, 1]),
            s("Broadcast", vec![1, 0]),
            s("Multiply", vec![2, 3]),
        ],
        vec![4],
        &[false],
    )
    .unwrap();
    let source = expression.source().unwrap();
    let reference = include_str!("../../../../metal/mtplx_qwen.metal");
    let mut compared = 0;
    for section in reference
        .split("[[host_name(\"kernel_qwen_mtplx_silu_bf16_")
        .skip(1)
    {
        let suffix = section.split_once('"').unwrap().0;
        let original = format!("[[host_name(\"kernel_qwen_mtplx_silu_bf16_{section}");
        let original = &original[..original.find("\n}\n").unwrap() + 3];
        let symbol = original
            .split_once("void ")
            .unwrap()
            .1
            .split_once('(')
            .unwrap()
            .0;
        let expected = original
            .replace(symbol, &format!("{}_{suffix}", expression.prefix))
            .replace(
                &format!("kernel_qwen_mtplx_silu_bf16_{suffix}"),
                &format!("{}_{suffix}", expression.prefix),
            );
        let marker = format!("[[host_name(\"{}_{suffix}\")]]", expression.prefix);
        let start = source.find(&marker).unwrap();
        let generated = &source[start..];
        let generated = &generated[..generated.find("\n}\n").unwrap() + 3];
        assert_eq!(generated, expected, "{suffix}");
        compared += 1;
    }
    assert_eq!(compared, 19);
}

#[test]
fn mtplx_compile_fusion_preserves_parent_and_output_boundaries() {
    use super::{
        binary::{Binary, binary},
        compiled::Replay,
        ops,
        stream::{Device, Stream},
        unary::{Unary, unary},
    };
    let gpu = Stream::new(0, Device::Gpu);
    let other = Stream::new(1, Device::Gpu);
    let x = Array::unallocated(&[4], Dtype::F32).unwrap();
    let scalar = Array::unallocated(&[], Dtype::F32).unwrap();
    let inputs = [x, scalar];
    let results = Replay::trace_compiled(&inputs, |a| {
        let square = unary(&a[0], Unary::Square, gpu)?;
        let sum = binary(&square, &a[1], Binary::Add, gpu)?;
        Ok(vec![square, sum.clone(), sum])
    })
    .unwrap()
    .run(&inputs)
    .unwrap();
    assert!(matches!(
        &*results[0].operation().unwrap(),
        Operation::Compiled(_)
    ));
    assert_eq!(results[0].outputs().len(), 2);
    assert_eq!(results[1].id(), results[2].id());
    assert!(std::rc::Rc::ptr_eq(
        &results[0].operation().unwrap(),
        &results[1].operation().unwrap()
    ));
    // Mixed-parent Broadcast is split; the outside consumer retains its view.
    let results = Replay::trace_compiled(&inputs, |a| {
        let broadcast = ops::broadcast_to(&a[1], &[4], gpu)?;
        let sum = binary(&a[0], &broadcast, Binary::Add, gpu)?;
        Ok(vec![ops::reshape(&broadcast, &[2, 2], gpu)?, sum])
    })
    .unwrap()
    .run(&inputs)
    .unwrap();
    assert!(matches!(
        &*results[0].inputs()[0].operation().unwrap(),
        Operation::Broadcast
    ));
    assert!(results[1].inputs().iter().any(|a| a.id() == inputs[1].id()));
    // A non-Broadcast intermediate with outside parents cannot be split.
    let results = Replay::trace_compiled(&inputs, |a| {
        let square = unary(&a[0], Unary::Square, gpu)?;
        Ok(vec![
            ops::reshape(&square, &[2, 2], gpu)?,
            binary(&square, &a[1], Binary::Add, gpu)?,
        ])
    })
    .unwrap()
    .run(&inputs)
    .unwrap();
    assert_eq!(results[0].inputs()[0].id(), results[1].inputs()[0].id());
    // Stream mismatch is a fusion boundary even though CSE ignores streams.
    let results = Replay::trace_compiled(&inputs, |a| {
        let square = unary(&a[0], Unary::Square, gpu)?;
        let negative = unary(&square, Unary::Negative, other)?;
        Ok(vec![binary(&negative, &a[1], Binary::Add, other)?])
    })
    .unwrap()
    .run(&inputs)
    .unwrap();
    assert_eq!(results[0].inputs()[0].stream(), Some(gpu));
    // A 25-node chain splits at the original depth=11 boundary: 11+11+3.
    let results = Replay::trace_compiled(&inputs, |a| {
        let mut x = a[0].clone();
        for _ in 0..25 {
            x = unary(&x, Unary::Negative, gpu)?;
        }
        Ok(vec![x])
    })
    .unwrap()
    .run(&inputs)
    .unwrap();
    let mut x = results[0].clone();
    let mut lengths = Vec::new();
    while let Some(op) = x.operation() {
        let Operation::Compiled(expression) = &*op else {
            panic!("unfused chain");
        };
        lengths.push(expression.statements.len());
        let next = x.inputs()[0].clone();
        x = next;
    }
    assert_eq!(lengths, [11, 11, 3]);
    // Broadcast-only outputs remain views rather than materialized kernels.
    let result = Replay::trace_compiled(&inputs, |a| {
        Ok(vec![ops::broadcast_to(&a[1], &[4096], gpu)?])
    })
    .unwrap()
    .run(&inputs)
    .unwrap();
    assert!(matches!(
        &*result[0].operation().unwrap(),
        Operation::Broadcast
    ));
}

#[test]
#[ignore = "requires Metal and external supervisor; generated fusion vs pinned SiLU plus multi-output/scalar dispatch"]
fn mtplx_compiled_fusion_gpu_matches_reference() {
    use super::{
        super::{
            gpu::Context,
            qwen_mtplx_tests::{capture_dispatches, pattern},
        },
        allocator,
        compiled::Replay,
        configure_sources, gdn, mlp, ops,
        stream::{Device, Streams},
    };
    use super::{
        binary::{Binary, binary},
        unary::{Unary, unary},
    };
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
    let gpu = streams.default_stream(Device::Gpu).unwrap();
    for rank in [1, 4, 8] {
        for reverse in [false, true] {
            eprintln!("compiled SiLU rank={rank} reverse={reverse}");
            let shape = vec![2; rank];
            let input = Array::new(&shape, Dtype::BF16, pattern(1 << rank, 37)).unwrap();
            let axes = (0..rank as i32)
                .step_by(2)
                .chain((1..rank as i32).step_by(2))
                .collect::<Vec<_>>();
            let input = ops::transpose(&input, &axes, gpu).unwrap();
            let input = if reverse {
                ops::slice(
                    &input,
                    &vec![1; rank],
                    &vec![-3; rank],
                    &vec![-1; rank],
                    gpu,
                )
                .unwrap()
            } else {
                input
            };
            ops::evaluate(&streams, std::slice::from_ref(&input), gpu, false).unwrap();
            let expected = mlp::silu(&input, gpu).unwrap();
            let graph = Replay::trace_compiled(std::slice::from_ref(&input), |a| {
                Ok(vec![mlp::silu(&a[0], gpu)?])
            })
            .unwrap();
            let actual = graph.run(std::slice::from_ref(&input)).unwrap().remove(0);
            let (_, calls) = capture_dispatches(|| {
                ops::evaluate(&streams, std::slice::from_ref(&actual), gpu, false).unwrap()
            });
            assert_eq!(calls.len(), 1);
            assert!(calls[0].0.starts_with("kernel_qwen_mtplx_compiled_"));
            let actual = ops::contiguous(&actual, false, gpu).unwrap();
            let expected = ops::contiguous(&expected, false, gpu).unwrap();
            ops::evaluate(&streams, &[actual.clone(), expected.clone()], gpu, false).unwrap();
            let mut a = vec![0; actual.layout().nbytes()];
            let mut b = vec![0; expected.layout().nbytes()];
            actual.buffer().read(actual.offset(), &mut a).unwrap();
            expected.buffer().read(expected.offset(), &mut b).unwrap();
            assert_eq!(a, b);
        }
    }
    for rows in [1, 4, 17] {
        eprintln!("compiled nested compute_g rows={rows}");
        let args = [
            Array::new(&[48], Dtype::BF16, pattern(48, 62)).unwrap(),
            Array::new(&[1, rows, 48], Dtype::BF16, pattern(rows as u32 * 48, 60)).unwrap(),
            Array::new(&[48], Dtype::BF16, pattern(48, 63)).unwrap(),
        ];
        let expected = gdn::compute_g(&args[0], &args[1], &args[2], gpu).unwrap();
        let graph = Replay::trace_compiled(&args, |a| {
            Ok(vec![gdn::compute_g(&a[0], &a[1], &a[2], gpu)?])
        })
        .unwrap();
        let actual = graph.run(&args).unwrap().remove(0);
        let (_, calls) = capture_dispatches(|| {
            ops::evaluate(&streams, std::slice::from_ref(&actual), gpu, false).unwrap()
        });
        assert_eq!(calls.len(), 1, "{calls:?}");
        assert!(calls[0].0.starts_with("kernel_qwen_mtplx_compiled_"));
        ops::evaluate(&streams, std::slice::from_ref(&expected), gpu, false).unwrap();
        let mut a = vec![0; actual.layout().nbytes()];
        let mut b = vec![0; expected.layout().nbytes()];
        actual.buffer().read(actual.offset(), &mut a).unwrap();
        expected.buffer().read(expected.offset(), &mut b).unwrap();
        assert_eq!(a, b, "nested compute_g rows={rows}");
    }
    for n in [8, 65536] {
        for embedded in [false, true] {
            eprintln!("compiled multiple outputs n={n} embedded={embedded}");
            let data = (0..n)
                .flat_map(|i| ((i % 8) as f32).to_le_bytes())
                .collect::<Vec<_>>();
            let input = Array::new(&[n], Dtype::F32, super::scalar_buffer(&data).unwrap()).unwrap();
            let scalar = Array::new(
                &[],
                Dtype::F32,
                super::scalar_buffer(&2f32.to_le_bytes()).unwrap(),
            )
            .unwrap();
            let args = if embedded {
                vec![input]
            } else {
                vec![input, scalar.clone()]
            };
            let graph = Replay::trace_compiled(&args, |a| {
                let square = unary(&a[0], Unary::Square, gpu)?;
                let sum = binary(
                    &square,
                    if embedded { &scalar } else { &a[1] },
                    Binary::Add,
                    gpu,
                )?;
                Ok(vec![square, sum])
            })
            .unwrap();
            let roots = graph.run(&args).unwrap();
            // The pinned evaluator selects its pipeline before allocating
            // outputs (metal/compiled.cpp:384,419). Fresh ArrayDesc data_size
            // is zero, even when the logical shape exceeds 65536 elements.
            assert!(roots.iter().all(|a| a.data_size() == 0));
            let (_, calls) =
                capture_dispatches(|| ops::evaluate(&streams, &roots, gpu, false).unwrap());
            assert_eq!(calls.len(), 1, "{calls:?}");
            assert!(calls[0].0.ends_with("_contiguous"), "{calls:?}");
            assert_eq!(calls[0].1, [n as u32, 1, 1]);
            for (j, root) in roots.iter().enumerate() {
                let mut bytes = vec![0; root.layout().nbytes()];
                root.buffer().read(root.offset(), &mut bytes).unwrap();
                for (i, b) in bytes.chunks_exact(4).enumerate() {
                    assert_eq!(
                        f32::from_le_bytes(b.try_into().unwrap()),
                        ((i % 8) * (i % 8)) as f32 + (j * 2) as f32
                    );
                }
            }
        }
    }
}