From cf7d1f2d241c222731f40b830e66a6f55753b71a Mon Sep 17 00:00:00 2001 From: Brandon Schneider Date: Wed, 20 May 2026 18:42:48 -0500 Subject: [PATCH] feat(compile-bridge): add Q0.2 GPU enumeration mode --- .../shim/lake_compile_bridge/src/main.rs | 65 ++++- .../lake_compile_bridge/src/q02_dispatch.rs | 271 ++++++++++++++++++ .../src/shaders/q02_enumeration.wgsl | 129 +++++++++ 3 files changed, 450 insertions(+), 15 deletions(-) create mode 100644 4-Infrastructure/shim/lake_compile_bridge/src/q02_dispatch.rs create mode 100644 4-Infrastructure/shim/lake_compile_bridge/src/shaders/q02_enumeration.wgsl diff --git a/4-Infrastructure/shim/lake_compile_bridge/src/main.rs b/4-Infrastructure/shim/lake_compile_bridge/src/main.rs index 7e828425..10a4d719 100644 --- a/4-Infrastructure/shim/lake_compile_bridge/src/main.rs +++ b/4-Infrastructure/shim/lake_compile_bridge/src/main.rs @@ -21,27 +21,38 @@ use clap::Parser; use serde::Serialize; mod gpu; +mod q02_dispatch; // ── Theorem Registry ──────────────────────────────────────────────────── /// FixedPoint theorems that can be GPU-verified. const THEOREMS: &[(&str, u32)] = &[ - ("zero_mul", 0), - ("mul_zero", 1), - ("add_zero", 2), - ("zero_add", 3), - ("sub_self", 4), - ("one_mul", 5), - ("mul_one", 6), - ("add_comm", 7), + ("zero_mul", 0), + ("mul_zero", 1), + ("add_zero", 2), + ("zero_add", 3), + ("sub_self", 4), + ("one_mul", 5), + ("mul_one", 6), + ("add_comm", 7), ("neg_involutive", 8), - ("sub_via_neg", 9), + ("sub_via_neg", 9), +]; + +/// Q0_2 theorems for exhaustive GPU verification. +const Q02_THEOREMS: &[(&str, u32)] = &[ + ("q0_2_mul_self_nonneg", 10), + ("q0_2_mul_nonneg", 11), + ("q0_2_add_nonneg", 12), ]; // ── CLI ───────────────────────────────────────────────────────────────── #[derive(Parser, Debug)] -#[command(name = "lake_compile_bridge", about = "GPU-accelerated Lean build bridge")] +#[command( + name = "lake_compile_bridge", + about = "GPU-accelerated Lean build bridge" +)] struct Args { /// Lake build target (default: "Semantics.FixedPoint") #[arg(short, long, default_value = "Semantics.FixedPoint")] @@ -59,6 +70,10 @@ struct Args { #[arg(short, long, default_value = "build_receipt.json")] receipt: String, + /// Run Q0_2 exhaustive enumeration instead of FixedPoint random sampling + #[arg(long, default_value_t = false)] + q02: bool, + /// Dry run: print what would be done without running #[arg(long, default_value_t = false)] dry_run: bool, @@ -139,19 +154,39 @@ fn main() -> anyhow::Result<()> { }; // ── Stage 3: GPU theorem verification ───────────────────────────── - let theorems = if gpu_available && !args.dry_run { + let theorems = if args.q02 { + if gpu_available && !args.dry_run { + eprintln!( + " dispatching {} Q0_2 theorems to GPU...", + Q02_THEOREMS.len() + ); + q02_dispatch::verify_q02_on_gpu()? + } else { + eprintln!(" (GPU unavailable or dry-run; marking all Q0_2 theorems as untested)"); + Q02_THEOREMS + .iter() + .map(|(name, id)| TheoremReceipt { + name: name.to_string(), + theorem_id: *id, + tested: 0, + passed: false, + }) + .collect::>() + } + } else if gpu_available && !args.dry_run { eprintln!(" dispatching {} theorems to GPU...", THEOREMS.len()); gpu::verify_theorems_on_gpu(THEOREMS, args.vectors)? } else { eprintln!(" (GPU unavailable or dry-run; marking all theorems as untested)"); - THEOREMS.iter().map(|(name, id)| { - TheoremReceipt { + THEOREMS + .iter() + .map(|(name, id)| TheoremReceipt { name: name.to_string(), theorem_id: *id, tested: 0, passed: false, - } - }).collect::>() + }) + .collect::>() }; let total = theorems.len(); diff --git a/4-Infrastructure/shim/lake_compile_bridge/src/q02_dispatch.rs b/4-Infrastructure/shim/lake_compile_bridge/src/q02_dispatch.rs new file mode 100644 index 00000000..60292a45 --- /dev/null +++ b/4-Infrastructure/shim/lake_compile_bridge/src/q02_dispatch.rs @@ -0,0 +1,271 @@ +// ── Q0_2 GPU Dispatch ──────────────────────────────────────── +// Exhaustive Q0_2 value enumeration via wgpu + WGSL +// +// Q0_2 values: {0, 1/4, 1/2, 3/4} encoded as Q16_16: +// [0x00000000, 0x00004000, 0x00008000, 0x0000C000] +// +// Theorems: +// 10 (q0_2_mul_self_nonneg): 4 vectors (each value paired with itself) +// 11 (q0_2_mul_nonneg): 16 vectors (all 4x4 pairs) +// 12 (q0_2_add_nonneg): 16 vectors (all 4x4 pairs) + +use crate::TheoremReceipt; +use std::borrow::Cow; +use wgpu::util::DeviceExt; + +/// Q0_2 theorem definitions for GPU dispatch. +pub const Q02_THEOREMS: &[(&str, u32)] = &[ + ("q0_2_mul_self_nonneg", 10), + ("q0_2_mul_nonneg", 11), + ("q0_2_add_nonneg", 12), +]; + +/// The four Q0_2 values as Q16_16 fixed-point. +const Q02_VALUES: [u32; 4] = [ + 0x00000000, // 0 + 0x00004000, // 1/4 + 0x00008000, // 1/2 + 0x0000C000, // 3/4 +]; + +/// Verify Q0_2 theorems on the GPU using exhaustive enumeration. +pub fn verify_q02_on_gpu() -> anyhow::Result> { + let instance = wgpu::Instance::default(); + + let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::HighPerformance, + compatible_surface: None, + force_fallback_adapter: false, + })) + .ok_or_else(|| anyhow::anyhow!("No GPU adapter found"))?; + + let mut limits = wgpu::Limits::default(); + limits.max_storage_buffer_binding_size = adapter.limits().max_storage_buffer_binding_size; + limits.max_buffer_size = adapter.limits().max_buffer_size; + limits.max_compute_invocations_per_workgroup = + adapter.limits().max_compute_invocations_per_workgroup; + + let (device, queue) = pollster::block_on(adapter.request_device( + &wgpu::DeviceDescriptor { + label: None, + required_features: wgpu::Features::empty(), + required_limits: limits, + }, + None, + ))?; + + // ── Generate test vectors ───────────────────────────────────────── + // Exhaustive Q0_2 enumeration: + // theorem 10 (mul_self_nonneg): 4 vectors (value paired with itself) + // theorem 11 (mul_nonneg): 16 vectors (all 4x4 pairs) + // theorem 12 (add_nonneg): 16 vectors (all 4x4 pairs) + // + // We generate 16 vectors total (theorem 10 only uses first 4). + + let num_theorems = Q02_THEOREMS.len() as u32; + let num_vectors = 16u32; + + let mut test_vectors: Vec = Vec::with_capacity((num_vectors * 3) as usize); + for i in 0..4 { + for j in 0..4 { + let a = Q02_VALUES[i]; + let b = Q02_VALUES[j]; + test_vectors.push(a); + test_vectors.push(b); + test_vectors.push(0); // expected (unused for property-based checks) + } + } + + // ── Create buffers ──────────────────────────────────────────────── + + let vectors_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Q0_2 Test Vectors"), + contents: bytemuck::cast_slice(&test_vectors), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }); + + let mut batch_data: Vec = Vec::with_capacity((num_theorems * 4) as usize); + for &(_name, id) in Q02_THEOREMS { + batch_data.push(id); + batch_data.push(num_vectors); + batch_data.push(0); + batch_data.push(0); + } + let batches_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Q0_2 Theorem Batches"), + contents: bytemuck::cast_slice(&batch_data), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }); + + let mut results_init: Vec = Vec::with_capacity((num_theorems * 4) as usize); + for &(_name, id) in Q02_THEOREMS { + results_init.push(id); + results_init.push(1); + results_init.push(num_vectors); + results_init.push(0); + } + let results_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Q0_2 Results"), + contents: bytemuck::cast_slice(&results_init), + usage: wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_SRC + | wgpu::BufferUsages::COPY_DST, + }); + + let staging_size = (num_theorems * 4 * 4) as u64; + let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Q0_2 Staging"), + size: staging_size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + // ── Shader module ───────────────────────────────────────────────── + let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("Q0_2 Enumeration Shader"), + source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!( + "shaders/q02_enumeration.wgsl" + ))), + }); + + // ── Bind group layout ───────────────────────────────────────────── + let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: None, + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 2, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + // ── Pipeline ────────────────────────────────────────────────────── + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: None, + bind_group_layouts: &[&bind_group_layout], + push_constant_ranges: &[], + }); + + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: Some(&pipeline_layout), + module: &shader, + entry_point: "main", + }); + + // ── Bind group ──────────────────────────────────────────────────── + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: vectors_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: batches_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: results_buffer.as_entire_binding(), + }, + ], + }); + + // ── Dispatch ────────────────────────────────────────────────────── + let mut encoder = + device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + + { + let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: None, + timestamp_writes: None, + }); + compute_pass.set_pipeline(&pipeline); + compute_pass.set_bind_group(0, &bind_group, &[]); + + let workgroups_y = ((num_vectors + 63) / 64).max(1); + compute_pass.dispatch_workgroups(num_theorems, workgroups_y, 1); + } + + encoder.copy_buffer_to_buffer(&results_buffer, 0, &staging_buffer, 0, staging_size); + + queue.submit(Some(encoder.finish())); + + // ── Readback ────────────────────────────────────────────────────── + let (sender, receiver) = std::sync::mpsc::channel(); + let buffer_slice = staging_buffer.slice(..); + buffer_slice.map_async(wgpu::MapMode::Read, move |v| { + let _ = sender.send(v); + }); + + device.poll(wgpu::Maintain::Wait); + receiver + .recv() + .map_err(|e| anyhow::anyhow!("Q0_2 GPU readback channel error: {:?}", e))? + .map_err(|e| anyhow::anyhow!("Q0_2 GPU buffer map error: {:?}", e))?; + + let mapped = buffer_slice.get_mapped_range(); + let result_bytes: Vec = mapped.to_vec(); + staging_buffer.unmap(); + + let result_u32s: Vec = result_bytes + .chunks_exact(4) + .map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect(); + + let mut receipts = Vec::with_capacity(Q02_THEOREMS.len()); + for (i, &(name, _id)) in Q02_THEOREMS.iter().enumerate() { + let base = i * 4; + let theorem_id = result_u32s[base]; + let _passed_flag = result_u32s[base + 1]; + let total = result_u32s[base + 2]; + let failed = result_u32s[base + 3]; + + let passed_val = failed == 0; + receipts.push(TheoremReceipt { + name: name.to_string(), + theorem_id, + tested: total, + passed: passed_val, + }); + + if passed_val { + eprintln!(" Q0_2 ✓ {} passed ({} vectors)", name, total); + } else { + eprintln!( + " Q0_2 ✗ {} FAILED ({}/{} vectors failed)", + name, failed, total + ); + } + } + + Ok(receipts) +} diff --git a/4-Infrastructure/shim/lake_compile_bridge/src/shaders/q02_enumeration.wgsl b/4-Infrastructure/shim/lake_compile_bridge/src/shaders/q02_enumeration.wgsl new file mode 100644 index 00000000..4b71a3ff --- /dev/null +++ b/4-Infrastructure/shim/lake_compile_bridge/src/shaders/q02_enumeration.wgsl @@ -0,0 +1,129 @@ +// ── Q0_2 Exhaustive Enumeration Verification Shader ─────────────── +// Each workgroup verifies one Q0_2 theorem across all 4 Q0_2 values +// (16 pairs for binary operations). +// +// Q0_2 is a 2-bit fixed-point type with 4 values: +// 0 → 0x00000000 +// 0.25 → 0x00004000 +// 0.5 → 0x00008000 +// 0.75 → 0x0000C000 +// +// Theorem IDs: +// 10: q0_2_mul_self_nonneg — ∀ a ∈ Q0_2, (a*a).toInt ≥ 0 +// 11: q0_2_mul_nonneg — ∀ a,b ∈ Q0_2, (a*b).toInt ≥ 0 +// 12: q0_2_add_nonneg — ∀ a,b ∈ Q0_2, (a+b).toInt ≥ 0 + +struct TestVector { + a: u32, + b: u32, + expected: u32, +} + +struct TheoremBatch { + theorem_id: u32, + count: u32, + padding: u32, + padding2: u32, +} + +struct TheoremResult { + theorem_id: u32, + passed: u32, + total: u32, + failed: u32, +} + +@group(0) @binding(0) var vectors: array; +@group(0) @binding(1) var batches: array; +@group(0) @binding(2) var results: array; + +// ── Q0_2 Constants ───────────────────────────────────────────────── + +const Q0_2_VALUES: array = array( + 0x00000000u, + 0x00004000u, + 0x00008000u, + 0x0000C000u, +); + +// ── Q16_16 Arithmetic (matches FixedPoint.lean) ───────────────────── + +fn q16_add(a: u32, b: u32) -> u32 { + let s: u32 = a + b; + if (a < 0x80000000u && b < 0x80000000u && s >= 0x80000000u) { + return 0x7FFFFFFFu; + } + if (a >= 0x80000000u && b >= 0x80000000u && s < 0x80000000u) { + return 0x80000000u; + } + return s; +} + +fn q16_mul(a: u32, b: u32) -> u32 { + let prod: u64 = u64(a) * u64(b); + return u32(prod >> 16u); +} + +fn q16_to_int(a: u32) -> i32 { + if (a >= 0x80000000u) { + return i32(a) - 0x100000000i; + } + return i32(a); +} + +// ── Q0_2 Verification Kernels ─────────────────────────────────────── + +fn check_q0_2_mul_self_nonneg(idx: u32) -> bool { + let a = Q0_2_VALUES[idx]; + let result = q16_mul(a, a); + return q16_to_int(result) >= 0i; +} + +fn check_q0_2_mul_nonneg(idx: u32) -> bool { + let ai = idx / 4u; + let bi = idx % 4u; + let a = Q0_2_VALUES[ai]; + let b = Q0_2_VALUES[bi]; + let result = q16_mul(a, b); + return q16_to_int(result) >= 0i; +} + +fn check_q0_2_add_nonneg(idx: u32) -> bool { + let ai = idx / 4u; + let bi = idx % 4u; + let a = Q0_2_VALUES[ai]; + let b = Q0_2_VALUES[bi]; + let result = q16_add(a, b); + return q16_to_int(result) >= 0i; +} + +// ── Main Dispatch ─────────────────────────────────────────────────── + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let batch_idx = gid.x; + let local_idx = gid.y; + + if (batch_idx >= arrayLength(&batches)) { + return; + } + + let batch = batches[batch_idx]; + + if (local_idx >= batch.count) { + return; + } + + var pass: bool = false; + + switch (batch.theorem_id) { + case 10u: { pass = check_q0_2_mul_self_nonneg(local_idx); } + case 11u: { pass = check_q0_2_mul_nonneg(local_idx); } + case 12u: { pass = check_q0_2_add_nonneg(local_idx); } + default: { pass = true; } + } + + if (!pass) { + atomicAdd(&results[batch_idx].failed, 1u); + } +}