SilverSight/python/dna_braid.wgsl

/**
 * dna_braid.wgsl — WebGPU Compute Shader: Braid Sort on DNA Sequences
 *
 * Treats GPU actions as triangle math:
 *   - Each braid strand = triangle vertex
 *   - Each crossing = triangle rotation (compare-swap)
 *   - Eigensolid convergence = sorted output
 *   - Workgroup = triangle mesh
 *
 * The 8 Hachimoji bases (A,B,C,G,P,S,T,Z) map to 8 triangle vertices.
 * A braid crossing swaps two adjacent vertices if they're out of order.
 * The eigensolid is the fixed point where no crossings remain.
 *
 * WebGPU compute shader — runs on any GPU with WebGPU support.
 * Zero-copy: reads/writes directly to GPU buffer.
 */

// ============================================================
// CONSTANTS
// ============================================================

const N_BASES: u32 = 8u;
const BASE_A: u32 = 0u;
const BASE_B: u32 = 1u;
const BASE_C: u32 = 2u;
const BASE_G: u32 = 3u;
const BASE_P: u32 = 4u;
const BASE_S: u32 = 5u;
const BASE_T: u32 = 6u;
const BASE_Z: u32 = 7u;

// Workgroup size (must be power of 2 for radix sort)
const WORKGROUP_SIZE: u32 = 256u;

// ============================================================
// BINDINGS
// ============================================================

@group(0) @binding(0) var<storage, read> sequences: array<u32>;       // packed DNA sequences
@group(0) @binding(1) var<storage, read_write> indices: array<u32>;   // sort indices (in/out)
@group(0) @binding(2) var<storage, read_write> scratch: array<u32>;   // scratch buffer
@group(0) @binding(3) var<uniform> params: Params;                     // parameters

struct Params {
    n_sequences: u32,      // number of sequences
    seq_length: u32,       // bases per sequence
    radix_pass: u32,       // current radix pass (0 = LSD)
    _pad: u32,             // alignment
};

// ============================================================
// TRIANGLE MATH: braid crossing as triangle rotation
// ============================================================

/// Braid crossing: compare-swap two adjacent elements.
/// This is a triangle rotation in the permutation space.
/// If a > b, rotate the triangle (swap a and b).
fn braid_cross(a: u32, b: u32) -> vec2<u32> {
    // Triangle rotation: if out of order, swap
    if (a > b) {
        return vec2<u32>(b, a);  // rotated
    }
    return vec2<u32>(a, b);      // unchanged
}

/// Extract a single digit (base) from a packed sequence.
/// Sequences are packed as base-8 digits in a u32.
/// digit_index 0 = most significant (leftmost) base.
fn extract_digit(sequence: u32, digit_index: u32, seq_length: u32) -> u32 {
    // LSD-first extraction: rightmost digit is index 0
    let shift = digit_index * 3u;  // 3 bits per base (base-8)
    return (sequence >> shift) & 0x7u;
}

/// Braid eigensolid check: is the sequence sorted at this digit?
/// Returns true if no crossings needed (converged).
fn is_eigensolid(a_digit: u32, b_digit: u32, a_idx: u32, b_idx: u32) -> bool {
    // Eigensolid: a_digit < b_digit, or equal with correct index order
    return (a_digit < b_digit) || (a_digit == b_digit && a_idx <= b_idx);
}

// ============================================================
// KERNEL 1: RADIX SORT — digit extraction
// ============================================================

/// Extract the current radix digit for all sequences.
/// Each thread handles one sequence.
@compute @workgroup_size(WORKGROUP_SIZE)
fn extract_digits(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx >= params.n_sequences) {
        return;
    }

    let seq_idx = indices[idx];
    let sequence = sequences[seq_idx];
    let digit = extract_digit(sequence, params.radix_pass, params.seq_length);

    // Store digit in scratch buffer (for counting sort)
    scratch[idx] = digit;
}

// ============================================================
// KERNEL 2: BRAID SORT — triangle mesh compare-swap
// ============================================================

/// Odd-even transposition sort (braid pattern).
/// Each workgroup handles a chunk of the array.
/// Alternates between odd and even phases.
/// Each comparison is a braid crossing (triangle rotation).
@compute @workgroup_size(WORKGROUP_SIZE)
fn braid_sort_odd(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x * 2u;  // even indices
    if (idx + 1u >= params.n_sequences) {
        return;
    }

    // Extract digits for this pair
    let seq_a = indices[idx];
    let seq_b = indices[idx + 1u];
    let digit_a = extract_digit(sequences[seq_a], params.radix_pass, params.seq_length);
    let digit_b = extract_digit(sequences[seq_b], params.radix_pass, params.seq_length);

    // Braid crossing: compare-swap
    let crossed = braid_cross(digit_a, digit_b);
    if (crossed.x != digit_a) {
        // Crossing occurred: swap indices
        indices[idx] = seq_b;
        indices[idx + 1u] = seq_a;
    }
}

@compute @workgroup_size(WORKGROUP_SIZE)
fn braid_sort_even(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x * 2u + 1u;  // odd indices
    if (idx + 1u >= params.n_sequences) {
        return;
    }

    let seq_a = indices[idx];
    let seq_b = indices[idx + 1u];
    let digit_a = extract_digit(sequences[seq_a], params.radix_pass, params.seq_length);
    let digit_b = extract_digit(sequences[seq_b], params.radix_pass, params.seq_length);

    let crossed = braid_cross(digit_a, digit_b);
    if (crossed.x != digit_a) {
        indices[idx] = seq_b;
        indices[idx + 1u] = seq_a;
    }
}

// ============================================================
// KERNEL 3: EIGENSOLID CHECK — convergence test
// ============================================================

/// Check if the braid has converged (eigensolid reached).
/// Each thread checks one pair. Writes 1 to scratch if crossing needed.
@compute @workgroup_size(WORKGROUP_SIZE)
fn eigensolid_check(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx + 1u >= params.n_sequences) {
        scratch[idx] = 0u;
        return;
    }

    let seq_a = indices[idx];
    let seq_b = indices[idx + 1u];
    let digit_a = extract_digit(sequences[seq_a], params.radix_pass, params.seq_length);
    let digit_b = extract_digit(sequences[seq_b], params.radix_pass, params.seq_length);

    // Eigensolid: no crossing needed = converged
    if (is_eigensolid(digit_a, digit_b, seq_a, seq_b)) {
        scratch[idx] = 0u;  // converged
    } else {
        scratch[idx] = 1u;  // needs crossing
    }
}

// ============================================================
// KERNEL 4: COUNTING SORT — radix distribution
// ============================================================

/// Counting sort for radix sort distribution phase.
/// Each thread handles one element, computes its bucket.
@compute @workgroup_size(WORKGROUP_SIZE)
fn counting_sort(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx >= params.n_sequences) {
        return;
    }

    let digit = scratch[idx];  // digit was extracted in extract_digits
    // Store (digit, index) pair for stable sort
    // Pack: high 3 bits = digit, low 29 bits = index
    scratch[idx] = (digit << 29u) | (idx & 0x1FFFFFFFu);
}

// ============================================================
// KERNEL 5: TRIANGLE MESH — parallel reduction
// ============================================================

/// Parallel reduction to find minimum energy solution.
/// Uses triangle math: each pair reduces to a single vertex.
/// The final vertex is the optimum.
@compute @workgroup_size(WORKGROUP_SIZE)
fn reduce_min(@builtin(global_invocation_id) gid: vec3<u32>,
              @builtin(local_invocation_id) lid: vec3<u32>) {
    // Workgroup-local reduction
    // Each thread starts with its own value
    // Pairs reduce like triangle vertices merging
    // After log2(WORKGROUP_SIZE) steps, one vertex remains

    // This is a placeholder for the reduction kernel
    // In practice, this would reduce the scratch buffer
    // to find the minimum-energy index
}

// ============================================================
// MAIN ENTRY POINTS
// ============================================================

/// Dispatch the full braid sort pipeline.
/// Call this from JavaScript/WebGPU API:
///   1. dispatch(extract_digits)     — extract radix digits
///   2. dispatch(braid_sort_odd/even) — braid crossing passes
///   3. dispatch(eigensolid_check)   — convergence test
///   4. repeat 2-3 until converged
///   5. dispatch(reduce_min)         — find optimum