CUDA acceleration for trace and constraint commitments

air/src/options.rs

@@ -129,6 +129,21 @@ impl ProvingOptions {
         self
     }
 
+    /// Sets partitions for this [ProvingOptions].
+    ///
+    /// Partitions can be provided to split traces during proving and distribute work across
+    /// multiple devices. The number of partitions should be equal to the number of devices.
+    pub const fn with_partitions(mut self, num_partitions: usize) -> Self {
+        let hash_rate = match self.hash_fn {
+            HashFunction::Blake3_192 => 6,
+            HashFunction::Blake3_256 => 8,
+            HashFunction::Rpo256 => 8,
+            HashFunction::Rpx256 => 8,
+        };
+        self.proof_options = self.proof_options.with_partitions(num_partitions, hash_rate);
+        self
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 

miden/Cargo.toml

@@ -55,6 +55,7 @@ executable = [
     "dep:rustyline",
     "dep:tracing-subscriber",
 ]
+cuda = ["prover/cuda", "std"]
 metal = ["prover/metal", "std"]
 std = ["assembly/std", "processor/std", "prover/std", "verifier/std"]
 # For internal use, not meant to be used by users

miden/src/cli/prove.rs

@@ -4,6 +4,8 @@ use assembly::diagnostics::{IntoDiagnostic, Report, WrapErr};
 use clap::Parser;
 use miden_vm::{internal::InputFile, ProvingOptions};
 use processor::{DefaultHost, ExecutionOptions, ExecutionOptionsError, Program};
+#[cfg(all(target_arch = "x86_64", feature = "cuda"))]
+use prover::get_num_of_gpus;
 use stdlib::StdLibrary;
 use tracing::instrument;
 
@@ -65,6 +67,11 @@ impl ProveCmd {
     pub fn get_proof_options(&self) -> Result<ProvingOptions, ExecutionOptionsError> {
         let exec_options =
             ExecutionOptions::new(Some(self.max_cycles), self.expected_cycles, self.trace, false)?;
+
+        let partitions = 1;
+        #[cfg(all(target_arch = "x86_64", feature = "cuda"))]
+        let partitions = get_num_of_gpus();
+
         Ok(match self.security.as_str() {
             "96bits" => {
                 if self.rpx {
@@ -82,7 +89,8 @@ impl ProveCmd {
             },
             other => panic!("{} is not a valid security setting", other),
         }
-        .with_execution_options(exec_options))
+        .with_execution_options(exec_options)
+        .with_partitions(partitions))
     }
 
     pub fn execute(&self) -> Result<(), Report> {

prover/Cargo.toml

@@ -17,7 +17,8 @@ edition.workspace = true
 async = ["winter-maybe-async/async"]
 concurrent = ["processor/concurrent", "std", "winter-prover/concurrent"]
 default = ["std"]
-metal = ["dep:miden-gpu", "dep:elsa", "dep:pollster", "concurrent", "std"]
+metal = ["dep:miden-gpu", "miden-gpu/metal", "dep:elsa", "dep:pollster", "concurrent", "std"]
+cuda = ["dep:miden-gpu", "miden-gpu/cuda", "std"]
 std = ["air/std", "processor/std", "winter-prover/std"]
 
 [dependencies]
@@ -29,5 +30,11 @@ winter-prover = { package = "winter-prover", version = "0.11", default-features
 
 [target.'cfg(all(target_arch = "aarch64", target_os = "macos"))'.dependencies]
 elsa = { version = "1.9", optional = true }
-miden-gpu = { version = "0.4", optional = true }
+miden-gpu = { path = "../../miden-gpu", optional = true }
 pollster = { version = "0.4", optional = true }
+
+[target.'cfg(target_arch = "x86_64")'.dependencies]
+miden-gpu = { path = "../../miden-gpu", optional = true }
+
+[dev-dependencies]
+serial_test = "3.1"

prover/README.md

@@ -45,6 +45,7 @@ Miden prover can be compiled with the following features:
 * `std` - enabled by default and relies on the Rust standard library.
 * `concurrent` - implies `std` and also enables multi-threaded proof generation.
 * `metal` - enables [Metal](https://en.wikipedia.org/wiki/Metal_(API))-based acceleration of proof generation (for recursive proofs) on supported platforms (e.g., Apple silicon).
+* `cuda` - enables [Cuda](https://en.wikipedia.org/wiki/CUDA)-based acceleration of proof generation (for recursive proofs) on supported platforms.
 * `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly.
     * Only the `wasm32-unknown-unknown` and `wasm32-wasip1` targets are officially supported.