Merge branch 'feat/symclaw-gpu-cubecl-0.10-cuda'

crates/symclaw-gpu/Cargo.toml

@@ -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 }

crates/symclaw-gpu/src/eval.rs

@@ -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 in_arg = || ArrayArg::from_raw_parts(input_handle.clone(), n_points);
-            let out_arg = ArrayArg::from_raw_parts::<f32>(&output_handle, n_points, 1);
+            let out_arg = || ArrayArg::from_raw_parts(output_handle.clone(), n_points);
             match op {
-                Op::Sin => sin_kernel::launch_unchecked::<R>(
+                Op::Sin => {
-                    &client, cube_count, cube_dim, in_arg, out_arg,
+                    sin_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
-                )
+                    true
-                .is_ok(),
+                }
-                Op::Cos => cos_kernel::launch_unchecked::<R>(
+                Op::Cos => {
-                    &client, cube_count, cube_dim, in_arg, out_arg,
+                    cos_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
-                )
+                    true
-                .is_ok(),
+                }
-                Op::Exp => exp_kernel::launch_unchecked::<R>(
+                Op::Exp => {
-                    &client, cube_count, cube_dim, in_arg, out_arg,
+                    exp_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
-                )
+                    true
-                .is_ok(),
+                }
-                Op::Sqrt => sqrt_kernel::launch_unchecked::<R>(
+                Op::Sqrt => {
-                    &client, cube_count, cube_dim, in_arg, out_arg,
+                    sqrt_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
-                )
+                    true
-                .is_ok(),
+                }
-                Op::Neg => neg_kernel::launch_unchecked::<R>(
+                Op::Neg => {
-                    &client, cube_count, cube_dim, in_arg, out_arg,
+                    neg_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
-                )
+                    true
-                .is_ok(),
+                }
-                Op::Abs => abs_kernel::launch_unchecked::<R>(
+                Op::Abs => {
-                    &client, cube_count, cube_dim, in_arg, out_arg,
+                    abs_kernel::launch_unchecked::<R>(&client, cube_count, cube_dim, in_arg(), out_arg());
-                )
+                    true
-                .is_ok(),
+                }
                 _ => 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())
     })

crates/symclaw-gpu/src/groebner.rs

@@ -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(handle.clone(), total),
-                        ArrayArg::from_raw_parts::<u32>(&params_handle, 5, 1),
+                        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]);
             }

crates/symclaw-gpu/src/linalg.rs

@@ -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(a_handle.clone(), a_f32.len()),
-                ArrayArg::from_raw_parts::<f32>(&b_handle, b_f32.len(), 1),
+                ArrayArg::from_raw_parts(b_handle.clone(), b_f32.len()),
-                ArrayArg::from_raw_parts::<f32>(&c_handle, output_len, 1),
+                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())
     })

crates/symclaw-gpu/src/monte_carlo.rs

@@ -211,22 +211,22 @@ pub fn gpu_monte_carlo_integrate_nd(
         let cube_count = CubeCount::Static((n_samples as u32).div_ceil(256), 1, 1);
 
         unsafe {
-            let _ = mc_eval_kernel::launch_unchecked::<R>(
+            mc_eval_kernel::launch_unchecked::<R>(
                 &client,
                 cube_count,
                 cube_dim,
-                ArrayArg::from_raw_parts::<u32>(&ops_h, n_ops, 1),
+                ArrayArg::from_raw_parts(ops_h.clone(), n_ops),
-                ArrayArg::from_raw_parts::<f32>(&consts_h, consts_len, 1),
+                ArrayArg::from_raw_parts(consts_h.clone(), consts_len),
-                ArrayArg::from_raw_parts::<f32>(&lo_h, n_vars, 1),
+                ArrayArg::from_raw_parts(lo_h.clone(), n_vars),
-                ArrayArg::from_raw_parts::<f32>(&hi_h, n_vars, 1),
+                ArrayArg::from_raw_parts(hi_h.clone(), n_vars),
-                ArrayArg::from_raw_parts::<f32>(&out_h, n_samples, 1),
+                ArrayArg::from_raw_parts(out_h.clone(), n_samples),
-                ScalarArg::new(n_ops as u32),
+                n_ops as u32,
-                ScalarArg::new(n_vars as u32),
+                n_vars as u32,
-                ScalarArg::new(BASE_SEED),
+                BASE_SEED,
             );
         }
 
-        let bytes = client.read_one(out_h);
+        let bytes = client.read_one_unchecked(out_h);
         f32::from_bytes(&bytes).to_vec()
     });
 

crates/symclaw-gpu/src/ntt.rs

@@ -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.
+/// GPU modular multiply: (a * b) % m using 16-bit splits, with the `<<16`
-#[cube]
+/// reductions fully INLINED (no nested `#[cube]` helper calls — the cuda/cpp
-fn shift_left_16(val: u32, m: u32) -> u32 {
+/// codegen mishandled the nested form, producing unreduced NTT results; wgpu
-    let mut s = val;
+/// was unaffected). Both a, b must be < m < 2^30 so every partial product fits
-    s = (s + s) % m;
+/// in u32.
-    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.
 #[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,27 +271,31 @@ 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.
+        // Modular multiply (v_raw * tw) % modulus.
-    // Split both into high/low 16-bit halves.
         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.
 pub fn compute_twiddles(n: usize, root: u64, modulus: u64) -> Vec<u32> {
@@ -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 data_arg = ArrayArg::from_raw_parts(data_handle.clone(), n);
-                let tw_arg = ArrayArg::from_raw_parts::<u32>(&tw_handle, n, 1);
+                let tw_arg = ArrayArg::from_raw_parts(tw_handle.clone(), n);
-                let _ = ntt_butterfly_kernel::launch_unchecked::<R>(
+                ntt_butterfly_kernel::launch_unchecked::<R>(
                     &client,
                     cube_count,
                     cube_dim,
                     data_arg,
                     tw_arg,
-                    ScalarArg::new(half),
+                    half,
-                    ScalarArg::new(step),
+                    step,
-                    ScalarArg::new(n_over_step),
+                    n_over_step,
-                    ScalarArg::new(m),
+                    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 arg = ArrayArg::from_raw_parts(handle.clone(), n);
-            let _ = scale_kernel::launch_unchecked::<R>(
+            scale_kernel::launch_unchecked::<R>(
                 &client,
                 cube_count,
                 cube_dim,
                 arg,
-                ScalarArg::new(n_inv),
+                n_inv,
-                ScalarArg::new(m),
+                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 aa = ArrayArg::from_raw_parts(ha.clone(), n);
-            let ba = ArrayArg::from_raw_parts::<u32>(&hb, n, 1);
+            let ba = ArrayArg::from_raw_parts(hb.clone(), n);
-            let ca = ArrayArg::from_raw_parts::<u32>(&hc, n, 1);
+            let ca = ArrayArg::from_raw_parts(hc.clone(), n);
-            let _ = pointwise_mul_kernel::launch_unchecked::<R>(
+            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()
     });
 

crates/symclaw-gpu/src/poly_gcd.rs

@@ -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(ch.clone(), poly.n_terms as usize),
-                    ArrayArg::from_raw_parts::<u32>(&eh, poly.exponents.len(), 1),
+                    ArrayArg::from_raw_parts(eh.clone(), poly.exponents.len()),
-                    ArrayArg::from_raw_parts::<u32>(&ph, flat_points.len(), 1),
+                    ArrayArg::from_raw_parts(ph.clone(), flat_points.len()),
-                    ArrayArg::from_raw_parts::<u32>(&oh, n_points, 1),
+                    ArrayArg::from_raw_parts(oh.clone(), n_points),
-                    ArrayArg::from_raw_parts::<u32>(&prm, 3, 1),
+                    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(ca.clone(), poly_a.n_terms as usize),
-                    ArrayArg::from_raw_parts::<u32>(&ea, poly_a.exponents.len(), 1),
+                    ArrayArg::from_raw_parts(ea.clone(), poly_a.exponents.len()),
-                    ArrayArg::from_raw_parts::<u32>(&cb, poly_b.n_terms as usize, 1),
+                    ArrayArg::from_raw_parts(cb.clone(), poly_b.n_terms as usize),
-                    ArrayArg::from_raw_parts::<u32>(&eb, poly_b.exponents.len(), 1),
+                    ArrayArg::from_raw_parts(eb.clone(), poly_b.exponents.len()),
-                    ArrayArg::from_raw_parts::<u32>(&ph, flat_points.len(), 1),
+                    ArrayArg::from_raw_parts(ph.clone(), flat_points.len()),
-                    ArrayArg::from_raw_parts::<u32>(&oh, 2 * n_points, 1),
+                    ArrayArg::from_raw_parts(oh.clone(), 2 * n_points),
-                    ArrayArg::from_raw_parts::<u32>(&prm, 5, 1),
+                    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(ch.clone(), coeffs.len()),
-                    ArrayArg::from_raw_parts::<u32>(&ph, n_points, 1),
+                    ArrayArg::from_raw_parts(ph.clone(), n_points),
-                    ArrayArg::from_raw_parts::<u32>(&oh, n_points, 1),
+                    ArrayArg::from_raw_parts(oh.clone(), n_points),
-                    ArrayArg::from_raw_parts::<u32>(&prm, 2, 1),
+                    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()
         })
     }