qualia_core_db/wgsl_forge/emit/

wgsl.rs

use std::fmt::Write;

use super::GeneratedShader;
use crate::wgsl_forge::{ForgeError, KernelSpec, Op, Schedule, FORGE_SCHEMA_VERSION};

pub fn emit_wgsl(kernel: &KernelSpec, schedule: Schedule) -> Result<GeneratedShader, ForgeError> {
    kernel.validate()?;
    let semantic_hash = kernel.semantic_hash()?;
    let mut source = String::with_capacity(2_048);
    writeln!(
        source,
        "// Generated by Qualia WGSL Forge schema {FORGE_SCHEMA_VERSION}."
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    writeln!(
        source,
        "// Kernel: {}@{}",
        kernel.id, kernel.semantic_version
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    writeln!(source, "// Semantic hash: {semantic_hash}")
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    writeln!(
        source,
        "// Schedule: workgroup={}, items={}, vector={}",
        schedule.workgroup_size, schedule.items_per_invocation, schedule.vector_width
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;

    emit_kernel_body(&mut source, kernel, schedule)?;

    let source_hash = blake3::hash(source.as_bytes()).to_hex().to_string();
    Ok(GeneratedShader {
        kernel_id: kernel.id.clone(),
        semantic_hash,
        source_hash,
        schedule,
        source,
    })
}

fn emit_kernel_body(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    if kernel.id == "topk" {
        return emit_topk_wgsl(source, kernel, schedule);
    }

    if kernel.id == "fused-ffn" {
        return emit_ffn_wgsl(source, kernel, schedule);
    }

    if kernel.id == "p64-project" {
        return emit_p64_wgsl(source, kernel, schedule);
    }

    if kernel.id == "ternary-gemv" {
        return emit_ternary_gemv_wgsl(source, kernel, schedule);
    }

    // gemm / gemv / fft are lowered through the compute-graph IR: the KernelSpec becomes a
    // one-node ComputeGraph (`to_graph`) and `lower_graph` walks it, dispatching the node
    // to the WgslDelegateLowerer. In Phase 1 the leaf emission DELEGATES to the proven
    // emit_*_wgsl functions, so the generated bytes are byte-identical by construction (the
    // certify-cache source_hash is unchanged). This is the seam that later phases use to
    // replace the per-kernel.id string emitters with native graph templates and to lower
    // the same graph to other backends. See docs/plans/dag-ir-forge.md.
    if matches!(kernel.id.as_str(), "gemm" | "gemv" | "fft") {
        let graph = kernel.to_graph()?;
        let mut lowerer = WgslDelegateLowerer {
            source,
            kernel,
            schedule,
        };
        return crate::wgsl_forge::ir::graph::lower_graph(&graph, &mut lowerer);
    }

    if kernel.id == "ray-probe" {
        // The ray-query enable extension must precede any other item.
        writeln!(source, "enable wgpu_ray_query;")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    if kernel.id == "affine-f32" {
        writeln!(
            source,
            r#"
struct AffineParams {{
    length: u32,
    scale: f32,
    bias: f32,
    _pad: u32,
}}"#
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    writeln!(source, "").map_err(|error| ForgeError::Emission(error.to_string()))?;
    for buffer in &kernel.buffers {
        // Acceleration structures bind without an address space.
        if buffer.element == crate::wgsl_forge::ir::BufferElement::AccelerationStructure {
            writeln!(
                source,
                "@group({}) @binding({}) var {}: acceleration_structure;",
                buffer.group, buffer.binding, buffer.name
            )
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
            continue;
        }
        let access = match buffer.access {
            crate::wgsl_forge::ir::BufferAccess::StorageRead => "storage, read",
            crate::wgsl_forge::ir::BufferAccess::StorageReadWrite => "storage, read_write",
            crate::wgsl_forge::ir::BufferAccess::Uniform => "uniform",
        };
        let type_decl = match buffer.element {
            crate::wgsl_forge::ir::BufferElement::Scalar(
                crate::wgsl_forge::ir::ScalarType::F32,
            ) => "array<f32>",
            crate::wgsl_forge::ir::BufferElement::Scalar(
                crate::wgsl_forge::ir::ScalarType::U32,
            ) => "array<u32>",
            crate::wgsl_forge::ir::BufferElement::Scalar(
                crate::wgsl_forge::ir::ScalarType::I32,
            ) => "array<i32>",
            crate::wgsl_forge::ir::BufferElement::Scalar(
                crate::wgsl_forge::ir::ScalarType::U64Words,
            ) => "array<vec2<u32>>",
            crate::wgsl_forge::ir::BufferElement::AffineParams => "AffineParams",
            crate::wgsl_forge::ir::BufferElement::P64Words64 => "array<P64Words64>",
            crate::wgsl_forge::ir::BufferElement::AccelerationStructure => {
                unreachable!("handled above")
            }
        };
        writeln!(
            source,
            "@group({}) @binding({}) var<{}> {}: {};",
            buffer.group, buffer.binding, access, buffer.name, type_decl
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    writeln!(
        source,
        r#"
const ITEMS_PER_INVOCATION: u32 = {}u;
const VECTOR_WIDTH: u32 = {}u;

@compute @workgroup_size({})
fn {}(@builtin(global_invocation_id) gid: vec3<u32>) {{"#,
        schedule.items_per_invocation,
        schedule.vector_width,
        schedule.workgroup_size,
        kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;

    if kernel.id == "ray-probe" {
        // One ray per invocation: load the 8-float ray, intersect, store the hit.
        writeln!(source, "    let i = gid.x;")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(source, "    if (i >= arrayLength(&hits)) {{ return; }}")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(source, "    let base = i * 8u;")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(
            source,
            "    let origin = vec3<f32>(rays[base], rays[base + 1u], rays[base + 2u]);"
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(
            source,
            "    let direction = vec3<f32>(rays[base + 3u], rays[base + 4u], rays[base + 5u]);"
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(source, "    let t_min = rays[base + 6u];")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(source, "    let t_max = rays[base + 7u];")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
        emit_ops(source, &kernel.ops, "    ")?;
        writeln!(source, "    hits[i] = hit_t;")
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(source, "}}").map_err(|error| ForgeError::Emission(error.to_string()))?;
    } else if kernel.id == "affine-f32" {
        // Keeping the optimized vectorized wrapper for affine-f32 specifically
        writeln!(
            source,
            "    for (var item: u32 = 0u; item < ITEMS_PER_INVOCATION; item = item + 1u) {{"
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(
            source,
            "        let global_id = (gid.x * ITEMS_PER_INVOCATION + item) * VECTOR_WIDTH;"
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;

        if schedule.vector_width == 1 {
            writeln!(source, "        if (global_id < params.length) {{")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
            emit_ops(source, &kernel.ops, "            ")?;
            writeln!(source, "        }}")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
        } else {
            writeln!(
                source,
                "        if (global_id + {}u < params.length) {{",
                schedule.vector_width - 1
            )
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
            let constructor = (0..schedule.vector_width)
                .map(|index| format!("input[global_id + {index}u]"))
                .collect::<Vec<_>>()
                .join(", ");
            writeln!(
                source,
                "            let value = vec{}<f32>({});",
                schedule.vector_width, constructor
            )
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
            writeln!(source, "            let transformed = value * vec{}<f32>(params.scale) + vec{}<f32>(params.bias);", schedule.vector_width, schedule.vector_width).map_err(|error| ForgeError::Emission(error.to_string()))?;
            for index in 0..schedule.vector_width {
                let components = ["x", "y", "z", "w"];
                writeln!(
                    source,
                    "            output[global_id + {index}u] = transformed.{};",
                    components[index as usize]
                )
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            writeln!(source, "        }} else {{")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
            writeln!(source, "            for (var component: u32 = 0u; component < VECTOR_WIDTH; component = component + 1u) {{").map_err(|error| ForgeError::Emission(error.to_string()))?;
            writeln!(
                source,
                "                let global_id = global_id + component;"
            )
            .map_err(|error| ForgeError::Emission(error.to_string()))?;
            writeln!(source, "                if (global_id < params.length) {{")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
            emit_ops(source, &kernel.ops, "                    ")?;
            writeln!(source, "                }}")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
            writeln!(source, "            }}")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
            writeln!(source, "        }}")
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
        }
        writeln!(source, "    }}\n}}").map_err(|error| ForgeError::Emission(error.to_string()))?;
    } else {
        // Generic loop for fused-ffn and other kernels
        writeln!(
            source,
            "    for (var item: u32 = 0u; item < ITEMS_PER_INVOCATION; item = item + 1u) {{"
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
        writeln!(
            source,
            "        let global_id = gid.x * ITEMS_PER_INVOCATION + item;"
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
        emit_ops(source, &kernel.ops, "        ")?;
        writeln!(source, "    }}\n}}").map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    Ok(())
}

/// Emits the top-k reduction kernel: one workgroup per block, the `k` largest
/// values of each block written to `output` in descending order.
///
/// This is the first kernel to exercise workgroup-shared memory and barrier
/// synchronisation. The shared arrays come from the IR (`kernel.shared_memory`)
/// and `Op::Barrier` lowers to `workgroupBarrier()`; the reduction control flow
/// is specialised here because it is not expressible in the scalar op set.
fn emit_topk_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;

    writeln!(
        source,
        r#"
struct TopKParams {{
    length: u32,
    k: u32,
    block_size: u32,
    _pad: u32,
}}
"#
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;

    for buffer in &kernel.buffers {
        let access = match buffer.access {
            crate::wgsl_forge::ir::BufferAccess::StorageRead => "storage, read",
            crate::wgsl_forge::ir::BufferAccess::StorageReadWrite => "storage, read_write",
            crate::wgsl_forge::ir::BufferAccess::Uniform => "uniform",
        };
        let type_decl = match buffer.access {
            crate::wgsl_forge::ir::BufferAccess::Uniform => "TopKParams",
            _ => "array<f32>",
        };
        writeln!(
            source,
            "@group({}) @binding({}) var<{}> {}: {};",
            buffer.group, buffer.binding, access, buffer.name, type_decl
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    writeln!(source, "").map_err(|error| ForgeError::Emission(error.to_string()))?;
    for shared in &kernel.shared_memory {
        writeln!(
            source,
            "var<workgroup> {}: array<{}, {}>;",
            shared.name,
            shared.element.wgsl_name(),
            shared.length.resolve(wg)
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    writeln!(
        source,
        r#"
@compute @workgroup_size({wg})
fn {entry}(
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(workgroup_id) wid: vec3<u32>,
) {{
    let tid = lid.x;
    let block = wid.x;
    let base = block * {wg}u;
    let gidx = base + tid;
    // Sentinel = f32::MIN (0xff7fffff); threads past the input load it so they
    // never win, and selected slots are reset to it between extractions.
    let sentinel = bitcast<f32>(0xff7fffffu);
    var v = sentinel;
    if (gidx < params.length) {{
        v = input[gidx];
    }}
    s_val[tid] = v;
    s_idx[tid] = tid;
    workgroupBarrier();

    for (var i: u32 = 0u; i < params.k; i = i + 1u) {{
        r_val[tid] = s_val[tid];
        r_idx[tid] = s_idx[tid];
        workgroupBarrier();
        // Tree arg-max reduction over the working copy.
        for (var stride: u32 = {wg}u / 2u; stride > 0u; stride = stride / 2u) {{
            if (tid < stride) {{
                if (r_val[tid + stride] > r_val[tid]) {{
                    r_val[tid] = r_val[tid + stride];
                    r_idx[tid] = r_idx[tid + stride];
                }}
            }}
            workgroupBarrier();
        }}
        if (tid == 0u) {{
            output[block * params.k + i] = r_val[0];
            s_val[r_idx[0]] = sentinel;
        }}
        workgroupBarrier();
    }}
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;

    Ok(())
}

/// Fused feed-forward network: one invocation per output element computes
/// `out[o] = sum_h w2[o,h] * gelu(sum_i w1[h,i] * input[i])`. Self-contained
/// (no shared memory); dimensions come from the uniform params block.
fn emit_ffn_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;
    writeln!(
        source,
        r#"
struct FfnParams {{
    input_size: u32,
    hidden_size: u32,
    output_size: u32,
    _pad: u32,
}}

@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read> w1: array<f32>;
@group(0) @binding(2) var<storage, read> w2: array<f32>;
@group(0) @binding(3) var<storage, read_write> output: array<f32>;
@group(0) @binding(4) var<uniform> params: FfnParams;

@compute @workgroup_size({wg})
fn {entry}(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let o = gid.x;
    if (o >= params.output_size) {{ return; }}
    var acc = 0.0;
    for (var h: u32 = 0u; h < params.hidden_size; h = h + 1u) {{
        var hv = 0.0;
        let w1_row = h * params.input_size;
        for (var i: u32 = 0u; i < params.input_size; i = i + 1u) {{
            hv = hv + w1[w1_row + i] * input[i];
        }}
        // GELU(hv)
        let g = 0.5 * hv * (1.0 + tanh(0.7978845608 * (hv + 0.044715 * hv * hv * hv)));
        acc = acc + w2[o * params.hidden_size + h] * g;
    }}
    output[o] = acc;
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    Ok(())
}

/// P64 descriptor projection: one invocation per record projects its 16 packed
/// u32 words onto the 16-element weight vector.
fn emit_p64_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;
    writeln!(
        source,
        r#"
struct P64Words64 {{
    lanes: array<vec4<u32>, 4>,
}}

@group(0) @binding(0) var<storage, read> input: array<P64Words64>;
@group(0) @binding(1) var<storage, read> weights: array<f32>;
@group(0) @binding(2) var<storage, read_write> output: array<f32>;

@compute @workgroup_size({wg})
fn {entry}(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let r = gid.x;
    if (r >= arrayLength(&output)) {{ return; }}
    let rec = input[r];
    var acc = 0.0;
    for (var w: u32 = 0u; w < 16u; w = w + 1u) {{
        let word = rec.lanes[w / 4u][w % 4u];
        acc = acc + weights[w] * f32(word);
    }}
    output[r] = acc;
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    Ok(())
}

/// BitNet-style ternary GEMV with on-the-fly dequant: one invocation per output
/// row `o` computes `out[o] = scale[o] * sum_i ternary(w[o,i]) * x[i]`. Weights
/// are 2-bit-packed ternary codes (16 codes per u32, low-to-high lanes; code
/// `0->0.0, 1->+1.0, 2->-1.0, 3->0.0`), `ceil(K/16)` words per row laid out
/// contiguously. Self-contained (no shared memory); dimensions come from the
/// uniform params block (m, k, k_words, _pad).
fn emit_ternary_gemv_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;
    writeln!(
        source,
        r#"
struct TernaryGemvParams {{
    m: u32,
    k: u32,
    k_words: u32,
    _pad: u32,
}}

@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w_packed: array<u32>;
@group(0) @binding(2) var<storage, read> scale: array<f32>;
@group(0) @binding(3) var<storage, read_write> output: array<f32>;
@group(0) @binding(4) var<uniform> params: TernaryGemvParams;

@compute @workgroup_size({wg})
fn {entry}(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let o = gid.x;
    if (o >= params.m) {{ return; }}
    var acc = 0.0;
    let row_base = o * params.k_words;
    for (var word_idx: u32 = 0u; word_idx < params.k_words; word_idx = word_idx + 1u) {{
        let word = w_packed[row_base + word_idx];
        let lane_base = word_idx * 16u;
        for (var lane: u32 = 0u; lane < 16u; lane = lane + 1u) {{
            let i = lane_base + lane;
            if (i >= params.k) {{ break; }}
            // Extract the 2-bit ternary code for this lane (low-to-high).
            let code = (word >> (lane * 2u)) & 3u;
            // Map code -> value: 0->0.0, 1->+1.0, 2->-1.0, 3->0.0 (unused).
            var tern = 0.0;
            if (code == 1u) {{
                tern = 1.0;
            }} else if (code == 2u) {{
                tern = -1.0;
            }}
            acc = acc + tern * x[i];
        }}
    }}
    output[o] = scale[o] * acc;
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    Ok(())
}

/// General dense row-major GEMM, all f32: one invocation per output element
/// `o = i*N + j` computes `C[i][j] = sum_k A[i*K+k] * B[k*N+j]`. Self-contained
/// (no shared memory); dimensions come from the uniform params block (m, n, k, _pad).
/// The K dimension is `params.k`; the inner accumulation index is `kk` so it never
/// shadows the dimension name.
/// Phase-1 WGSL lowerer for the compute-graph IR. Its `matmul`/`gemv`/`fft` methods
/// **delegate** to the existing, certified `emit_*_wgsl` functions, so routing a
/// gemm/gemv/fft `KernelSpec` through `to_graph` → `lower_graph` produces byte-identical
/// WGSL. Native graph-template nodes (Reduce/Broadcast, …) use [`WgslGraphLowerer`]
/// instead; unimplemented op-classes inherit the trait's default `Err` (never a silent
/// no-op).
struct WgslDelegateLowerer<'a> {
    source: &'a mut String,
    kernel: &'a KernelSpec,
    schedule: Schedule,
}

impl crate::wgsl_forge::ir::graph::Lowerer for WgslDelegateLowerer<'_> {
    fn matmul(
        &mut self,
        _node: &crate::wgsl_forge::ir::graph::GraphNode,
    ) -> Result<(), ForgeError> {
        emit_gemm_wgsl(self.source, self.kernel, self.schedule)
    }
    fn gemv(&mut self, _node: &crate::wgsl_forge::ir::graph::GraphNode) -> Result<(), ForgeError> {
        emit_gemv_wgsl(self.source, self.kernel, self.schedule)
    }
    fn fft(&mut self, _node: &crate::wgsl_forge::ir::graph::GraphNode) -> Result<(), ForgeError> {
        emit_fft_wgsl(self.source, self.kernel, self.schedule)
    }
}

/// Phase-2 WGSL lowerer for **native** compute-graph nodes — those with no backing
/// `KernelSpec` (there is no legacy standalone reduce/broadcast kernel to delegate to).
/// Each method emits the node's WGSL template directly from its op payload + per-node
/// schedule. Op-classes not yet built inherit the trait default (`Err`).
struct WgslGraphLowerer<'a> {
    source: &'a mut String,
}

impl crate::wgsl_forge::ir::graph::Lowerer for WgslGraphLowerer<'_> {
    fn reduce(&mut self, node: &crate::wgsl_forge::ir::graph::GraphNode) -> Result<(), ForgeError> {
        use crate::wgsl_forge::ir::graph::OpNode;
        if let OpNode::Reduce { op, .. } = node.op {
            self.source
                .push_str(&crate::wgsl_forge::graph_ops::reduce::reduce_wgsl(
                    op,
                    node.sched.workgroup_size,
                ));
            Ok(())
        } else {
            Err(ForgeError::Emission(
                "WgslGraphLowerer::reduce received a non-Reduce node".to_string(),
            ))
        }
    }

    fn broadcast(
        &mut self,
        node: &crate::wgsl_forge::ir::graph::GraphNode,
    ) -> Result<(), ForgeError> {
        self.source
            .push_str(&crate::wgsl_forge::graph_ops::broadcast::broadcast_wgsl(
                node.sched.workgroup_size,
            ));
        Ok(())
    }

    fn elementwise(
        &mut self,
        node: &crate::wgsl_forge::ir::graph::GraphNode,
    ) -> Result<(), ForgeError> {
        use crate::wgsl_forge::ir::graph::OpNode;
        if let OpNode::Elementwise { f } = node.op {
            let src = crate::wgsl_forge::graph_ops::elementwise::elementwise_wgsl(
                f,
                node.sched.workgroup_size,
            )?;
            self.source.push_str(&src);
            Ok(())
        } else {
            Err(ForgeError::Emission(
                "WgslGraphLowerer::elementwise received a non-Elementwise node".to_string(),
            ))
        }
    }
}

/// Emit a complete WGSL module for a **pure compute-graph** (no backing `KernelSpec`) —
/// the native-node path (Reduce/Broadcast in Phase 2). A single-node graph produces one
/// kernel module; multi-node graph emission (per-node functions + a driver) is Phase 4.
pub fn emit_graph_wgsl(
    graph: &crate::wgsl_forge::ir::graph::ComputeGraph,
    schedule: Schedule,
) -> Result<GeneratedShader, ForgeError> {
    let mut source = String::with_capacity(1_024);
    writeln!(
        source,
        "// Generated by Qualia WGSL Forge schema {FORGE_SCHEMA_VERSION} (compute-graph)."
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    let mut lowerer = WgslGraphLowerer {
        source: &mut source,
    };
    crate::wgsl_forge::ir::graph::lower_graph(graph, &mut lowerer)?;
    let source_hash = blake3::hash(source.as_bytes()).to_hex().to_string();
    Ok(GeneratedShader {
        kernel_id: "graph".to_string(),
        semantic_hash: source_hash.clone(),
        source_hash,
        schedule,
        source,
    })
}

fn emit_gemm_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;
    writeln!(
        source,
        r#"
struct GemmParams {{
    m: u32,
    n: u32,
    k: u32,
    _pad: u32,
}}

@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;
@group(0) @binding(3) var<uniform> params: GemmParams;

@compute @workgroup_size({wg})
fn {entry}(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let o = gid.x;
    if (o >= params.m * params.n) {{ return; }}
    let row = o / params.n;
    let col = o % params.n;
    var acc = 0.0;
    let a_row = row * params.k;
    for (var kk: u32 = 0u; kk < params.k; kk = kk + 1u) {{
        acc = acc + a[a_row + kk] * b[kk * params.n + col];
    }}
    c[o] = acc;
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    Ok(())
}

/// Dense row-major matrix-vector product, all f32: one invocation per output ROW
/// `i` computes `y[i] = sum_j A[i*N+j] * x[j]`. Self-contained (no shared memory);
/// dimensions come from the uniform params block (m, n, _pad0, _pad1). The inner
/// accumulation index is `j`.
fn emit_gemv_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;
    writeln!(
        source,
        r#"
struct GemvParams {{
    m: u32,
    n: u32,
    _pad0: u32,
    _pad1: u32,
}}

@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> x: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> params: GemvParams;

@compute @workgroup_size({wg})
fn {entry}(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let i = gid.x;
    if (i >= params.m) {{ return; }}
    var acc = 0.0;
    let a_row = i * params.n;
    for (var j: u32 = 0u; j < params.n; j = j + 1u) {{
        acc = acc + a[a_row + j] * x[j];
    }}
    y[i] = acc;
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;
    Ok(())
}

/// Forward DFT via iterative radix-2 Decimation-In-Time over ONE workgroup of
/// `N = workgroup_size` threads (one complex element per thread; `N` a power of
/// two). Complex data is interleaved f32: element `j` is
/// `(input[2*j], input[2*j+1]) = (real, imag)`, so the input/output buffers hold
/// `2*N` f32. Stage 1 is a bit-reversal load into the `s_re`/`s_im` shared arrays;
/// then `log2(N)` butterfly stages run, each with a `workgroupBarrier()`. The same
/// forward sign convention `exp(-2*pi*i*k/m)` as the [`crate::wgsl_forge::oracle::dft_cpu`]
/// reference is used so GPU and CPU compute the identical transform. The shared
/// arrays come from the IR (`kernel.shared_memory`, sized to the workgroup size)
/// and `Op::Barrier` is the reusable IR primitive; the butterfly control flow is
/// specialised here because it is not expressible in the scalar op set.
fn emit_fft_wgsl(
    source: &mut String,
    kernel: &KernelSpec,
    schedule: Schedule,
) -> Result<(), ForgeError> {
    let wg = schedule.workgroup_size;

    writeln!(
        source,
        r#"
struct FftParams {{
    n: u32,
    log2n: u32,
    _pad0: u32,
    _pad1: u32,
}}
"#
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;

    for buffer in &kernel.buffers {
        let access = match buffer.access {
            crate::wgsl_forge::ir::BufferAccess::StorageRead => "storage, read",
            crate::wgsl_forge::ir::BufferAccess::StorageReadWrite => "storage, read_write",
            crate::wgsl_forge::ir::BufferAccess::Uniform => "uniform",
        };
        let type_decl = match buffer.access {
            crate::wgsl_forge::ir::BufferAccess::Uniform => "FftParams",
            _ => "array<f32>",
        };
        writeln!(
            source,
            "@group({}) @binding({}) var<{}> {}: {};",
            buffer.group, buffer.binding, access, buffer.name, type_decl
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    writeln!(source, "").map_err(|error| ForgeError::Emission(error.to_string()))?;
    for shared in &kernel.shared_memory {
        writeln!(
            source,
            "var<workgroup> {}: array<{}, {}>;",
            shared.name,
            shared.element.wgsl_name(),
            shared.length.resolve(wg)
        )
        .map_err(|error| ForgeError::Emission(error.to_string()))?;
    }

    // `-2*pi` as a valid WGSL f32 literal (the `f` suffix keeps it f32). The
    // butterfly indexing and the forward twiddle `exp(-2*pi*i*k/m)` follow the
    // certified algorithm exactly.
    writeln!(
        source,
        r#"
@compute @workgroup_size({wg})
fn {entry}(
    @builtin(local_invocation_id) local_id: vec3<u32>,
    @builtin(global_invocation_id) global_id: vec3<u32>,
) {{
    let t = local_id.x;                 // thread index 0..N-1
    let n = params.n;
    let logn = params.log2n;
    // 1. Bit-reversal load: element t goes to reverse(t).
    let rev = reverseBits(t) >> (32u - logn);
    s_re[rev] = input[2u * t];
    s_im[rev] = input[2u * t + 1u];
    workgroupBarrier();
    // 2. log2(N) butterfly stages. Stage s: span = 1<<s, m = 2*span.
    //    N/2 butterflies; threads 0..N/2-1 active.
    for (var s: u32 = 0u; s < logn; s = s + 1u) {{
        let span = 1u << s;
        let m = span << 1u;
        if (t < (n >> 1u)) {{
            let k = t & (span - 1u);                  // position within the butterfly group
            let j = ((t >> s) << (s + 1u)) + k;       // lower index = (t/span)*m + k
            let jp = j + span;                        // upper index
            // twiddle w = exp(-2*pi*i*k/m)
            let ang = -6.28318548f * f32(k) / f32(m);
            let wr = cos(ang);
            let wi = sin(ang);
            let ur = s_re[j];
            let ui = s_im[j];
            let vr = s_re[jp];
            let vi = s_im[jp];
            // v' = v * w
            let tr = vr * wr - vi * wi;
            let ti = vr * wi + vi * wr;
            s_re[j] = ur + tr;
            s_im[j] = ui + ti;
            s_re[jp] = ur - tr;
            s_im[jp] = ui - ti;
        }}
        workgroupBarrier();
    }}
    // 3. Store.
    output[2u * t] = s_re[t];
    output[2u * t + 1u] = s_im[t];
}}"#,
        wg = wg,
        entry = kernel.entry_point
    )
    .map_err(|error| ForgeError::Emission(error.to_string()))?;

    Ok(())
}

fn emit_ops(source: &mut String, ops: &[Op], indent: &str) -> Result<(), ForgeError> {
    for op in ops {
        match op {
            Op::StructLoad {
                buffer,
                field,
                destination,
            } => {
                writeln!(source, "{indent}let {destination} = {buffer}.{field};")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Load {
                buffer,
                index,
                destination,
            } => {
                writeln!(source, "{indent}let {destination} = {buffer}[{index}];")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Store {
                buffer,
                index,
                value,
            } => {
                writeln!(source, "{indent}{buffer}[{index}] = {value};")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Fma {
                a,
                b,
                c,
                destination,
            } => {
                writeln!(source, "{indent}let {destination} = {a} * {b} + {c};")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Mul {
                left,
                right,
                destination,
            } => {
                writeln!(source, "{indent}let {destination} = {left} * {right};")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Add {
                left,
                right,
                destination,
            } => {
                writeln!(source, "{indent}let {destination} = {left} + {right};")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::DotProduct {
                left_buffer,
                left_base,
                right_buffer,
                right_base,
                len,
                destination,
            } => {
                writeln!(source, "{indent}var {destination} = 0.0;")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(
                    source,
                    "{indent}for (var i: u32 = 0u; i < {len}; i = i + 1u) {{"
                )
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(source, "{indent}    {destination} = {destination} + {left_buffer}[{left_base} + i] * {right_buffer}[{right_base} + i];").map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(source, "{indent}}}")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Loop {
                induction_var,
                start,
                end,
                step,
                body,
            } => {
                writeln!(source, "{indent}for (var {induction_var}: u32 = {start}; {induction_var} < {end}; {induction_var} = {induction_var} + {step}) {{").map_err(|error| ForgeError::Emission(error.to_string()))?;
                emit_ops(source, body, &format!("{indent}    "))?;
                writeln!(source, "{indent}}}")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Relu {
                operand,
                destination,
            } => {
                writeln!(source, "{indent}let {destination} = max(0.0, {operand});")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Gelu {
                operand,
                destination,
            } => {
                // GELU approximation: 0.5 * x * (1.0 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
                writeln!(source, "{indent}let {destination} = 0.5 * {operand} * (1.0 + tanh(0.7978845608 * ({operand} + 0.044715 * {operand} * {operand} * {operand})));").map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::MatrixMultiply { .. } => {
                // No scalar WGSL lowering exists for a dense GEMM op; fail loudly
                // rather than silently emit nothing. Tensor-core GEMM is delivered
                // via the cooperative-matrix tile or the CUDA WMMA path.
                return Err(ForgeError::Emission(
                    "Op::MatrixMultiply has no scalar WGSL lowering; use coopmat::matmul_tc_wgsl (cooperative-matrix) or the CUDA WMMA path".to_string(),
                ));
            }
            Op::Barrier => {
                writeln!(source, "{indent}workgroupBarrier();")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Intrinsic(crate::wgsl_forge::ir::Intrinsic::RayQuery {
                acceleration_structure,
                origin,
                direction,
                t_min,
                t_max,
                destination,
            }) => {
                // Hardware ray-query: initialise, proceed to completion, read the
                // committed hit. Traversal must loop until `rayQueryProceed` returns
                // false — a single call can leave a multi-node BVH partially walked.
                writeln!(source, "{indent}var rq: ray_query;")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(source, "{indent}rayQueryInitialize(&rq, {acceleration_structure}, RayDesc(0u, 0xFFu, {t_min}, {t_max}, {origin}, {direction}));").map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(
                    source,
                    "{indent}loop {{ if (!rayQueryProceed(&rq)) {{ break; }} }}"
                )
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(
                    source,
                    "{indent}let committed = rayQueryGetCommittedIntersection(&rq);"
                )
                .map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(source, "{indent}var {destination} = -1.0;")
                    .map_err(|error| ForgeError::Emission(error.to_string()))?;
                writeln!(source, "{indent}if (committed.kind != RAY_QUERY_INTERSECTION_NONE) {{ {destination} = committed.t; }}").map_err(|error| ForgeError::Emission(error.to_string()))?;
            }
            Op::Intrinsic(_) => {
                return Err(ForgeError::Emission(
                    "Intrinsics not implemented for WGSL yet".to_string(),
                ));
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::wgsl_forge::{BuiltinKernel, Schedule};

    #[test]
    fn emission_is_byte_deterministic() {
        let kernel = BuiltinKernel::AffineF32.spec();
        let first = emit_wgsl(&kernel, Schedule::default()).unwrap();
        let second = emit_wgsl(&kernel, Schedule::default()).unwrap();
        assert_eq!(first.source, second.source);
    }

    /// Phase-1 DAG-IR seam (docs/plans/dag-ir-forge.md §8.2): routing gemm/gemv/fft through
    /// `KernelSpec::to_graph` → `lower_graph` (serviced by `WgslDelegateLowerer`) must emit
    /// the EXACT body the legacy leaf emitter produces — byte-identical, since the certify
    /// cache key folds in the source_hash. This also guards the bridge: if `to_graph`
    /// produced the wrong op-class, the routed body would not match the right leaf.
    #[test]
    fn graph_routing_is_byte_equal_for_gemm_gemv_fft() {
        let sched = Schedule::default();
        let cases: [(&str, KernelSpec); 3] = [
            ("gemm", BuiltinKernel::Gemm.spec()),
            ("gemv", BuiltinKernel::Gemv.spec()),
            ("fft", BuiltinKernel::Fft.spec()),
        ];
        for (id, kernel) in cases {
            let routed = emit_wgsl(&kernel, sched).unwrap();
            let mut legacy_body = String::new();
            match id {
                "gemm" => emit_gemm_wgsl(&mut legacy_body, &kernel, sched).unwrap(),
                "gemv" => emit_gemv_wgsl(&mut legacy_body, &kernel, sched).unwrap(),
                "fft" => emit_fft_wgsl(&mut legacy_body, &kernel, sched).unwrap(),
                _ => unreachable!(),
            }
            assert!(
                routed.source.ends_with(&legacy_body),
                "{id}: graph-routed WGSL body must be byte-identical to the legacy leaf body"
            );
        }
    }

    /// Phase-2: a native `Reduce` graph node lowers through `emit_graph_wgsl` →
    /// `WgslGraphLowerer` to the `reduce_wgsl` template (no `KernelSpec` involved). This
    /// proves the native-template lowering path, distinct from Phase 1's delegation.
    #[test]
    fn graph_native_reduce_lowers_to_template() {
        use crate::wgsl_forge::graph_ops::reduce::{reduce_wgsl, REDUCE_ENTRY};
        use crate::wgsl_forge::ir::graph::{
            Axis, ComputeGraph, DType, OpNode, RedKind, Shape, TensorRef,
        };

        let sched = Schedule {
            workgroup_size: 256,
            ..Default::default()
        };
        let mut g = ComputeGraph::new();
        let inp = TensorRef::external(Shape::new(&[1024]), DType::F32);
        g.push(
            OpNode::Reduce {
                op: RedKind::Sum,
                axis: Axis::Last,
            },
            &[inp],
            Shape::new(&[1]),
            DType::F32,
            sched,
        )
        .unwrap();
        let shader = emit_graph_wgsl(&g, sched).unwrap();
        assert!(shader.source.contains(REDUCE_ENTRY));
        assert!(
            shader.source.ends_with(&reduce_wgsl(RedKind::Sum, 256)),
            "native graph emission must equal the reduce template"
        );
    }

    #[test]
    fn p64_kernel_lays_64bit_fields_as_paired_u32_words() {
        // The p64 kernel carries 64-bit P64 fields. WGSL has no portable native u64,
        // so the generated module must lay them out as packed `u32` words. Assert the
        // exact packed layout the emitter produces (plan §1 / §10): a `P64Words64`
        // struct of `array<vec4<u32>, 4>` lanes, bound as `array<P64Words64>`, and the
        // per-word `f32(...)` projection over those `u32` lanes.
        let kernel = BuiltinKernel::P64Project.spec();
        let generated = emit_wgsl(&kernel, Schedule::default()).unwrap();
        let source = &generated.source;

        assert!(
            source.contains("lanes: array<vec4<u32>, 4>"),
            "expected paired/packed vec4<u32> word layout, got:\n{source}"
        );
        assert!(
            source.contains("var<storage, read> input: array<P64Words64>"),
            "expected the P64 record buffer bound as packed u32 words, got:\n{source}"
        );
        // The 64-bit data is read as u32 words and projected via f32(word) — i.e. no
        // native 64-bit scalar appears in the kernel body.
        assert!(
            source.contains("f32(word)"),
            "expected per-u32-word projection, got:\n{source}"
        );
        assert!(
            !source.contains("u64") && !source.contains("f64"),
            "the portable p64 kernel must not emit any native 64-bit scalar, got:\n{source}"
        );
    }
}