chore(deps): pin dev's Jolt crates to Jolt main (f5c1947b1)

.github/workflows/ci.yml

@@ -147,6 +147,14 @@ jobs:
       - name: Require one jolt-field identity and no akita-field
         run: scripts/check-shared-field-identity.sh
 
+  jolt-pins:
+    name: Jolt pins on main
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Require every a16z/jolt rev to be an ancestor of Jolt main
+        run: scripts/check-jolt-pins-on-main.sh
+
   doc:
     name: Documentation
     runs-on: ubuntu-latest

AGENTS.md

@@ -36,6 +36,7 @@ scripts/check-crate-deps.sh akita-zk-prover
 scripts/check-crate-deps.sh akita-labinius-verifier
 scripts/check-crate-deps.sh akita-labinius-prover
 scripts/check-shared-field-identity.sh
+scripts/check-jolt-pins-on-main.sh
 scripts/check-external-schedule-artifacts.sh
 cargo machete --with-metadata
 typos

book/src/usage/jolt-recursion.md

@@ -140,7 +140,7 @@ ZEROOS_GUEST_RUSTFLAGS=-Zunstable-options \
 
 The full command performs Jolt preprocessing, proves the guest execution, and
 verifies the resulting Jolt proof. Success ends with
-`Akita-in-Jolt proof OK`, `is_valid=true`, and `guest_panic=false`.
+`Akita-in-Jolt proof OK`, `accepted=true`, and `guest_panic=false`.
 
 The 32 variable target is intended for a server with enough memory for its
 large input and trace. Measure it with `--trace-only` first. The current

crates/akita-algebra/Cargo.toml

@@ -10,11 +10,11 @@ readme = "../../README.md"
 
 [dependencies]
 akita-error = { version = "0.1.0", path = "../akita-error" }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "asm",
   "solinas",
 ] }
-jolt-poly = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false }
+jolt-poly = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false }
 akita-serialization = { version = "0.1.0", path = "../akita-serialization" }
 rand_core = { version = "0.6", features = ["getrandom"] }
 rayon = { version = "1.10", optional = true }

crates/akita-challenges/Cargo.toml

@@ -11,7 +11,7 @@ readme = "../../README.md"
 [dependencies]
 akita-error = { version = "0.1.0", path = "../akita-error" }
 serde = { version = "1", features = ["derive"] }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "asm",
   "solinas",
 ] }

crates/akita-config/Cargo.toml

@@ -31,7 +31,7 @@ dev-protocol = ["akita-params/dev-protocol"]
 [dependencies]
 akita-error = { version = "0.1.0", path = "../akita-error" }
 akita-challenges = { version = "0.1.0", path = "../akita-challenges", default-features = false }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "asm",
   "solinas",
 ] }

crates/akita-cpu-backend/Cargo.toml

@@ -14,11 +14,11 @@ akita-error = { version = "0.1.0", path = "../akita-error" }
 akita-algebra = { version = "0.1.0", path = "../akita-algebra" }
 akita-challenges = { version = "0.1.0", path = "../akita-challenges", default-features = false }
 akita-config = { version = "0.1.0", path = "../akita-config", default-features = false }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "asm",
   "solinas",
 ] }
-jolt-poly = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false }
+jolt-poly = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false }
 akita-serialization = { version = "0.1.0", path = "../akita-serialization" }
 akita-sumcheck = { version = "0.1.0", path = "../akita-sumcheck", default-features = false }
 akita-transcript = { version = "0.1.0", path = "../akita-transcript", default-features = false }

crates/akita-labinius-prover/Cargo.toml

@@ -17,12 +17,12 @@ tracing = "0.1"
 akita-challenges = { version = "0.1.0", path = "../akita-challenges", default-features = false }
 akita-params = { version = "0.1.0", path = "../akita-params", default-features = false }
 akita-labinius-verifier = { version = "0.1.0", path = "../akita-labinius-verifier", default-features = false }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "asm",
   "solinas",
 ] }
 
-jolt-poly = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false }
+jolt-poly = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false }
 
 [dev-dependencies]
 akita-transcript = { version = "0.1.0", path = "../akita-transcript", default-features = false }

crates/akita-labinius-verifier/Cargo.toml

@@ -17,7 +17,7 @@ akita-transcript = { version = "0.1.0", path = "../akita-transcript", default-fe
 akita-challenges = { version = "0.1.0", path = "../akita-challenges", default-features = false }
 akita-params = { version = "0.1.0", path = "../akita-params", default-features = false }
 akita-types = { version = "0.1.0", path = "../akita-types", default-features = false, optional = true }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "asm",
   "solinas",
 ] }

crates/akita-metal/Cargo.toml

@@ -18,11 +18,11 @@ akita-cpu-backend = { version = "0.1.0", path = "../akita-cpu-backend", default-
 akita-error = { version = "0.1.0", path = "../akita-error" }
 akita-types = { version = "0.1.0", path = "../akita-types", default-features = false }
 bytemuck = { version = "1.16", features = ["derive"] }
-jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c", default-features = false, features = [
+jolt-field = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3", default-features = false, features = [
   "bytemuck",
   "solinas",
 ] }
-jolt-metal = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "c7c2d08d5bb0c7c80501f29f723e595af7fcff3c" }
+jolt-metal = { version = "=0.1.0", git = "https://github.com/a16z/jolt", rev = "f5c1947b15305e469d1602e8e92ff07c33d778e3" }
 thiserror = "2.0"
 
 [dev-dependencies]

crates/akita-metal/src/crt.rs

@@ -97,19 +97,27 @@ impl<const K: usize, const D: usize> DeviceCrtNtt<K, D> {
         let radix = DeviceBuffer::from_slice(metal.device(), &self.crt_weights::<F>())?;
         let scales = DeviceBuffer::from_slice(metal.device(), &[F::one()])?;
         let mut batch = Batch::new(metal.device())?;
-        batch.dispatch(
-            pipeline,
-            &[
-                Binding::buffer(residues),
-                Binding::buffer(out),
-                Binding::buffer(&self.primes),
-                Binding::buffer(&self.gamma),
-                Binding::buffer(&radix),
-                Binding::buffer(&scales),
-                Binding::value(&shape),
-            ],
-            Grid::linear(coefficients, COEFFICIENT_GROUP),
-        )?;
+        // SAFETY: `rows` makes `residues` exactly `elements` rings of K x D
+        // words, and `out` holds `elements * D` coefficients. Thread
+        // `index < coefficients` reads the residue words
+        // `element * K * D + i * D + coefficient` for `i < K`, `primes[i]`,
+        // `gamma[i * K + j]` and `radix[i]` (buffers of K, K x K and K entries for
+        // this K), and `scales[0]` (one limb). It writes only `out[index]`.
+        unsafe {
+            batch.dispatch_unchecked(
+                pipeline,
+                &[
+                    Binding::buffer(residues),
+                    Binding::buffer(out),
+                    Binding::buffer(&self.primes),
+                    Binding::buffer(&self.gamma),
+                    Binding::buffer(&radix),
+                    Binding::buffer(&scales),
+                    Binding::value(&shape),
+                ],
+                Grid::linear(coefficients, COEFFICIENT_GROUP),
+            )
+        }?;
         Ok(batch.commit_and_wait()?)
     }
 }

crates/akita-metal/src/decompose.rs

@@ -117,15 +117,22 @@ where
     };
     let mut out = DeviceDigitPlanes::zeroed(metal, digits, log_basis)?;
     let mut batch = Batch::new(metal.device())?;
-    batch.dispatch(
-        pipeline,
-        &[
-            Binding::buffer(coefficients),
-            Binding::buffer(out.buffer_mut()),
-            Binding::value(&shape),
-        ],
-        Grid::linear(count, COEFFICIENT_GROUP),
-    )?;
+    // SAFETY: thread `index < count` reads `coefficients[index]` and writes
+    // the plane words `(ring * levels + level) * ring_degree + coefficient` for
+    // `level < levels`, all below `count * levels = digits`, the length of
+    // `out`. `ring_degree` is a power of two below 2^32, so the kernel's shift
+    // and mask compute exactly that index. No two threads write the same word.
+    unsafe {
+        batch.dispatch_unchecked(
+            pipeline,
+            &[
+                Binding::buffer(coefficients),
+                Binding::buffer(out.buffer_mut()),
+                Binding::value(&shape),
+            ],
+            Grid::linear(count, COEFFICIENT_GROUP),
+        )
+    }?;
     let duration = batch.commit_and_wait()?;
     Ok((out, duration))
 }

crates/akita-metal/src/matvec.rs

@@ -479,42 +479,74 @@ where
         // and residue buffers.
         let mut batch = Batch::new(metal.device())?;
         for (shape, crt_shape) in shapes.iter().zip(&crt_shapes) {
-            batch.dispatch(
-                partials_pipeline,
-                &[
-                    Binding::buffer(planes.buffer()),
-                    Binding::buffer(&self.entries),
-                    Binding::buffer(&partials),
-                    Binding::buffer(&self.ntt.primes),
-                    Binding::buffer(&self.ntt.tables),
-                    Binding::value(shape),
-                ],
-                Grid::linear(partial_threads, lanes),
-            )?;
-            batch.dispatch(
-                finish_pipeline,
-                &[
-                    Binding::buffer(&partials),
-                    Binding::buffer(&residues),
-                    Binding::buffer(&self.ntt.primes),
-                    Binding::buffer(&self.ntt.tables),
-                    Binding::value(shape),
-                ],
-                Grid::linear(finish_threads, lanes),
-            )?;
-            batch.dispatch(
-                crt_pipeline,
-                &[
-                    Binding::buffer(&residues),
-                    Binding::buffer(out),
-                    Binding::buffer(&self.ntt.primes),
-                    Binding::buffer(&self.ntt.gamma),
-                    Binding::buffer(&radix),
-                    Binding::buffer(&scales),
-                    Binding::value(crt_shape),
-                ],
-                Grid::linear(coefficients, COEFFICIENT_GROUP),
-            )?;
+            // SAFETY: the grid is `chunks * K * block_tiles * row_tiles` full
+            // threadgroups of D/2 lanes. A threadgroup reads the digit planes
+            // `(block * cols + col) * D + lane` and `+ D/2` with `block < blocks` and
+            // `col < col_end <= cols`, below `planes.len() = blocks * cols * D`; the
+            // matrix entry `(row, col, limb, prime)` with `row < rows`, `limb < limbs`
+            // and `prime < K`, below `entries.len() = rows * cols * limbs * K * D`; and
+            // words below D of each table in prime `prime`'s block. It writes only its
+            // own `(chunk, limb, block, row, prime)` row of `partials`, which holds
+            // `chunks * limbs * blocks * rows * K * D` words. In the threadgroup array
+            // `x` of `BLOCK_TILE * D` words, each lane writes and reads back only its own
+            // two words per array, and `PlainForward::apply` orders the cross-lane
+            // accesses with its barriers.
+            unsafe {
+                batch.dispatch_unchecked(
+                    partials_pipeline,
+                    &[
+                        Binding::buffer(planes.buffer()),
+                        Binding::buffer(&self.entries),
+                        Binding::buffer(&partials),
+                        Binding::buffer(&self.ntt.primes),
+                        Binding::buffer(&self.ntt.tables),
+                        Binding::value(shape),
+                    ],
+                    Grid::linear(partial_threads, lanes),
+                )
+            }?;
+            // SAFETY: the grid is `residue_rows` full threadgroups of D/2 lanes, one
+            // per (limb, block, row, prime) row. Row `group` reads `partials` words
+            // `chunk * stride + group * D + lane` and `+ D/2` for `chunk < chunks`,
+            // with `stride = residue_rows * D`, below `partials.len()`; and words below D
+            // of prime `group % K`'s tables. It writes only its own row of `residues`,
+            // which holds `residue_rows * D` words. Lanes share only the threadgroup
+            // array, whose accesses the barriers in `ntt.h` order.
+            unsafe {
+                batch.dispatch_unchecked(
+                    finish_pipeline,
+                    &[
+                        Binding::buffer(&partials),
+                        Binding::buffer(&residues),
+                        Binding::buffer(&self.ntt.primes),
+                        Binding::buffer(&self.ntt.tables),
+                        Binding::value(shape),
+                    ],
+                    Grid::linear(finish_threads, lanes),
+                )
+            }?;
+            // SAFETY: `out` holds `coefficients = blocks * rows * D` elements and
+            // `residues` holds `limbs` groups of `blocks * rows` rings of K x D words.
+            // Thread `index < coefficients` reads the residue words
+            // `limb * coefficients * K + element * K * D + i * D + coefficient` for
+            // `limb < limbs` and `i < K`, below `residues.len()`; `primes[i]`,
+            // `gamma[i * K + j]`, `radix[i]` and `scales[limb]` (`limb_scales` returns
+            // `limbs` entries). It reads and writes only `out[index]`.
+            unsafe {
+                batch.dispatch_unchecked(
+                    crt_pipeline,
+                    &[
+                        Binding::buffer(&residues),
+                        Binding::buffer(out),
+                        Binding::buffer(&self.ntt.primes),
+                        Binding::buffer(&self.ntt.gamma),
+                        Binding::buffer(&radix),
+                        Binding::buffer(&scales),
+                        Binding::value(crt_shape),
+                    ],
+                    Grid::linear(coefficients, COEFFICIENT_GROUP),
+                )
+            }?;
         }
         Ok(batch.commit_and_wait()?)
     }

crates/akita-metal/src/ntt.rs

@@ -177,16 +177,25 @@ impl<const K: usize, const D: usize> DeviceCrtNtt<K, D> {
             .ok_or_else(|| shape_overflow(rows))?;
         let pipeline = metal.pipeline(&ring_degree_kernel(template, D))?;
         let mut batch = Batch::new(metal.device())?;
-        batch.dispatch(
-            pipeline,
-            &[
-                Binding::buffer(data),
-                Binding::buffer(&self.primes),
-                Binding::buffer(&self.tables),
-                Binding::value(&batch_shape),
-            ],
-            Grid::linear(threads, lanes),
-        )?;
+        // SAFETY: `rows` makes `data` exactly `rows` (element, prime) rows of D
+        // words, and the grid is `rows` full threadgroups of D/2 lanes. Lane `lane`
+        // of row `row` reads and writes `data[row * D + lane]` and
+        // `data[row * D + lane + D/2]`, and reads `primes[row % K]` and words below
+        // D of each table in prime `row % K`'s block of the `[K][7][D]` tables.
+        // Lanes share only the threadgroup array, whose accesses the barriers in
+        // `ntt.h` order.
+        unsafe {
+            batch.dispatch_unchecked(
+                pipeline,
+                &[
+                    Binding::buffer(data),
+                    Binding::buffer(&self.primes),
+                    Binding::buffer(&self.tables),
+                    Binding::value(&batch_shape),
+                ],
+                Grid::linear(threads, lanes),
+            )
+        }?;
         Ok(batch.commit_and_wait()?)
     }
 }