feat(compile-bridge): add Q0.2 GPU enumeration mode

This commit is contained in:
Brandon Schneider 2026-05-20 18:42:48 -05:00
parent b02bb47b53
commit c1369611de
3 changed files with 450 additions and 15 deletions

View file

@ -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::<Vec<_>>()
}
} 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::<Vec<_>>()
})
.collect::<Vec<_>>()
};
let total = theorems.len();

View file

@ -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<Vec<TheoremReceipt>> {
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<u32> = 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<u32> = 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<u32> = 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<u8> = mapped.to_vec();
staging_buffer.unmap();
let result_u32s: Vec<u32> = 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)
}

View file

@ -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<storage, read> vectors: array<TestVector>;
@group(0) @binding(1) var<storage, read> batches: array<TheoremBatch>;
@group(0) @binding(2) var<storage, read_write> results: array<TheoremResult>;
// Q0_2 Constants
const Q0_2_VALUES: array<u32, 4> = array<u32, 4>(
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<u32>) {
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);
}
}