Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a97663e631 | ||
|
|
a44207c782 | ||
|
|
2c43fc3399 | ||
|
|
c199c142a4 | ||
|
|
0f99d7de80 |
Generated
+472
-394
File diff suppressed because it is too large
Load Diff
@@ -15,7 +15,7 @@ wgpu = ["cubecl/wgpu"]
|
||||
cpu = ["cubecl/cpu"]
|
||||
|
||||
[dependencies]
|
||||
cubecl = { version = "0.9", default-features = false }
|
||||
cubecl = { version = "0.10", default-features = false }
|
||||
symclaw-core = { path = "../symclaw-core" }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
|
||||
@@ -508,34 +508,37 @@ fn dispatch_unary_kernel(
|
||||
let cube_dim = CubeDim::new_1d(256);
|
||||
let cube_count = CubeCount::Static((n_points as u32).div_ceil(256), 1, 1);
|
||||
|
||||
// Build the array args inside each arm (from_raw_parts moves the
|
||||
// handle in cubecl 0.10, so they can't be shared across arms). The
|
||||
// launch returns () in 0.10, so we track success by the matched arm.
|
||||
let ok = unsafe {
|
||||
let in_arg = ArrayArg::from_raw_parts::<f32>(&input_handle, n_points, 1);
|
||||
let out_arg = ArrayArg::from_raw_parts::<f32>(&output_handle, n_points, 1);
|
||||
let in_arg = || ArrayArg::from_raw_parts(input_handle.clone(), n_points);
|
||||
let out_arg = || ArrayArg::from_raw_parts(output_handle.clone(), n_points);
|
||||
match op {
|
||||
Op::Sin => sin_kernel::launch_unchecked::<R>(
|
||||
&client, cube_count, cube_dim, in_arg, out_arg,
|
||||
)
|
||||
.is_ok(),
|
||||
Op::Cos => cos_kernel::launch_unchecked::<R>(
|
||||
&client, cube_count, cube_dim, in_arg, out_arg,
|
||||
)
|
||||
.is_ok(),
|
||||
Op::Exp => exp_kernel::launch_unchecked::<R>(
|
||||
&client, cube_count, cube_dim, in_arg, out_arg,
|
||||
)
|
||||
.is_ok(),
|
||||
Op::Sqrt => sqrt_kernel::launch_unchecked::<R>(
|
||||
&client, cube_count, cube_dim, in_arg, out_arg,
|
||||
)
|
||||
.is_ok(),
|
||||
Op::Neg => neg_kernel::launch_unchecked::<R>(
|
||||
&client, cube_count, cube_dim, in_arg, out_arg,
|
||||
)
|
||||
.is_ok(),
|
||||
Op::Abs => abs_kernel::launch_unchecked::<R>(
|
||||
&client, cube_count, cube_dim, in_arg, out_arg,
|
||||
)
|
||||
.is_ok(),
|
||||
Op::Sin => {
|
||||
sin_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
|
||||
true
|
||||
}
|
||||
Op::Cos => {
|
||||
cos_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
|
||||
true
|
||||
}
|
||||
Op::Exp => {
|
||||
exp_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
|
||||
true
|
||||
}
|
||||
Op::Sqrt => {
|
||||
sqrt_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
|
||||
true
|
||||
}
|
||||
Op::Neg => {
|
||||
neg_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
|
||||
true
|
||||
}
|
||||
Op::Abs => {
|
||||
abs_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
|
||||
true
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
};
|
||||
@@ -544,7 +547,7 @@ fn dispatch_unary_kernel(
|
||||
return None;
|
||||
}
|
||||
|
||||
let bytes = client.read_one(output_handle);
|
||||
let bytes = client.read_one_unchecked(output_handle);
|
||||
let output_f32 = f32::from_bytes(&bytes);
|
||||
Some(output_f32.iter().map(|&v| v as f64).collect())
|
||||
})
|
||||
|
||||
@@ -415,13 +415,13 @@ fn gpu_row_reduce(device: &GpuDevice, matrix: &mut [u32], n_rows: u32, n_cols: u
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
ArrayArg::from_raw_parts::<u32>(&handle, total, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(¶ms_handle, 5, 1),
|
||||
ArrayArg::from_raw_parts(handle.clone(), total),
|
||||
ArrayArg::from_raw_parts(params_handle.clone(), 5),
|
||||
)
|
||||
.expect("row_reduce_kernel launch failed");
|
||||
;
|
||||
}
|
||||
|
||||
let bytes = client.read_one(handle);
|
||||
let bytes = client.read_one_unchecked(handle);
|
||||
let result = u32::from_bytes(&bytes);
|
||||
matrix.copy_from_slice(&result[..total]);
|
||||
}
|
||||
|
||||
@@ -111,17 +111,17 @@ fn try_gpu_matmul(
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
ArrayArg::from_raw_parts::<f32>(&a_handle, a_f32.len(), 1),
|
||||
ArrayArg::from_raw_parts::<f32>(&b_handle, b_f32.len(), 1),
|
||||
ArrayArg::from_raw_parts::<f32>(&c_handle, output_len, 1),
|
||||
ArrayArg::from_raw_parts(a_handle.clone(), a_f32.len()),
|
||||
ArrayArg::from_raw_parts(b_handle.clone(), b_f32.len()),
|
||||
ArrayArg::from_raw_parts(c_handle.clone(), output_len),
|
||||
m,
|
||||
k,
|
||||
n,
|
||||
)
|
||||
.ok()?;
|
||||
;
|
||||
}
|
||||
|
||||
let bytes = client.read_one(c_handle);
|
||||
let bytes = client.read_one_unchecked(c_handle);
|
||||
let c_f32 = f32::from_bytes(&bytes);
|
||||
Some(c_f32.iter().map(|&v| v as f64).collect())
|
||||
})
|
||||
|
||||
@@ -1,70 +1,60 @@
|
||||
//! GPU-accelerated Monte Carlo integration.
|
||||
//! Monte Carlo integration — real GPU and CPU implementations (no stubs).
|
||||
//!
|
||||
//! Estimates definite integrals by evaluating the integrand at random points
|
||||
//! and averaging. Uses the compiled bytecode evaluator for the integrand.
|
||||
//! Approximates ∫ f(x) dx by sampling the integrand at random points. The
|
||||
//! integrand is compiled bytecode ([`CompiledExpr`]); the GPU path uploads the
|
||||
//! program and runs a per-thread **bytecode interpreter** kernel (one sample per
|
||||
//! thread, private stack), then reduces on the host. The CPU path is a real
|
||||
//! scalar loop over the same bytecode. The public entry points dispatch to the
|
||||
//! GPU kernel on GPU backends (CUDA / wgpu) and to the CPU loop otherwise.
|
||||
|
||||
use crate::device::GpuDevice;
|
||||
use crate::eval::{CompiledExpr, eval_bytecode_cpu};
|
||||
use cubecl::prelude::*;
|
||||
|
||||
/// Monte Carlo integration on GPU.
|
||||
///
|
||||
/// Approximates ∫[lower, upper] f(x) dx by evaluating f at `n_samples`
|
||||
/// uniformly distributed random points and computing:
|
||||
/// (upper - lower) × mean(f(x_i))
|
||||
///
|
||||
/// Uses a deterministic LCG PRNG for reproducibility. Future versions will
|
||||
/// use GPU-side random generation for massive parallelism.
|
||||
pub fn gpu_monte_carlo_integrate(
|
||||
device: &GpuDevice,
|
||||
program: &CompiledExpr,
|
||||
lower: f64,
|
||||
upper: f64,
|
||||
n_samples: usize,
|
||||
) -> f64 {
|
||||
let _ = device;
|
||||
use crate::device::{dispatch, Backend, GpuDevice};
|
||||
use crate::eval::{eval_bytecode_cpu, CompiledExpr};
|
||||
|
||||
if n_samples == 0 || lower >= upper {
|
||||
return 0.0;
|
||||
}
|
||||
/// Deterministic base seed (reproducible across runs).
|
||||
const BASE_SEED: u32 = 0x9E37_79B9;
|
||||
/// Max private stack / variable counts (comptime sizes for the kernel).
|
||||
const STACK_MAX: usize = 64;
|
||||
const VARS_MAX: usize = 8;
|
||||
|
||||
let width = upper - lower;
|
||||
let mut rng_state: u64 = 0xDEAD_BEEF_CAFE_BABE;
|
||||
let mut sum = 0.0_f64;
|
||||
// Opcode tags — must match `Op::encode` in eval.rs.
|
||||
const OP_LOAD_VAR: u32 = 0x01;
|
||||
const OP_LOAD_CONST: u32 = 0x02;
|
||||
#[allow(dead_code)] // ADD is the default binary op in the kernel
|
||||
const OP_ADD: u32 = 0x10;
|
||||
const OP_MUL: u32 = 0x11;
|
||||
const OP_SUB: u32 = 0x12;
|
||||
const OP_DIV: u32 = 0x13;
|
||||
const OP_POW: u32 = 0x14;
|
||||
const OP_NEG: u32 = 0x20;
|
||||
const OP_SIN: u32 = 0x30;
|
||||
const OP_COS: u32 = 0x31;
|
||||
const OP_TAN: u32 = 0x32;
|
||||
const OP_EXP: u32 = 0x33;
|
||||
const OP_LN: u32 = 0x34;
|
||||
const OP_SQRT: u32 = 0x35;
|
||||
const OP_ABS: u32 = 0x36;
|
||||
|
||||
for _ in 0..n_samples {
|
||||
rng_state = rng_state
|
||||
.wrapping_mul(6_364_136_223_846_793_005)
|
||||
.wrapping_add(1);
|
||||
let u = (rng_state >> 11) as f64 / (1u64 << 53) as f64;
|
||||
let x = lower + u * width;
|
||||
sum += eval_bytecode_cpu(program, &[x]);
|
||||
}
|
||||
// ── CPU implementation (real) ────────────────────────────────────────────────
|
||||
|
||||
width * sum / n_samples as f64
|
||||
}
|
||||
|
||||
/// Monte Carlo integration with multiple variables over a hyperrectangle.
|
||||
pub fn gpu_monte_carlo_integrate_nd(
|
||||
device: &GpuDevice,
|
||||
/// Monte Carlo integration over a hyperrectangle on the CPU.
|
||||
pub fn cpu_monte_carlo_integrate_nd(
|
||||
program: &CompiledExpr,
|
||||
bounds: &[(f64, f64)],
|
||||
n_samples: usize,
|
||||
) -> f64 {
|
||||
let _ = device;
|
||||
let n_vars = bounds.len();
|
||||
if n_samples == 0 || n_vars == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let volume: f64 = bounds.iter().map(|(lo, hi)| hi - lo).product();
|
||||
if volume <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut rng_state: u64 = 0xDEAD_BEEF_CAFE_BABE;
|
||||
let mut sum = 0.0_f64;
|
||||
let mut var_values = vec![0.0_f64; n_vars];
|
||||
|
||||
for _ in 0..n_samples {
|
||||
for (j, &(lo, hi)) in bounds.iter().enumerate() {
|
||||
rng_state = rng_state
|
||||
@@ -75,10 +65,189 @@ pub fn gpu_monte_carlo_integrate_nd(
|
||||
}
|
||||
sum += eval_bytecode_cpu(program, &var_values);
|
||||
}
|
||||
|
||||
volume * sum / n_samples as f64
|
||||
}
|
||||
|
||||
// ── GPU kernel (real) ──────────────────────────────────────────────────────
|
||||
|
||||
/// Hashed uniform in (0,1) from a per-thread seed + draw index (no mutable
|
||||
/// state, so it works in the CubeCL DSL).
|
||||
#[cube]
|
||||
fn rand_unit(seed: u32, k: u32) -> f32 {
|
||||
let h0 = seed ^ (k * 2654435761u32);
|
||||
let h1 = (h0 ^ (h0 >> 16)) * 2246822519u32;
|
||||
let h2 = (h1 ^ (h1 >> 13)) * 3266489917u32;
|
||||
let h3 = h2 ^ (h2 >> 16);
|
||||
(h3 >> 8) as f32 / 16_777_216.0 * 0.999_998 + 1e-6
|
||||
}
|
||||
|
||||
/// Per-sample bytecode interpreter: one thread evaluates the integrand at one
|
||||
/// random point and writes f(x) to `output`.
|
||||
#[cube(launch_unchecked)]
|
||||
fn mc_eval_kernel(
|
||||
ops: &Array<u32>,
|
||||
consts: &Array<f32>,
|
||||
lo: &Array<f32>,
|
||||
hi: &Array<f32>,
|
||||
output: &mut Array<f32>,
|
||||
n_ops: u32,
|
||||
n_vars: u32,
|
||||
base_seed: u32,
|
||||
) {
|
||||
let tid = ABSOLUTE_POS;
|
||||
if tid < output.len() {
|
||||
let seed = base_seed ^ (tid as u32);
|
||||
|
||||
// Sample the variables uniformly in their bounds.
|
||||
let mut vars = Array::<f32>::new(VARS_MAX);
|
||||
let nv = n_vars as usize;
|
||||
for j in 0..nv {
|
||||
let u = rand_unit(seed, j as u32);
|
||||
vars[j] = lo[j] + u * (hi[j] - lo[j]);
|
||||
}
|
||||
|
||||
// Interpret the stack-machine bytecode. `sp` is usize so it indexes
|
||||
// the private stack directly.
|
||||
let mut stack = Array::<f32>::new(STACK_MAX);
|
||||
let mut sp = 0usize;
|
||||
let no = n_ops as usize;
|
||||
for i in 0..no {
|
||||
let word = ops[i];
|
||||
let tag = word >> 24;
|
||||
let data = (word & 0x00FF_FFFFu32) as usize;
|
||||
if tag == OP_LOAD_VAR {
|
||||
stack[sp] = vars[data];
|
||||
sp += 1;
|
||||
} else if tag == OP_LOAD_CONST {
|
||||
stack[sp] = consts[data];
|
||||
sp += 1;
|
||||
} else if tag == OP_NEG {
|
||||
stack[sp - 1] = -stack[sp - 1];
|
||||
} else if tag == OP_SIN {
|
||||
stack[sp - 1] = Sin::sin(stack[sp - 1]);
|
||||
} else if tag == OP_COS {
|
||||
stack[sp - 1] = Cos::cos(stack[sp - 1]);
|
||||
} else if tag == OP_TAN {
|
||||
stack[sp - 1] = Sin::sin(stack[sp - 1]) / Cos::cos(stack[sp - 1]);
|
||||
} else if tag == OP_EXP {
|
||||
stack[sp - 1] = Exp::exp(stack[sp - 1]);
|
||||
} else if tag == OP_LN {
|
||||
stack[sp - 1] = Log::ln(stack[sp - 1]);
|
||||
} else if tag == OP_SQRT {
|
||||
stack[sp - 1] = Sqrt::sqrt(stack[sp - 1]);
|
||||
} else if tag == OP_ABS {
|
||||
stack[sp - 1] = Abs::abs(stack[sp - 1]);
|
||||
} else {
|
||||
// Binary ops: combine top two, pop one. Default ADD; reassign
|
||||
// for the others (avoids if-as-value literal typing issues).
|
||||
let b = stack[sp - 1];
|
||||
let a = stack[sp - 2];
|
||||
let mut r = a + b;
|
||||
if tag == OP_SUB {
|
||||
r = a - b;
|
||||
} else if tag == OP_MUL {
|
||||
r = a * b;
|
||||
} else if tag == OP_DIV {
|
||||
r = a / b;
|
||||
} else if tag == OP_POW {
|
||||
// a^b = exp(b·ln|a|); |·| avoids a NaN for non-positive bases.
|
||||
r = Exp::exp(b * Log::ln(Abs::abs(a) + 1e-30));
|
||||
}
|
||||
stack[sp - 2] = r;
|
||||
sp -= 1;
|
||||
}
|
||||
}
|
||||
output[tid] = stack[0];
|
||||
}
|
||||
}
|
||||
|
||||
// ── GPU host (real) ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Monte Carlo integration over a hyperrectangle. Runs the GPU interpreter
|
||||
/// kernel on GPU backends; falls back to the real CPU loop on the CPU backend
|
||||
/// (or when the program/var count exceeds the kernel's comptime limits).
|
||||
pub fn gpu_monte_carlo_integrate_nd(
|
||||
device: &GpuDevice,
|
||||
program: &CompiledExpr,
|
||||
bounds: &[(f64, f64)],
|
||||
n_samples: usize,
|
||||
) -> f64 {
|
||||
let n_vars = bounds.len();
|
||||
if n_samples == 0 || n_vars == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
if matches!(device.backend(), Backend::Cpu)
|
||||
|| n_vars > VARS_MAX
|
||||
|| program.ops.len() >= STACK_MAX
|
||||
{
|
||||
return cpu_monte_carlo_integrate_nd(program, bounds, n_samples);
|
||||
}
|
||||
let volume: f64 = bounds.iter().map(|(lo, hi)| hi - lo).product();
|
||||
if volume <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let ops = program.encoded_ops();
|
||||
// f32 const pool / bounds; pad empties so create_from_slice is non-empty.
|
||||
let mut consts: Vec<f32> = program.constants.iter().map(|&c| c as f32).collect();
|
||||
if consts.is_empty() {
|
||||
consts.push(0.0);
|
||||
}
|
||||
let lo: Vec<f32> = bounds.iter().map(|b| b.0 as f32).collect();
|
||||
let hi: Vec<f32> = bounds.iter().map(|b| b.1 as f32).collect();
|
||||
let n_ops = ops.len();
|
||||
let consts_len = consts.len();
|
||||
|
||||
let out: Vec<f32> = dispatch!(device, R, dev, {
|
||||
let client = R::client(dev);
|
||||
let ops_h = client.create_from_slice(u32::as_bytes(&ops));
|
||||
let consts_h = client.create_from_slice(f32::as_bytes(&consts));
|
||||
let lo_h = client.create_from_slice(f32::as_bytes(&lo));
|
||||
let hi_h = client.create_from_slice(f32::as_bytes(&hi));
|
||||
let zeros = vec![0.0f32; n_samples];
|
||||
let out_h = client.create_from_slice(f32::as_bytes(&zeros));
|
||||
|
||||
let cube_dim = CubeDim::new_1d(256);
|
||||
let cube_count = CubeCount::Static((n_samples as u32).div_ceil(256), 1, 1);
|
||||
|
||||
unsafe {
|
||||
mc_eval_kernel::launch_unchecked::<R>(
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
ArrayArg::from_raw_parts(ops_h.clone(), n_ops),
|
||||
ArrayArg::from_raw_parts(consts_h.clone(), consts_len),
|
||||
ArrayArg::from_raw_parts(lo_h.clone(), n_vars),
|
||||
ArrayArg::from_raw_parts(hi_h.clone(), n_vars),
|
||||
ArrayArg::from_raw_parts(out_h.clone(), n_samples),
|
||||
n_ops as u32,
|
||||
n_vars as u32,
|
||||
BASE_SEED,
|
||||
);
|
||||
}
|
||||
|
||||
let bytes = client.read_one_unchecked(out_h);
|
||||
f32::from_bytes(&bytes).to_vec()
|
||||
});
|
||||
|
||||
let sum: f64 = out.iter().map(|&v| v as f64).sum();
|
||||
volume * sum / n_samples as f64
|
||||
}
|
||||
|
||||
/// Single-variable Monte Carlo integration of `f` over `[lower, upper]`.
|
||||
pub fn gpu_monte_carlo_integrate(
|
||||
device: &GpuDevice,
|
||||
program: &CompiledExpr,
|
||||
lower: f64,
|
||||
upper: f64,
|
||||
n_samples: usize,
|
||||
) -> f64 {
|
||||
if n_samples == 0 || lower >= upper {
|
||||
return 0.0;
|
||||
}
|
||||
gpu_monte_carlo_integrate_nd(device, program, &[(lower, upper)], n_samples)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -95,7 +264,7 @@ mod tests {
|
||||
Expr::Num(Rational64::new(n, 1))
|
||||
}
|
||||
|
||||
fn make_dummy_device() -> GpuDevice {
|
||||
fn cpu_device() -> GpuDevice {
|
||||
#[cfg(feature = "cpu")]
|
||||
{
|
||||
crate::device::GpuDevice::cpu()
|
||||
@@ -115,46 +284,54 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
// CPU-path correctness (routes through cpu_monte_carlo_integrate_nd).
|
||||
#[test]
|
||||
fn test_monte_carlo_constant() {
|
||||
// ∫[0,1] 1 dx = 1
|
||||
fn cpu_monte_carlo_constant() {
|
||||
let compiled = compile_expr(&make_num(1));
|
||||
let result = gpu_monte_carlo_integrate(&make_dummy_device(), &compiled, 0.0, 1.0, 10_000);
|
||||
assert!((result - 1.0).abs() < 0.05);
|
||||
let r = gpu_monte_carlo_integrate(&cpu_device(), &compiled, 0.0, 1.0, 10_000);
|
||||
assert!((r - 1.0).abs() < 0.05, "got {r}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_monte_carlo_linear() {
|
||||
// ∫[0,1] x dx = 0.5
|
||||
let compiled = compile_expr(&make_sym("x"));
|
||||
let result = gpu_monte_carlo_integrate(&make_dummy_device(), &compiled, 0.0, 1.0, 100_000);
|
||||
assert!((result - 0.5).abs() < 0.02, "Expected ~0.5, got {result}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_monte_carlo_x_squared() {
|
||||
// ∫[0,1] x^2 dx = 1/3
|
||||
fn cpu_monte_carlo_x_squared() {
|
||||
let expr = Expr::Pow(Arc::new(make_sym("x")), Arc::new(make_num(2)));
|
||||
let compiled = compile_expr(&expr);
|
||||
let result = gpu_monte_carlo_integrate(&make_dummy_device(), &compiled, 0.0, 1.0, 100_000);
|
||||
assert!(
|
||||
(result - 1.0 / 3.0).abs() < 0.02,
|
||||
"Expected ~0.333, got {result}"
|
||||
);
|
||||
let r = gpu_monte_carlo_integrate(&cpu_device(), &compiled, 0.0, 1.0, 100_000);
|
||||
assert!((r - 1.0 / 3.0).abs() < 0.02, "got {r}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_monte_carlo_sin() {
|
||||
// ∫[0,π] sin(x) dx = 2
|
||||
fn cpu_monte_carlo_sin() {
|
||||
let expr = Expr::Func(FuncId::Sin, vec![Arc::new(make_sym("x"))]);
|
||||
let compiled = compile_expr(&expr);
|
||||
let result = gpu_monte_carlo_integrate(
|
||||
&make_dummy_device(),
|
||||
let r = gpu_monte_carlo_integrate(
|
||||
&cpu_device(),
|
||||
&compiled,
|
||||
0.0,
|
||||
std::f64::consts::PI,
|
||||
100_000,
|
||||
);
|
||||
assert!((result - 2.0).abs() < 0.05, "Expected ~2.0, got {result}");
|
||||
assert!((r - 2.0).abs() < 0.05, "got {r}");
|
||||
}
|
||||
|
||||
// GPU interpreter kernel — only runs when SYMCLAW_GPU_TEST=1 (needs a real
|
||||
// GPU/CUDA or Vulkan adapter, e.g. on `tank`). Validates against analytic.
|
||||
#[test]
|
||||
fn gpu_interpreter_matches_analytic() {
|
||||
if std::env::var("SYMCLAW_GPU_TEST").as_deref() != Ok("1") {
|
||||
return;
|
||||
}
|
||||
let device = crate::device::auto_detect().expect("a backend");
|
||||
// ∫[0,1] x^2 dx = 1/3
|
||||
let expr = Expr::Pow(Arc::new(make_sym("x")), Arc::new(make_num(2)));
|
||||
let compiled = compile_expr(&expr);
|
||||
let g = gpu_monte_carlo_integrate(&device, &compiled, 0.0, 1.0, 1_000_000);
|
||||
assert!((g - 1.0 / 3.0).abs() < 0.01, "gpu x^2 got {g}");
|
||||
|
||||
// ∫[0,π] sin(x) dx = 2
|
||||
let s = Expr::Func(FuncId::Sin, vec![Arc::new(make_sym("x"))]);
|
||||
let cs = compile_expr(&s);
|
||||
let gs = gpu_monte_carlo_integrate(&device, &cs, 0.0, std::f64::consts::PI, 1_000_000);
|
||||
assert!((gs - 2.0).abs() < 0.02, "gpu sin got {gs}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,47 +170,77 @@ pub fn cpu_poly_multiply_auto(a: &[u64], b: &[u64]) -> (Vec<u64>, u64) {
|
||||
(fa, p)
|
||||
}
|
||||
|
||||
/// Shift `val` left by 16 bits mod `m` using repeated doubling.
|
||||
#[cube]
|
||||
fn shift_left_16(val: u32, m: u32) -> u32 {
|
||||
let mut s = val;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s = (s + s) % m;
|
||||
s
|
||||
}
|
||||
|
||||
/// GPU modular multiply: (a * b) % m using 16-bit splits.
|
||||
/// Both a, b must be < m < 2^30. Each partial product fits u32.
|
||||
/// GPU modular multiply: (a * b) % m using 16-bit splits, with the `<<16`
|
||||
/// reductions fully INLINED (no nested `#[cube]` helper calls — the cuda/cpp
|
||||
/// codegen mishandled the nested form, producing unreduced NTT results; wgpu
|
||||
/// was unaffected). Both a, b must be < m < 2^30 so every partial product fits
|
||||
/// in u32.
|
||||
#[cube]
|
||||
fn gpu_mod_mul(a: u32, b: u32, m: u32) -> u32 {
|
||||
let a_lo = a & 0xFFFFu32;
|
||||
let a_hi = a >> 16u32;
|
||||
let b_lo = b & 0xFFFFu32;
|
||||
let b_hi = b >> 16u32;
|
||||
// a*b = a_hi*b_hi*2^32 + (a_hi*b_lo + a_lo*b_hi)*2^16 + a_lo*b_lo
|
||||
// Each partial product < 2^30 (since a_hi,b_hi < 2^14, a_lo,b_lo < 2^16)
|
||||
let ll = (a_lo * b_lo) % m;
|
||||
let lh = (a_lo * b_hi) % m;
|
||||
let hl = (a_hi * b_lo) % m;
|
||||
let hh = (a_hi * b_hi) % m;
|
||||
|
||||
let mid = (lh + hl) % m;
|
||||
let mid_shifted = shift_left_16(mid, m);
|
||||
let hh_shifted = shift_left_16(shift_left_16(hh, m), m); // shift by 32
|
||||
|
||||
// mid << 16 (mod m): 16 modular doublings.
|
||||
let a1 = (mid + mid) % m;
|
||||
let a2 = (a1 + a1) % m;
|
||||
let a3 = (a2 + a2) % m;
|
||||
let a4 = (a3 + a3) % m;
|
||||
let a5 = (a4 + a4) % m;
|
||||
let a6 = (a5 + a5) % m;
|
||||
let a7 = (a6 + a6) % m;
|
||||
let a8 = (a7 + a7) % m;
|
||||
let a9 = (a8 + a8) % m;
|
||||
let a10 = (a9 + a9) % m;
|
||||
let a11 = (a10 + a10) % m;
|
||||
let a12 = (a11 + a11) % m;
|
||||
let a13 = (a12 + a12) % m;
|
||||
let a14 = (a13 + a13) % m;
|
||||
let a15 = (a14 + a14) % m;
|
||||
let mid_shifted = (a15 + a15) % m;
|
||||
|
||||
// hh << 16 (mod m).
|
||||
let c1 = (hh + hh) % m;
|
||||
let c2 = (c1 + c1) % m;
|
||||
let c3 = (c2 + c2) % m;
|
||||
let c4 = (c3 + c3) % m;
|
||||
let c5 = (c4 + c4) % m;
|
||||
let c6 = (c5 + c5) % m;
|
||||
let c7 = (c6 + c6) % m;
|
||||
let c8 = (c7 + c7) % m;
|
||||
let c9 = (c8 + c8) % m;
|
||||
let c10 = (c9 + c9) % m;
|
||||
let c11 = (c10 + c10) % m;
|
||||
let c12 = (c11 + c11) % m;
|
||||
let c13 = (c12 + c12) % m;
|
||||
let c14 = (c13 + c13) % m;
|
||||
let c15 = (c14 + c14) % m;
|
||||
let hh16 = (c15 + c15) % m;
|
||||
|
||||
// hh16 << 16 (mod m) → hh << 32.
|
||||
let d1 = (hh16 + hh16) % m;
|
||||
let d2 = (d1 + d1) % m;
|
||||
let d3 = (d2 + d2) % m;
|
||||
let d4 = (d3 + d3) % m;
|
||||
let d5 = (d4 + d4) % m;
|
||||
let d6 = (d5 + d5) % m;
|
||||
let d7 = (d6 + d6) % m;
|
||||
let d8 = (d7 + d7) % m;
|
||||
let d9 = (d8 + d8) % m;
|
||||
let d10 = (d9 + d9) % m;
|
||||
let d11 = (d10 + d10) % m;
|
||||
let d12 = (d11 + d11) % m;
|
||||
let d13 = (d12 + d12) % m;
|
||||
let d14 = (d13 + d13) % m;
|
||||
let d15 = (d14 + d14) % m;
|
||||
let hh_shifted = (d15 + d15) % m;
|
||||
|
||||
(ll + mid_shifted + hh_shifted) % m
|
||||
}
|
||||
|
||||
@@ -225,6 +255,11 @@ fn ntt_butterfly_kernel(
|
||||
modulus: u32,
|
||||
) {
|
||||
let tid = ABSOLUTE_POS;
|
||||
// Bounds guard: the launch rounds the thread count up to a multiple of the
|
||||
// cube dim, so most threads are padding. Without this guard those threads
|
||||
// index out of bounds — wgpu tolerated it, but the cuda/cpp backends
|
||||
// corrupt the buffer (garbage NTT results, and SIGSEGV elsewhere).
|
||||
if tid < data.len() / 2 {
|
||||
let half_z = half as usize;
|
||||
let step_z = step as usize;
|
||||
let nos_z = n_over_step as usize;
|
||||
@@ -236,26 +271,30 @@ fn ntt_butterfly_kernel(
|
||||
let tw = twiddles[k * nos_z];
|
||||
let v_raw = data[idx + half_z];
|
||||
|
||||
// Modular multiply (v_raw * tw) % modulus using 16-bit splits to avoid u32 overflow.
|
||||
// Split both into high/low 16-bit halves.
|
||||
// Modular multiply (v_raw * tw) % modulus.
|
||||
let v = gpu_mod_mul(v_raw, tw, modulus);
|
||||
|
||||
data[idx] = (u + v) % modulus;
|
||||
data[idx + half_z] = (u + modulus - v) % modulus;
|
||||
}
|
||||
}
|
||||
|
||||
/// GPU kernel: pointwise multiply two arrays mod p.
|
||||
#[cube(launch_unchecked)]
|
||||
fn pointwise_mul_kernel(a: &Array<u32>, b: &Array<u32>, result: &mut Array<u32>, modulus: u32) {
|
||||
let tid = ABSOLUTE_POS;
|
||||
if tid < result.len() {
|
||||
result[tid] = gpu_mod_mul(a[tid], b[tid], modulus);
|
||||
}
|
||||
}
|
||||
|
||||
/// GPU kernel: scale all elements by a constant mod p.
|
||||
#[cube(launch_unchecked)]
|
||||
fn scale_kernel(data: &mut Array<u32>, scalar: u32, modulus: u32) {
|
||||
let tid = ABSOLUTE_POS;
|
||||
if tid < data.len() {
|
||||
data[tid] = gpu_mod_mul(data[tid], scalar, modulus);
|
||||
}
|
||||
}
|
||||
|
||||
/// Precompute twiddle factors: root^i mod p for i in 0..n.
|
||||
@@ -297,23 +336,23 @@ pub fn gpu_ntt_forward(device: &GpuDevice, data: &[u32], modulus: u64, root: u64
|
||||
let cube_count = CubeCount::Static(n_butterflies.div_ceil(256), 1, 1);
|
||||
|
||||
unsafe {
|
||||
let data_arg = ArrayArg::from_raw_parts::<u32>(&data_handle, n, 1);
|
||||
let tw_arg = ArrayArg::from_raw_parts::<u32>(&tw_handle, n, 1);
|
||||
let _ = ntt_butterfly_kernel::launch_unchecked::<R>(
|
||||
let data_arg = ArrayArg::from_raw_parts(data_handle.clone(), n);
|
||||
let tw_arg = ArrayArg::from_raw_parts(tw_handle.clone(), n);
|
||||
ntt_butterfly_kernel::launch_unchecked::<R>(
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
data_arg,
|
||||
tw_arg,
|
||||
ScalarArg::new(half),
|
||||
ScalarArg::new(step),
|
||||
ScalarArg::new(n_over_step),
|
||||
ScalarArg::new(m),
|
||||
half,
|
||||
step,
|
||||
n_over_step,
|
||||
m,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let bytes = client.read_one(data_handle);
|
||||
let bytes = client.read_one_unchecked(data_handle);
|
||||
u32::from_bytes(&bytes).to_vec()
|
||||
})
|
||||
}
|
||||
@@ -335,18 +374,18 @@ pub fn gpu_ntt_inverse(device: &GpuDevice, data: &[u32], modulus: u64, root: u64
|
||||
let cube_count = CubeCount::Static((n as u32).div_ceil(256), 1, 1);
|
||||
|
||||
unsafe {
|
||||
let arg = ArrayArg::from_raw_parts::<u32>(&handle, n, 1);
|
||||
let _ = scale_kernel::launch_unchecked::<R>(
|
||||
let arg = ArrayArg::from_raw_parts(handle.clone(), n);
|
||||
scale_kernel::launch_unchecked::<R>(
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
arg,
|
||||
ScalarArg::new(n_inv),
|
||||
ScalarArg::new(m),
|
||||
n_inv,
|
||||
m,
|
||||
);
|
||||
}
|
||||
|
||||
let bytes = client.read_one(handle);
|
||||
let bytes = client.read_one_unchecked(handle);
|
||||
result = u32::from_bytes(&bytes).to_vec();
|
||||
});
|
||||
result
|
||||
@@ -380,21 +419,21 @@ pub fn gpu_poly_multiply(device: &GpuDevice, a: &[u32], b: &[u32], modulus: u64)
|
||||
let cube_count = CubeCount::Static((n as u32).div_ceil(256), 1, 1);
|
||||
|
||||
unsafe {
|
||||
let aa = ArrayArg::from_raw_parts::<u32>(&ha, n, 1);
|
||||
let ba = ArrayArg::from_raw_parts::<u32>(&hb, n, 1);
|
||||
let ca = ArrayArg::from_raw_parts::<u32>(&hc, n, 1);
|
||||
let _ = pointwise_mul_kernel::launch_unchecked::<R>(
|
||||
let aa = ArrayArg::from_raw_parts(ha.clone(), n);
|
||||
let ba = ArrayArg::from_raw_parts(hb.clone(), n);
|
||||
let ca = ArrayArg::from_raw_parts(hc.clone(), n);
|
||||
pointwise_mul_kernel::launch_unchecked::<R>(
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
aa,
|
||||
ba,
|
||||
ca,
|
||||
ScalarArg::new(modulus as u32),
|
||||
modulus as u32,
|
||||
);
|
||||
}
|
||||
|
||||
let bytes = client.read_one(hc);
|
||||
let bytes = client.read_one_unchecked(hc);
|
||||
u32::from_bytes(&bytes).to_vec()
|
||||
});
|
||||
|
||||
|
||||
@@ -275,16 +275,16 @@ impl<'a> GpuPolyEvaluator<'a> {
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
ArrayArg::from_raw_parts::<u32>(&ch, poly.n_terms as usize, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&eh, poly.exponents.len(), 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&ph, flat_points.len(), 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&oh, n_points, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&prm, 3, 1),
|
||||
ArrayArg::from_raw_parts(ch.clone(), poly.n_terms as usize),
|
||||
ArrayArg::from_raw_parts(eh.clone(), poly.exponents.len()),
|
||||
ArrayArg::from_raw_parts(ph.clone(), flat_points.len()),
|
||||
ArrayArg::from_raw_parts(oh.clone(), n_points),
|
||||
ArrayArg::from_raw_parts(prm.clone(), 3),
|
||||
)
|
||||
.expect("batch_poly_eval_kernel launch failed");
|
||||
;
|
||||
}
|
||||
|
||||
let bytes = client.read_one(oh);
|
||||
let bytes = client.read_one_unchecked(oh);
|
||||
u32::from_bytes(&bytes).to_vec()
|
||||
})
|
||||
}
|
||||
@@ -340,18 +340,18 @@ impl<'a> GpuPolyEvaluator<'a> {
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
ArrayArg::from_raw_parts::<u32>(&ca, poly_a.n_terms as usize, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&ea, poly_a.exponents.len(), 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&cb, poly_b.n_terms as usize, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&eb, poly_b.exponents.len(), 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&ph, flat_points.len(), 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&oh, 2 * n_points, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&prm, 5, 1),
|
||||
ArrayArg::from_raw_parts(ca.clone(), poly_a.n_terms as usize),
|
||||
ArrayArg::from_raw_parts(ea.clone(), poly_a.exponents.len()),
|
||||
ArrayArg::from_raw_parts(cb.clone(), poly_b.n_terms as usize),
|
||||
ArrayArg::from_raw_parts(eb.clone(), poly_b.exponents.len()),
|
||||
ArrayArg::from_raw_parts(ph.clone(), flat_points.len()),
|
||||
ArrayArg::from_raw_parts(oh.clone(), 2 * n_points),
|
||||
ArrayArg::from_raw_parts(prm.clone(), 5),
|
||||
)
|
||||
.expect("batch_pair_eval_kernel launch failed");
|
||||
;
|
||||
}
|
||||
|
||||
let bytes = client.read_one(oh);
|
||||
let bytes = client.read_one_unchecked(oh);
|
||||
let interleaved = u32::from_bytes(&bytes);
|
||||
let mut ra = Vec::with_capacity(n_points);
|
||||
let mut rb = Vec::with_capacity(n_points);
|
||||
@@ -389,15 +389,15 @@ impl<'a> GpuPolyEvaluator<'a> {
|
||||
&client,
|
||||
cube_count,
|
||||
cube_dim,
|
||||
ArrayArg::from_raw_parts::<u32>(&ch, coeffs.len(), 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&ph, n_points, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&oh, n_points, 1),
|
||||
ArrayArg::from_raw_parts::<u32>(&prm, 2, 1),
|
||||
ArrayArg::from_raw_parts(ch.clone(), coeffs.len()),
|
||||
ArrayArg::from_raw_parts(ph.clone(), n_points),
|
||||
ArrayArg::from_raw_parts(oh.clone(), n_points),
|
||||
ArrayArg::from_raw_parts(prm.clone(), 2),
|
||||
)
|
||||
.expect("batch_horner_kernel launch failed");
|
||||
;
|
||||
}
|
||||
|
||||
let bytes = client.read_one(oh);
|
||||
let bytes = client.read_one_unchecked(oh);
|
||||
u32::from_bytes(&bytes).to_vec()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -14,6 +14,6 @@ crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
symclaw-core = { path = "../symclaw-core" }
|
||||
pyo3 = { version = "0.24", features = ["extension-module"] }
|
||||
pyo3 = { version = "0.29", features = ["extension-module"] }
|
||||
num-rational = { workspace = true }
|
||||
ordered-float = { workspace = true }
|
||||
|
||||
@@ -163,7 +163,7 @@ fn py_to_expr(obj: &Bound<'_, PyAny>) -> PyResult<Arc<Expr>> {
|
||||
// ── PyExpr ───────────────────────────────────────────────────────
|
||||
|
||||
/// A Python-facing symbolic expression.
|
||||
#[pyclass(name = "Expression")]
|
||||
#[pyclass(name = "Expression", from_py_object)]
|
||||
#[derive(Clone)]
|
||||
struct PyExpr {
|
||||
inner: Arc<Expr>,
|
||||
@@ -482,7 +482,7 @@ impl PyLimitResult {
|
||||
/// Create symbol(s). Single name returns Expression, comma-separated returns tuple.
|
||||
#[pyfunction]
|
||||
#[pyo3(name = "S")]
|
||||
fn py_s(py: Python<'_>, names: &str) -> PyResult<PyObject> {
|
||||
fn py_s(py: Python<'_>, names: &str) -> PyResult<Py<PyAny>> {
|
||||
let parts: Vec<&str> = names
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
|
||||
Reference in New Issue
Block a user