sort of almost
This commit is contained in:
illuzen
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -569,6 +569,7 @@ dependencies = [
|
||||
"qp-poseidon-constants 1.0.1",
|
||||
"qp-poseidon-core 1.0.1",
|
||||
"rand 0.8.5",
|
||||
"rand_chacha 0.3.1",
|
||||
"thiserror 1.0.69",
|
||||
"wgpu",
|
||||
]
|
||||
|
||||
@@ -30,6 +30,8 @@ qp-plonky2-field = { path = "../../../qp-plonky2/field" }
|
||||
wgpu = { version = "27.0.1" } # GPU compute library
|
||||
futures = "0.3" # For async executor
|
||||
bytemuck = "1.16" # For buffer mapping
|
||||
rand = { workspace = true, features = ["std", "std_rng"] }
|
||||
rand_chacha = "0.3" # For deterministic random test generation
|
||||
|
||||
[dev-dependencies]
|
||||
hex = { workspace = true }
|
||||
|
||||
276
crates/engine-gpu/src/gf_mul_test.wgsl
Normal file
276
crates/engine-gpu/src/gf_mul_test.wgsl
Normal file
@@ -0,0 +1,276 @@
|
||||
// WGSL Goldilocks Field and Operations for Testing
|
||||
|
||||
// Goldilocks field element represented as [low_32, high_32]
|
||||
struct GoldilocksField {
|
||||
low: u32,
|
||||
high: u32,
|
||||
}
|
||||
|
||||
// Struct to represent a pair of GoldilocksField values (replaces vec2<GoldilocksField>)
|
||||
struct GoldilocksFieldPair {
|
||||
first: GoldilocksField,
|
||||
second: GoldilocksField,
|
||||
}
|
||||
|
||||
// Field modulus P = 2^64 - 2^32 + 1
|
||||
// P_LOW = 1
|
||||
// P_HIGH = 0xFFFFFFFF
|
||||
const P_LOW: u32 = 1u;
|
||||
const P_HIGH: u32 = 0xFFFFFFFFu;
|
||||
|
||||
// EPSILON = 2^32 - 1
|
||||
const EPSILON: u32 = 0xFFFFFFFFu;
|
||||
|
||||
fn gf_zero() -> GoldilocksField {
|
||||
return GoldilocksField(0u, 0u);
|
||||
}
|
||||
|
||||
fn gf_one() -> GoldilocksField {
|
||||
return GoldilocksField(1u, 0u);
|
||||
}
|
||||
|
||||
// Helper for 64-bit addition with carry
|
||||
// Returns vec2<u32>(sum, carry_out)
|
||||
fn u32_add_with_carry(a: u32, b: u32, carry_in: u32) -> vec2<u32> {
|
||||
let sum_temp = a + b;
|
||||
let carry_0 = select(0u, 1u, sum_temp < a); // Carry from a + b
|
||||
let sum = sum_temp + carry_in;
|
||||
let carry_1 = select(0u, 1u, sum < sum_temp); // Carry from (a+b) + carry_in
|
||||
return vec2<u32>(sum, carry_0 + carry_1);
|
||||
}
|
||||
|
||||
// Helper for 64-bit subtraction with borrow
|
||||
// Returns vec2<u32>(diff, borrow_out)
|
||||
fn u32_sub_with_borrow(a: u32, b: u32, borrow_in: u32) -> vec2<u32> {
|
||||
let diff_temp = a - b;
|
||||
let borrow_0 = select(0u, 1u, diff_temp > a); // Borrow from a - b
|
||||
let diff = diff_temp - borrow_in;
|
||||
let borrow_1 = select(0u, 1u, diff > diff_temp); // Borrow from (a-b) - borrow_in
|
||||
return vec2<u32>(diff, borrow_0 + borrow_1);
|
||||
}
|
||||
|
||||
// Multiplies two u32s to produce a 64-bit result (low, high)
|
||||
fn u32_mul_extended(a: u32, b: u32) -> GoldilocksField {
|
||||
let a_low = a & 0xFFFFu;
|
||||
let a_high = a >> 16u;
|
||||
let b_low = b & 0xFFFFu;
|
||||
let b_high = b >> 16u;
|
||||
|
||||
let p0 = a_low * b_low;
|
||||
let p1 = a_low * b_high;
|
||||
let p2 = a_high * b_low;
|
||||
let p3 = a_high * b_high;
|
||||
|
||||
// Add carries
|
||||
let c1 = p0 >> 16u;
|
||||
let sum1 = p1 + p2 + c1;
|
||||
let c2 = sum1 >> 16u;
|
||||
|
||||
let low = (sum1 << 16u) | (p0 & 0xFFFFu);
|
||||
let high = p3 + c2;
|
||||
|
||||
return GoldilocksField(low, high);
|
||||
}
|
||||
|
||||
// Multiplies two GoldilocksField values (representing 64-bit numbers).
|
||||
// Returns a 128-bit result as (GoldilocksField_low_64_bits, GoldilocksField_high_64_bits).
|
||||
fn gf_mul_u64_by_u64(a: GoldilocksField, b: GoldilocksField) -> GoldilocksFieldPair {
|
||||
let a0 = a.low;
|
||||
let a1 = a.high;
|
||||
let b0 = b.low;
|
||||
let b1 = b.high;
|
||||
|
||||
let p0 = u32_mul_extended(a0, b0); // 64-bit: (p0.low, p0.high)
|
||||
let p1 = u32_mul_extended(a1, b0); // 64-bit: (p1.low, p1.high)
|
||||
let p2 = u32_mul_extended(a0, b1); // 64-bit: (p2.low, p2.high)
|
||||
let p3 = u32_mul_extended(a1, b1); // 64-bit: (p3.low, p3.high)
|
||||
|
||||
// Combine the products:
|
||||
// Full product is p3 * 2^64 + (p1 + p2) * 2^32 + p0
|
||||
|
||||
// Low 32 bits of the 128-bit result
|
||||
let res0 = p0.low;
|
||||
|
||||
// Next 32 bits (from p0.high, p1.low, p2.low)
|
||||
let sum1_temp = u32_add_with_carry(p0.high, p1.low, 0u);
|
||||
let sum1_final = u32_add_with_carry(sum1_temp.x, p2.low, sum1_temp.y);
|
||||
let res1 = sum1_final.x;
|
||||
let carry_to_res2_from_sum1 = sum1_final.y;
|
||||
|
||||
// Next 32 bits (from p1.high, p2.high, p3.low, carry_to_res2_from_sum1)
|
||||
let sum2_temp = u32_add_with_carry(p1.high, p2.high, 0u);
|
||||
let sum2_temp2 = u32_add_with_carry(sum2_temp.x, p3.low, sum2_temp.y);
|
||||
let sum2_final = u32_add_with_carry(sum2_temp2.x, carry_to_res2_from_sum1, sum2_temp2.y);
|
||||
let res2 = sum2_final.x;
|
||||
let carry_to_res3_from_sum2 = sum2_final.y;
|
||||
|
||||
// High 32 bits (from p3.high, carry_to_res3_from_sum2)
|
||||
let res3 = p3.high + carry_to_res3_from_sum2; // This might overflow, but it's the highest part.
|
||||
|
||||
let low_64 = GoldilocksField(res0, res1);
|
||||
let high_64 = GoldilocksField(res2, res3);
|
||||
|
||||
return GoldilocksFieldPair(low_64, high_64);
|
||||
}
|
||||
|
||||
// Reduces a 128-bit number (represented as two GoldilocksField values) modulo P.
|
||||
// Based on plonky2's reduce128, adapted for WGSL u32.
|
||||
fn gf_reduce128(val_low: GoldilocksField, val_high: GoldilocksField) -> GoldilocksField {
|
||||
// x_lo is (val_low.high:val_low.low)
|
||||
// x_hi is (val_high.high:val_high.low)
|
||||
|
||||
// x_hi_hi = val_high.high
|
||||
// x_hi_lo = val_high.low
|
||||
|
||||
// t0 = x_lo - x_hi_hi
|
||||
var t0_low = val_low.low;
|
||||
var t0_high = val_low.high;
|
||||
var borrow_from_t0 = 0u;
|
||||
|
||||
// Subtract x_hi_hi (val_high.high) from x_lo (val_low.high:val_low.low)
|
||||
let sub_res_low = u32_sub_with_borrow(t0_low, val_high.high, 0u);
|
||||
t0_low = sub_res_low.x;
|
||||
borrow_from_t0 = sub_res_low.y; // Borrow from t0_high if t0_low < val_high.high
|
||||
|
||||
let sub_res_high = u32_sub_with_borrow(t0_high, 0u, borrow_from_t0); // t0_high - borrow
|
||||
t0_high = sub_res_high.x;
|
||||
borrow_from_t0 = sub_res_high.y; // Borrow from beyond 64 bits
|
||||
|
||||
// if borrow { t0 -= EPSILON; }
|
||||
if (borrow_from_t0 != 0u) {
|
||||
// This means t0 is negative. Add P to it.
|
||||
// P = (P_HIGH:P_LOW)
|
||||
let add_p_low = u32_add_with_carry(t0_low, P_LOW, 0u);
|
||||
t0_low = add_p_low.x;
|
||||
let add_p_high = u32_add_with_carry(t0_high, P_HIGH, add_p_low.y);
|
||||
t0_high = add_p_high.x;
|
||||
}
|
||||
|
||||
// t1 = x_hi_lo * EPSILON
|
||||
// x_hi_lo is val_high.low
|
||||
// t1 = val_high.low * EPSILON
|
||||
// This is `val_high.low * (2^32 - 1) = (val_high.low << 32) - val_high.low`
|
||||
var t1_low: u32;
|
||||
var t1_high: u32;
|
||||
if (val_high.low == 0u) {
|
||||
t1_low = 0u;
|
||||
t1_high = 0u;
|
||||
} else {
|
||||
t1_low = 0u - val_high.low; // This wraps around, equivalent to (2^32 - val_high.low)
|
||||
t1_high = val_high.low - 1u;
|
||||
}
|
||||
let t1 = GoldilocksField(t1_low, t1_high);
|
||||
|
||||
// t2 = t0 + t1
|
||||
var result_low = t0_low;
|
||||
var result_high = t0_high;
|
||||
|
||||
let add_t1_low = u32_add_with_carry(result_low, t1.low, 0u);
|
||||
result_low = add_t1_low.x;
|
||||
let add_t1_high = u32_add_with_carry(result_high, t1.high, add_t1_low.y);
|
||||
result_high = add_t1_high.x;
|
||||
let final_carry = add_t1_high.y;
|
||||
|
||||
var final_result = GoldilocksField(result_low, result_high);
|
||||
|
||||
// If there's a final_carry, it means we effectively added 2^64, so add EPSILON.
|
||||
if (final_carry != 0u) {
|
||||
let add_epsilon_res = u32_add_with_carry(final_result.low, EPSILON, 0u);
|
||||
final_result.low = add_epsilon_res.x;
|
||||
final_result.high = final_result.high + add_epsilon_res.y;
|
||||
}
|
||||
|
||||
// Final reduction: if result >= P, subtract P.
|
||||
if (final_result.high > P_HIGH || (final_result.high == P_HIGH && final_result.low >= P_LOW)) {
|
||||
let sub_p_res = u32_sub_with_borrow(final_result.low, P_LOW, 0u);
|
||||
final_result.low = sub_p_res.x;
|
||||
final_result.high = final_result.high - P_HIGH - sub_p_res.y;
|
||||
}
|
||||
return final_result;
|
||||
}
|
||||
|
||||
// Proper Goldilocks field addition for all values
|
||||
fn gf_add(a: GoldilocksField, b: GoldilocksField) -> GoldilocksField {
|
||||
var sum_low_res = u32_add_with_carry(a.low, b.low, 0u);
|
||||
var sum_high_res = u32_add_with_carry(a.high, b.high, sum_low_res.y);
|
||||
|
||||
var result = GoldilocksField(sum_low_res.x, sum_high_res.x);
|
||||
var carry_high = sum_high_res.y;
|
||||
|
||||
// If there's a carry_high, it means the sum is >= 2^64.
|
||||
// In Goldilocks field, 2^64 = 2^32 - 1 (mod P).
|
||||
// So if carry_high is 1, we effectively add EPSILON (2^32 - 1) to the result.
|
||||
if (carry_high != 0u) {
|
||||
let add_epsilon_res = u32_add_with_carry(result.low, EPSILON, 0u);
|
||||
result.low = add_epsilon_res.x;
|
||||
result.high = result.high + add_epsilon_res.y; // This could cause another carry
|
||||
}
|
||||
|
||||
// Final reduction: if result >= P, subtract P.
|
||||
if (result.high > P_HIGH || (result.high == P_HIGH && result.low >= P_LOW)) {
|
||||
let sub_p_res = u32_sub_with_borrow(result.low, P_LOW, 0u);
|
||||
result.low = sub_p_res.x;
|
||||
result.high = result.high - P_HIGH - sub_p_res.y;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// Proper Goldilocks field multiplication for all values
|
||||
fn gf_mul(a: GoldilocksField, b: GoldilocksField) -> GoldilocksField {
|
||||
// Handle special cases first
|
||||
if (a.high == 0u && a.low == 0u) { return gf_zero(); }
|
||||
if (b.high == 0u && b.low == 0u) { return gf_zero(); }
|
||||
if (a.high == 0u && a.low == 1u) { return b; }
|
||||
if (b.high == 0u && b.low == 1u) { return a; }
|
||||
|
||||
// Step 1: Handle small values directly (both high parts are 0)
|
||||
if (a.high == 0u && b.high == 0u) {
|
||||
let product_64 = u32_mul_extended(a.low, b.low);
|
||||
return gf_reduce128(product_64, gf_zero());
|
||||
}
|
||||
|
||||
// Step 2: Handle mixed cases (one has high=0, other has high!=0)
|
||||
if (a.high == 0u || b.high == 0u) {
|
||||
// Multiply small * large using shift-and-add
|
||||
var large_val: GoldilocksField;
|
||||
var small_val: u32;
|
||||
|
||||
if (a.high == 0u) {
|
||||
large_val = b;
|
||||
small_val = a.low;
|
||||
} else {
|
||||
large_val = a;
|
||||
small_val = b.low;
|
||||
}
|
||||
|
||||
var result = gf_zero();
|
||||
var power_of_two = large_val;
|
||||
var remaining = small_val;
|
||||
|
||||
// Binary multiplication: decompose small_val into powers of 2
|
||||
for (var bit = 0u; bit < 32u; bit++) {
|
||||
if ((remaining & 1u) != 0u) {
|
||||
result = gf_add(result, power_of_two);
|
||||
}
|
||||
remaining = remaining >> 1u;
|
||||
if (remaining == 0u) { break; }
|
||||
power_of_two = gf_add(power_of_two, power_of_two); // double
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// Step 3: General case (both a.high != 0u and b.high != 0u)
|
||||
let product_128 = gf_mul_u64_by_u64(a, b);
|
||||
return gf_reduce128(product_128.first, product_128.second);
|
||||
}
|
||||
|
||||
@group(0) @binding(0) var<storage, read> input_a: array<GoldilocksField>;
|
||||
@group(0) @binding(1) var<storage, read> input_b: array<GoldilocksField>;
|
||||
@group(0) @binding(2) var<storage, read_write> output: array<GoldilocksField>;
|
||||
|
||||
@compute @workgroup_size(1)
|
||||
fn gf_mul_test(@builtin(global_invocation_id) global_id: vec3<u32>) {
|
||||
let i = global_id.x;
|
||||
output[i] = gf_mul(input_a[i], input_b[i]);
|
||||
}
|
||||
@@ -2,6 +2,8 @@ use bytemuck;
|
||||
use futures::executor::block_on;
|
||||
use wgpu::{self, util::DeviceExt};
|
||||
|
||||
mod tests;
|
||||
|
||||
// Extract Poseidon2 constants and generate WGSL code
|
||||
fn generate_wgsl_constants() {
|
||||
println!("Extracting Poseidon2 constants for WGSL...");
|
||||
@@ -785,5 +787,10 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Run gf_mul tests
|
||||
if let Err(e) = tests::test_gf_mul(&device, &queue).await {
|
||||
eprintln!("gf_mul tests failed: {}", e);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1,154 +0,0 @@
|
||||
// Poseidon2 Constants for WGSL
|
||||
// Extracted from qp-poseidon-constants
|
||||
// Width=12, External rounds: 4+4, Internal rounds: 22
|
||||
|
||||
// Internal round constants (22 values)
|
||||
const INTERNAL_CONSTANTS: array<array<u32, 2>, 22> = array<array<u32, 2>, 22>(
|
||||
array<u32, 2>(2018170979u, 2549578122u),
|
||||
array<u32, 2>(794875120u, 3520249608u),
|
||||
array<u32, 2>(2677723654u, 1772320679u),
|
||||
array<u32, 2>(2743438884u, 2849007878u),
|
||||
array<u32, 2>(518907317u, 693269760u),
|
||||
array<u32, 2>(293328710u, 1484055617u),
|
||||
array<u32, 2>(2834138828u, 2315799483u),
|
||||
array<u32, 2>(1558078501u, 1039128420u),
|
||||
array<u32, 2>(2266808631u, 966316006u),
|
||||
array<u32, 2>(3393728842u, 1045622667u),
|
||||
array<u32, 2>(2245828300u, 2521440415u),
|
||||
array<u32, 2>(751064958u, 1070374632u),
|
||||
array<u32, 2>(3490278765u, 2390340773u),
|
||||
array<u32, 2>(3526960470u, 2224174634u),
|
||||
array<u32, 2>(639988950u, 4000511088u),
|
||||
array<u32, 2>(1839350858u, 504240201u),
|
||||
array<u32, 2>(559852230u, 255489215u),
|
||||
array<u32, 2>(2713771731u, 453385078u),
|
||||
array<u32, 2>(1745082278u, 422331096u),
|
||||
array<u32, 2>(2544763488u, 4141129721u),
|
||||
array<u32, 2>(2700752774u, 1052996327u),
|
||||
array<u32, 2>(4063512019u, 1429786100u)
|
||||
);
|
||||
|
||||
// Initial external round constants (4 rounds x 12 elements)
|
||||
const INITIAL_EXTERNAL_CONSTANTS: array<array<array<u32, 2>, 12>, 4> = array<array<array<u32, 2>, 12>, 4>(
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(2539329031u, 3221415792u),
|
||||
array<u32, 2>(4262746426u, 3164936845u),
|
||||
array<u32, 2>(3883202553u, 1922272763u),
|
||||
array<u32, 2>(3761386668u, 3841130025u),
|
||||
array<u32, 2>(1411081289u, 3588274735u),
|
||||
array<u32, 2>(4090250945u, 3962812520u),
|
||||
array<u32, 2>(1100826458u, 1215155029u),
|
||||
array<u32, 2>(1489773809u, 1813820067u),
|
||||
array<u32, 2>(2585015995u, 3824356688u),
|
||||
array<u32, 2>(2378857513u, 3651555078u),
|
||||
array<u32, 2>(2864423342u, 3852156759u),
|
||||
array<u32, 2>(1531416540u, 708695120u)
|
||||
),
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(1987505445u, 2913073372u),
|
||||
array<u32, 2>(630903574u, 655361351u),
|
||||
array<u32, 2>(3412085911u, 1258046621u),
|
||||
array<u32, 2>(1456979578u, 1461113191u),
|
||||
array<u32, 2>(523722014u, 526769230u),
|
||||
array<u32, 2>(925368168u, 922771817u),
|
||||
array<u32, 2>(4074853328u, 3855135279u),
|
||||
array<u32, 2>(273563309u, 4248797356u),
|
||||
array<u32, 2>(1762266526u, 3450728622u),
|
||||
array<u32, 2>(1115336254u, 1107677022u),
|
||||
array<u32, 2>(4174699389u, 3986946237u),
|
||||
array<u32, 2>(3534029317u, 3543582418u)
|
||||
),
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(2101518210u, 2109009407u),
|
||||
array<u32, 2>(2323647532u, 2311068777u),
|
||||
array<u32, 2>(3016422935u, 3015476255u),
|
||||
array<u32, 2>(1378256883u, 3683616468u),
|
||||
array<u32, 2>(2029516952u, 2022063001u),
|
||||
array<u32, 2>(644616330u, 642915770u),
|
||||
array<u32, 2>(2580628271u, 2576506160u),
|
||||
array<u32, 2>(1689124307u, 1689124307u),
|
||||
array<u32, 2>(4016144568u, 4016127928u),
|
||||
array<u32, 2>(1335766254u, 1335740398u),
|
||||
array<u32, 2>(1465316391u, 1465316391u),
|
||||
array<u32, 2>(4119764157u, 4119617533u)
|
||||
),
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(507988087u, 507988087u),
|
||||
array<u32, 2>(1650295309u, 1650295309u),
|
||||
array<u32, 2>(2867864750u, 2867864750u),
|
||||
array<u32, 2>(1679868924u, 1679868924u),
|
||||
array<u32, 2>(999842784u, 999842784u),
|
||||
array<u32, 2>(603209949u, 603209949u),
|
||||
array<u32, 2>(2869577370u, 2869577370u),
|
||||
array<u32, 2>(3936230090u, 3936230090u),
|
||||
array<u32, 2>(3435906572u, 3435906572u),
|
||||
array<u32, 2>(2433830883u, 2433830883u),
|
||||
array<u32, 2>(1537056815u, 1537056815u),
|
||||
array<u32, 2>(3757386231u, 3757386231u)
|
||||
)
|
||||
);
|
||||
|
||||
// Terminal external round constants (4 rounds x 12 elements)
|
||||
const TERMINAL_EXTERNAL_CONSTANTS: array<array<array<u32, 2>, 12>, 4> = array<array<array<u32, 2>, 12>, 4>(
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(2067639406u, 2067639406u),
|
||||
array<u32, 2>(342838134u, 342838134u),
|
||||
array<u32, 2>(2148923528u, 2148923528u),
|
||||
array<u32, 2>(1836351170u, 1836351170u),
|
||||
array<u32, 2>(2714619123u, 2714619123u),
|
||||
array<u32, 2>(4142963247u, 4142963247u),
|
||||
array<u32, 2>(884199780u, 884199780u),
|
||||
array<u32, 2>(2970893770u, 2970893770u),
|
||||
array<u32, 2>(1697177254u, 1697177254u),
|
||||
array<u32, 2>(249070999u, 249070999u),
|
||||
array<u32, 2>(485491595u, 485491595u),
|
||||
array<u32, 2>(1718641338u, 1718641338u)
|
||||
),
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(812285441u, 1616511558u),
|
||||
array<u32, 2>(1074098767u, 1073693711u),
|
||||
array<u32, 2>(3896701389u, 3896638389u),
|
||||
array<u32, 2>(4169265998u, 4169265934u),
|
||||
array<u32, 2>(4095009929u, 4095009865u),
|
||||
array<u32, 2>(1835728118u, 1835728054u),
|
||||
array<u32, 2>(4176537227u, 4176537163u),
|
||||
array<u32, 2>(1478766714u, 1478766650u),
|
||||
array<u32, 2>(376881709u, 376881645u),
|
||||
array<u32, 2>(555692597u, 555692533u),
|
||||
array<u32, 2>(2968574966u, 2968574902u),
|
||||
array<u32, 2>(3635110935u, 3635110871u)
|
||||
),
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(2545243919u, 3830388836u),
|
||||
array<u32, 2>(1989298984u, 1989298920u),
|
||||
array<u32, 2>(151916301u, 151916237u),
|
||||
array<u32, 2>(2556322548u, 1995322484u),
|
||||
array<u32, 2>(238805938u, 238805874u),
|
||||
array<u32, 2>(2490398575u, 2490398511u),
|
||||
array<u32, 2>(2566222314u, 2566222250u),
|
||||
array<u32, 2>(814653888u, 3828870544u),
|
||||
array<u32, 2>(522453199u, 522453135u),
|
||||
array<u32, 2>(188424356u, 188424292u),
|
||||
array<u32, 2>(1649066023u, 1649065959u),
|
||||
array<u32, 2>(1595705478u, 1595705414u)
|
||||
),
|
||||
array<array<u32, 2>, 12>(
|
||||
array<u32, 2>(2466398251u, 4275150062u),
|
||||
array<u32, 2>(1928488963u, 1928488899u),
|
||||
array<u32, 2>(3871900792u, 3871900728u),
|
||||
array<u32, 2>(1463617809u, 1463617745u),
|
||||
array<u32, 2>(4175880892u, 4175880828u),
|
||||
array<u32, 2>(752107605u, 752107541u),
|
||||
array<u32, 2>(2101968497u, 2101968433u),
|
||||
array<u32, 2>(2239882257u, 2239882193u),
|
||||
array<u32, 2>(1389589688u, 1389589624u),
|
||||
array<u32, 2>(1344537748u, 1344537684u),
|
||||
array<u32, 2>(3172231929u, 3172231865u),
|
||||
array<u32, 2>(1780012361u, 3777952458u)
|
||||
)
|
||||
);
|
||||
|
||||
// Helper function to create GoldilocksField from constant array
|
||||
fn gf_from_const(val: array<u32, 2>) -> GoldilocksField {
|
||||
return GoldilocksField(val[0], val[1]);
|
||||
}
|
||||
425
crates/engine-gpu/src/tests.rs
Normal file
425
crates/engine-gpu/src/tests.rs
Normal file
@@ -0,0 +1,425 @@
|
||||
use plonky2::field::goldilocks_field::GoldilocksField;
|
||||
use plonky2::field::types::{Field, Field64, PrimeField64};
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
use wgpu::util::DeviceExt;
|
||||
|
||||
// A simple struct to hold a test case for gf_mul
|
||||
struct GfMulTestCase {
|
||||
a: GoldilocksField,
|
||||
b: GoldilocksField,
|
||||
expected: GoldilocksField,
|
||||
}
|
||||
|
||||
// Represents the GoldilocksField in a WGSL-compatible format (two u32s)
|
||||
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
|
||||
#[repr(C)]
|
||||
struct GfWgls {
|
||||
low: u32,
|
||||
high: u32,
|
||||
}
|
||||
|
||||
impl From<GoldilocksField> for GfWgls {
|
||||
fn from(gf: GoldilocksField) -> Self {
|
||||
let val = gf.0;
|
||||
Self {
|
||||
low: val as u32,
|
||||
high: (val >> 32) as u32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_gf_mul_test_vectors() -> Vec<GfMulTestCase> {
|
||||
println!("Generating gf_mul test vectors...");
|
||||
let mut vectors = Vec::new();
|
||||
|
||||
// Basic edge cases
|
||||
println!("Adding basic edge cases...");
|
||||
|
||||
// Case 1: 0 * x = 0
|
||||
let a1 = GoldilocksField::ZERO;
|
||||
let b1 = GoldilocksField::from_canonical_u64(123456789);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a1,
|
||||
b: b1,
|
||||
expected: a1 * b1,
|
||||
});
|
||||
|
||||
// Case 2: 1 * x = x
|
||||
let a2 = GoldilocksField::ONE;
|
||||
let b2 = GoldilocksField::from_canonical_u64(987654321);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a2,
|
||||
b: b2,
|
||||
expected: a2 * b2,
|
||||
});
|
||||
|
||||
// Case 3: x * 1 = x (commutative check)
|
||||
let a3 = GoldilocksField::from_noncanonical_u64(0xDEADBEEFCAFEBABE);
|
||||
let b3 = GoldilocksField::ONE;
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a3,
|
||||
b: b3,
|
||||
expected: a3 * b3,
|
||||
});
|
||||
|
||||
// Case 4: Small numbers (no overflow in u32*u32)
|
||||
let a4 = GoldilocksField::from_canonical_u64(100);
|
||||
let b4 = GoldilocksField::from_canonical_u64(200);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a4,
|
||||
b: b4,
|
||||
expected: a4 * b4,
|
||||
});
|
||||
|
||||
// Boundary cases
|
||||
println!("Adding boundary cases...");
|
||||
|
||||
// Case 5: u32::MAX
|
||||
let a5 = GoldilocksField::from_canonical_u64(u32::MAX as u64);
|
||||
let b5 = GoldilocksField::from_canonical_u64(u32::MAX as u64);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a5,
|
||||
b: b5,
|
||||
expected: a5 * b5,
|
||||
});
|
||||
|
||||
// Case 6: Just above u32::MAX
|
||||
let a6 = GoldilocksField::from_canonical_u64(u32::MAX as u64 + 1);
|
||||
let b6 = GoldilocksField::from_canonical_u64(2);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a6,
|
||||
b: b6,
|
||||
expected: a6 * b6,
|
||||
});
|
||||
|
||||
// Case 7: Powers of 2
|
||||
let a7 = GoldilocksField::from_canonical_u64(1u64 << 32);
|
||||
let b7 = GoldilocksField::from_canonical_u64(1u64 << 31);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a7,
|
||||
b: b7,
|
||||
expected: a7 * b7,
|
||||
});
|
||||
|
||||
// Case 8: Large powers of 2
|
||||
let a8 = GoldilocksField::from_canonical_u64(1u64 << 62);
|
||||
let b8 = GoldilocksField::from_canonical_u64(1u64 << 1);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a8,
|
||||
b: b8,
|
||||
expected: a8 * b8,
|
||||
});
|
||||
|
||||
// Field modulus edge cases
|
||||
println!("Adding field modulus edge cases...");
|
||||
|
||||
// Case 9: Field modulus - 1
|
||||
let a9 = GoldilocksField::from_canonical_u64(GoldilocksField::ORDER - 1);
|
||||
let b9 = GoldilocksField::from_canonical_u64(2);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a9,
|
||||
b: b9,
|
||||
expected: a9 * b9,
|
||||
});
|
||||
|
||||
// Case 10: Field modulus - 1 squared
|
||||
let a10 = GoldilocksField::from_canonical_u64(GoldilocksField::ORDER - 1);
|
||||
let b10 = GoldilocksField::from_canonical_u64(GoldilocksField::ORDER - 1);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a10,
|
||||
b: b10,
|
||||
expected: a10 * b10,
|
||||
});
|
||||
|
||||
// Case 11: Large numbers that will cause reduction
|
||||
let a11 = GoldilocksField::from_noncanonical_u64(0xABCDEF1234567890);
|
||||
let b11 = GoldilocksField::from_noncanonical_u64(0x1122334455667788);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a11,
|
||||
b: b11,
|
||||
expected: a11 * b11,
|
||||
});
|
||||
|
||||
// Mixed size cases
|
||||
println!("Adding mixed size cases...");
|
||||
|
||||
// Case 12: Small * Large
|
||||
let a12 = GoldilocksField::from_canonical_u64(3);
|
||||
let b12 = GoldilocksField::from_noncanonical_u64(0xFFFFFFFFFFFFFFF0);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a12,
|
||||
b: b12,
|
||||
expected: a12 * b12,
|
||||
});
|
||||
|
||||
// Case 13: Large * Small
|
||||
let a13 = GoldilocksField::from_noncanonical_u64(0xFFFFFFFFFFFFFFF0);
|
||||
let b13 = GoldilocksField::from_canonical_u64(7);
|
||||
vectors.push(GfMulTestCase {
|
||||
a: a13,
|
||||
b: b13,
|
||||
expected: a13 * b13,
|
||||
});
|
||||
|
||||
// Random test cases
|
||||
println!("Adding random test cases...");
|
||||
|
||||
// Use a fixed seed for reproducible tests
|
||||
let mut rng = ChaCha8Rng::seed_from_u64(0x123456789ABCDEF0);
|
||||
|
||||
// Generate 50 random test cases
|
||||
for _ in 0..50 {
|
||||
// Use from_noncanonical_u64 to handle any u64 value safely
|
||||
let a_val = rng.gen::<u64>();
|
||||
let b_val = rng.gen::<u64>();
|
||||
|
||||
let a = GoldilocksField::from_noncanonical_u64(a_val);
|
||||
let b = GoldilocksField::from_noncanonical_u64(b_val);
|
||||
|
||||
vectors.push(GfMulTestCase {
|
||||
a,
|
||||
b,
|
||||
expected: a * b,
|
||||
});
|
||||
}
|
||||
|
||||
// Special random cases focusing on problematic ranges
|
||||
println!("Adding focused random cases...");
|
||||
|
||||
// Cases where one operand has high=0, other has high!=0
|
||||
for _ in 0..10 {
|
||||
let small_val = rng.gen::<u32>() as u64;
|
||||
// Generate a large value and use from_noncanonical_u64 to handle safely
|
||||
let large_val = rng.gen_range((1u64 << 32)..u64::MAX);
|
||||
|
||||
let a = GoldilocksField::from_canonical_u64(small_val);
|
||||
let b = GoldilocksField::from_noncanonical_u64(large_val);
|
||||
|
||||
vectors.push(GfMulTestCase {
|
||||
a,
|
||||
b,
|
||||
expected: a * b,
|
||||
});
|
||||
}
|
||||
|
||||
// Cases where both operands are large (both high!=0)
|
||||
for _ in 0..10 {
|
||||
let a_val = rng.gen_range((1u64 << 32)..u64::MAX);
|
||||
let b_val = rng.gen_range((1u64 << 32)..u64::MAX);
|
||||
|
||||
let a = GoldilocksField::from_noncanonical_u64(a_val);
|
||||
let b = GoldilocksField::from_noncanonical_u64(b_val);
|
||||
|
||||
vectors.push(GfMulTestCase {
|
||||
a,
|
||||
b,
|
||||
expected: a * b,
|
||||
});
|
||||
}
|
||||
|
||||
// Cases near field boundaries
|
||||
for _ in 0..5 {
|
||||
let offset = rng.gen_range(1..1000);
|
||||
let near_max = GoldilocksField::ORDER - offset;
|
||||
let other = rng.gen_range(2..100);
|
||||
|
||||
let a = GoldilocksField::from_canonical_u64(near_max);
|
||||
let b = GoldilocksField::from_canonical_u64(other);
|
||||
|
||||
vectors.push(GfMulTestCase {
|
||||
a,
|
||||
b,
|
||||
expected: a * b,
|
||||
});
|
||||
}
|
||||
|
||||
println!("Generated {} total test vectors.", vectors.len());
|
||||
vectors
|
||||
}
|
||||
|
||||
pub async fn test_gf_mul(
|
||||
device: &wgpu::Device,
|
||||
queue: &wgpu::Queue,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
println!("\n--- Running gf_mul tests ---");
|
||||
|
||||
let test_vectors = generate_gf_mul_test_vectors();
|
||||
let total_tests = test_vectors.len();
|
||||
let mut passed_tests = 0;
|
||||
let mut failed_tests = Vec::new();
|
||||
|
||||
println!("Running {} test cases...", total_tests);
|
||||
|
||||
// Load the full mining shader code
|
||||
let shader_source = include_str!("gf_mul_test.wgsl");
|
||||
|
||||
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||||
label: Some("gf_mul test shader"),
|
||||
source: wgpu::ShaderSource::Wgsl(shader_source.into()),
|
||||
});
|
||||
|
||||
// Create pipeline
|
||||
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||||
label: Some("gf_mul test pipeline"),
|
||||
layout: None, // Let wgpu infer the layout
|
||||
module: &shader,
|
||||
entry_point: Some("gf_mul_test"),
|
||||
compilation_options: Default::default(),
|
||||
cache: None,
|
||||
});
|
||||
|
||||
let bind_group_layout = pipeline.get_bind_group_layout(0);
|
||||
|
||||
for (i, vector) in test_vectors.iter().enumerate() {
|
||||
let a_wgls: GfWgls = vector.a.into();
|
||||
let b_wgls: GfWgls = vector.b.into();
|
||||
|
||||
let input_a_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("Input A Buffer"),
|
||||
contents: bytemuck::cast_slice(&[a_wgls]),
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
});
|
||||
|
||||
let input_b_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("Input B Buffer"),
|
||||
contents: bytemuck::cast_slice(&[b_wgls]),
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
});
|
||||
|
||||
let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("Output Buffer"),
|
||||
size: std::mem::size_of::<GfWgls>() as u64,
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
|
||||
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: Some("Test Bind Group"),
|
||||
layout: &bind_group_layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 0,
|
||||
resource: input_a_buffer.as_entire_binding(),
|
||||
},
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 1,
|
||||
resource: input_b_buffer.as_entire_binding(),
|
||||
},
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 2,
|
||||
resource: output_buffer.as_entire_binding(),
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
|
||||
label: Some("Command Encoder"),
|
||||
});
|
||||
|
||||
{
|
||||
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
|
||||
label: Some("Compute Pass"),
|
||||
timestamp_writes: None,
|
||||
});
|
||||
pass.set_pipeline(&pipeline);
|
||||
pass.set_bind_group(0, &bind_group, &[]);
|
||||
pass.dispatch_workgroups(1, 1, 1);
|
||||
}
|
||||
|
||||
let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("Staging Buffer"),
|
||||
size: output_buffer.size(),
|
||||
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
|
||||
encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging_buffer, 0, output_buffer.size());
|
||||
queue.submit(Some(encoder.finish()));
|
||||
|
||||
let slice = staging_buffer.slice(..);
|
||||
slice.map_async(wgpu::MapMode::Read, |_| ());
|
||||
device
|
||||
.poll(wgpu::PollType::Wait {
|
||||
submission_index: None,
|
||||
timeout: None,
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let data = slice.get_mapped_range();
|
||||
let result_wgls: GfWgls = bytemuck::from_bytes::<GfWgls>(&data).clone();
|
||||
drop(data);
|
||||
staging_buffer.unmap();
|
||||
|
||||
let gpu_result_u64 = (result_wgls.high as u64) << 32 | (result_wgls.low as u64);
|
||||
let gpu_result = GoldilocksField(gpu_result_u64);
|
||||
|
||||
let expected_wgls: GfWgls = vector.expected.into();
|
||||
|
||||
// Progress indicator every 10 tests
|
||||
if i % 10 == 0 || i < 20 {
|
||||
println!(
|
||||
"Progress: {}/{} ({:.1}%)",
|
||||
i + 1,
|
||||
total_tests,
|
||||
(i + 1) as f32 / total_tests as f32 * 100.0
|
||||
);
|
||||
}
|
||||
|
||||
// The GPU result might not be canonical, so we need to canonicalize it before comparing.
|
||||
if gpu_result.to_canonical_u64() == vector.expected.to_canonical_u64() {
|
||||
passed_tests += 1;
|
||||
|
||||
// Only show details for first few tests or if verbose mode
|
||||
if i < 5 {
|
||||
println!("Test case {} ✅ PASSED", i + 1);
|
||||
println!(
|
||||
" a: 0x{:016x}, b: 0x{:016x} = 0x{:016x}",
|
||||
vector.a.0, vector.b.0, vector.expected.0
|
||||
);
|
||||
}
|
||||
} else {
|
||||
failed_tests.push(i + 1);
|
||||
println!("Test case {} ❌ FAILED", i + 1);
|
||||
println!(
|
||||
" a: 0x{:016x} ({}, {})",
|
||||
vector.a.0, a_wgls.low, a_wgls.high
|
||||
);
|
||||
println!(
|
||||
" b: 0x{:016x} ({}, {})",
|
||||
vector.b.0, b_wgls.low, b_wgls.high
|
||||
);
|
||||
println!(
|
||||
" CPU expected: 0x{:016x} ({}, {})",
|
||||
vector.expected.0, expected_wgls.low, expected_wgls.high
|
||||
);
|
||||
println!(
|
||||
" GPU result: 0x{:016x} ({}, {})",
|
||||
gpu_result.0, result_wgls.low, result_wgls.high
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Print final summary
|
||||
println!("\n=== TEST SUMMARY ===");
|
||||
println!("Total tests: {}", total_tests);
|
||||
println!(
|
||||
"Passed: {} ({:.1}%)",
|
||||
passed_tests,
|
||||
passed_tests as f32 / total_tests as f32 * 100.0
|
||||
);
|
||||
|
||||
if failed_tests.is_empty() {
|
||||
println!("🎉 All tests PASSED!");
|
||||
} else {
|
||||
println!(
|
||||
"Failed: {} ({:.1}%)",
|
||||
failed_tests.len(),
|
||||
failed_tests.len() as f32 / total_tests as f32 * 100.0
|
||||
);
|
||||
println!("Failed test cases: {:?}", failed_tests);
|
||||
return Err(format!("{} out of {} tests failed", failed_tests.len(), total_tests).into());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user