Revise error norm simulation and GPU test coverage

.gitignore

@@ -12,4 +12,5 @@
 /revision_logs
 /*.log
 /*.status
+/tmp_*.sh
 /bench

Cargo.toml

@@ -31,7 +31,7 @@ digest = "0.10"
 bitvec = "1"
 keccak-asm = { version = "0.1.4" }
 tempfile = { version = "3.19.1" }
-tracing-subscriber = "0.3"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
 [dev-dependencies]
 sequential-test = "0.2.4"

build.rs

@@ -62,7 +62,7 @@ fn main() {
             .flag("-std=c++17")
             .flag("-Xcompiler")
             .flag("-fPIC")
-            .flag(&format!("-arch=sm_{cuda_arch}"));
+            .flag(format!("-arch=sm_{cuda_arch}"));
         if !debug_build {
             build.flag("-lineinfo");
         }

cuda/src/matrix/MatrixDecompose.cu

@@ -29,6 +29,55 @@ namespace
     constexpr size_t kDecomposeMaxGridZ = 65535;
 }
 
+__device__ __forceinline__ int64_t centered_lift_u64(uint64_t residue, uint64_t modulus)
+{
+    __int128 value = static_cast<__int128>(residue);
+    if (static_cast<unsigned __int128>(residue) * 2U > static_cast<unsigned __int128>(modulus))
+    {
+        value -= static_cast<__int128>(modulus);
+    }
+    return static_cast<int64_t>(value);
+}
+
+__device__ __forceinline__ int64_t balanced_digit_step(int64_t value, int64_t base, int64_t *next)
+{
+    int64_t quotient = value / base;
+    int64_t remainder = value % base;
+    if (remainder < 0)
+    {
+        remainder += base;
+        quotient -= 1;
+    }
+    const int64_t half = base / 2;
+    if (remainder < half)
+    {
+        *next = quotient;
+        return remainder;
+    }
+    if (remainder > half)
+    {
+        *next = quotient + 1;
+        return remainder - base;
+    }
+    if ((quotient & 1LL) == 0)
+    {
+        *next = quotient;
+        return half;
+    }
+    *next = quotient + 1;
+    return half - base;
+}
+
+__device__ __forceinline__ uint64_t signed_digit_to_residue(int64_t digit, uint64_t modulus)
+{
+    if (digit >= 0)
+    {
+        return static_cast<uint64_t>(digit) % modulus;
+    }
+    const uint64_t magnitude = static_cast<uint64_t>(-digit) % modulus;
+    return magnitude == 0 ? 0 : modulus - magnitude;
+}
+
 __global__ void matrix_decompose_all_slots_kernel(
     const uint8_t *src_base,
     uint8_t *const *dst_bases,
@@ -45,8 +94,10 @@ __global__ void matrix_decompose_all_slots_kernel(
     size_t out_cols,
     size_t log_base_q,
     uint32_t src_bits,
+    uint64_t src_modulus,
     uint32_t base_bits,
     uint32_t digits_per_tower,
+    bool balanced,
     size_t src_digit_offset_base,
     size_t poly_offset,
     size_t slot_offset)
@@ -92,18 +143,39 @@ __global__ void matrix_decompose_all_slots_kernel(
 
     const uint64_t residue =
         matrix_load_limb_u64(src_base, poly_idx, coeff_idx, src_stride_bytes, src_coeff_bytes);
-    const uint32_t shift = digit_idx * base_bits;
-    uint64_t mask = 0;
-    if (shift < src_bits)
+    uint64_t digit = 0;
+    if (balanced)
     {
-        const uint32_t remaining = src_bits - shift;
-        const uint32_t digit_bits = min(base_bits, remaining);
-        mask = digit_bits >= 64 ? ~uint64_t{0} : ((uint64_t{1} << digit_bits) - 1);
+        int64_t value = centered_lift_u64(residue, src_modulus);
+        int64_t signed_digit = 0;
+        const int64_t base = int64_t{1} << base_bits;
+        for (uint32_t idx = 0; idx <= digit_idx; ++idx)
+        {
+            int64_t next = 0;
+            const int64_t current_digit = balanced_digit_step(value, base, &next);
+            if (idx == digit_idx)
+            {
+                signed_digit = current_digit;
+            }
+            value = next;
+        }
+        digit = signed_digit_to_residue(signed_digit, out_modulus);
     }
-    uint64_t digit = shift >= 64 ? 0 : ((residue >> shift) & mask);
-    if (out_modulus != 0 && digit >= out_modulus)
+    else
     {
-        digit %= out_modulus;
+        const uint32_t shift = digit_idx * base_bits;
+        uint64_t mask = 0;
+        if (shift < src_bits)
+        {
+            const uint32_t remaining = src_bits - shift;
+            const uint32_t digit_bits = min(base_bits, remaining);
+            mask = digit_bits >= 64 ? ~uint64_t{0} : ((uint64_t{1} << digit_bits) - 1);
+        }
+        digit = shift >= 64 ? 0 : ((residue >> shift) & mask);
+        if (out_modulus != 0 && digit >= out_modulus)
+        {
+            digit %= out_modulus;
+        }
     }
 
     const size_t row = poly_idx / src_cols;
@@ -128,8 +200,10 @@ int launch_decompose_all_slots_kernel(
     size_t out_cols,
     size_t log_base_q,
     uint32_t src_bits,
+    uint64_t src_modulus,
     uint32_t base_bits,
     uint32_t digits_per_tower,
+    bool balanced,
     size_t src_digit_offset_base,
     cudaStream_t stream)
 {
@@ -194,8 +268,10 @@ int launch_decompose_all_slots_kernel(
                 out_cols,
                 log_base_q,
                 src_bits,
+                src_modulus,
                 base_bits,
                 digits_per_tower,
+                balanced,
                 src_digit_offset_base,
                 poly_offset,
                 slot_offset);
@@ -1284,8 +1360,10 @@ static int gpu_matrix_decompose_base_impl(
             out->cols,
             out_log_base_q,
             src_bits,
+            src->ctx->moduli[src_idx],
             base_bits,
             digits_per_tower,
+            !small,
             src_digit_offset_base,
             dispatch_stream);
         if (status != 0)

src/bench_estimator/mod.rs

@@ -913,7 +913,7 @@ mod tests {
             &self,
             _src_slots: &[(u32, Option<u32>)],
         ) -> CircuitBenchEstimate {
-            bench(6.0, 6.5, 17)
+            bench(6.0, 6.6, 17)
         }
 
         fn estimate_slot_reduce(

src/circuit/poly_circuit/eval.rs

@@ -13,6 +13,33 @@ impl<P: Poly> PolyCircuit<P> {
         slot_transfer_evaluator: Option<&dyn SlotTransferEvaluator<E>>,
         parallel_gates: Option<usize>,
     ) -> Vec<E>
+    where
+        E: Evaluable<P = P>,
+        E::Compact: Send + Sync,
+        PE: PltEvaluator<E>,
+        PE: Sync,
+    {
+        let one_compact = one.to_compact();
+        let input_compacts = inputs.into_iter().map(E::to_compact).collect();
+        self.eval_from_compacts(
+            params,
+            one_compact,
+            input_compacts,
+            plt_evaluator,
+            slot_transfer_evaluator,
+            parallel_gates,
+        )
+    }
+
+    pub fn eval_from_compacts<E, PE>(
+        &self,
+        params: &E::Params,
+        one_compact: E::Compact,
+        input_compacts: Vec<E::Compact>,
+        plt_evaluator: Option<&PE>,
+        slot_transfer_evaluator: Option<&dyn SlotTransferEvaluator<E>>,
+        parallel_gates: Option<usize>,
+    ) -> Vec<E>
     where
         E: Evaluable<P = P>,
         E::Compact: Send + Sync,
@@ -21,10 +48,9 @@ impl<P: Poly> PolyCircuit<P> {
     {
         let (call_id_base, gate_id_base) = self.eval_gate_id_bases();
         let parallel_gates = crate::env::resolve_circuit_parallel_gates(parallel_gates);
-        let one_compact = Arc::new(one.to_compact());
+        let one_compact = Arc::new(one_compact);
         let one = Arc::new(E::from_compact(params, one_compact.as_ref()));
-        let input_compacts =
-            inputs.into_iter().map(|input| Arc::new(input.to_compact())).collect::<Vec<_>>();
+        let input_compacts = input_compacts.into_iter().map(Arc::new).collect::<Vec<_>>();
         let scoped_gate_ids = self.build_scoped_gate_id_map(&call_id_base, &gate_id_base);
         let call_prefix = ScopedGateKey::from(0u32);
         let outputs = std::thread::scope(|scope| {
@@ -902,32 +928,11 @@ impl<P: Poly> PolyCircuit<P> {
                             .iter()
                             .copied()
                             .for_each(|gate_id| eval_gate(gate_id, eval_params, eval_one));
-                    } else if !regular_chunks.is_empty() {
-                        let mut loaded_curr = Some(load_chunk(regular_chunks[0]));
-                        let mut computed_prev: Option<Vec<ComputedGateCtx<E>>> = None;
-                        for chunk_idx in 0..regular_chunks.len() {
-                            let to_store = computed_prev.take();
-                            let to_compute =
-                                loaded_curr.take().expect("loaded chunk missing in pipeline");
-                            let next_chunk = regular_chunks.get(chunk_idx + 1).copied();
-                            let ((), (computed_curr, loaded_next)) = rayon::join(
-                                || {
-                                    if let Some(computed_chunk) = to_store {
-                                        store_chunk(computed_chunk);
-                                    }
-                                },
-                                || {
-                                    rayon::join(
-                                        || compute_chunk(to_compute),
-                                        || next_chunk.map(load_chunk),
-                                    )
-                                },
-                            );
-                            computed_prev = Some(computed_curr);
-                            loaded_curr = loaded_next;
-                        }
-                        if let Some(last_chunk) = computed_prev.take() {
-                            store_chunk(last_chunk);
+                    } else {
+                        for chunk in regular_chunks {
+                            let loaded = load_chunk(chunk);
+                            let computed = compute_chunk(loaded);
+                            store_chunk(computed);
                         }
                     }
                 }