feat: Adding GPU test

.github/workflows/test-gpu.yaml

@@ -0,0 +1,129 @@
+name: Test on GPU
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  zksync-airbender-build:
+    runs-on: [ ubuntu-latest ]
+    strategy:
+      matrix:
+        cuda: [ "12.5.0-devel-ubuntu20.04" ]
+        package: [ "zksmith" ]
+    container:
+      image: nvidia/cuda:${{ matrix.cuda }}
+    env:
+      BELLMAN_CUDA_DIR: ${{ github.workspace }}/bellman-cuda
+    steps:
+      - name: Prepare environment
+        env:
+          DEBIAN_FRONTEND: noninteractive
+        run: |
+          apt update && apt install -y \
+            pkg-config libclang-dev build-essential lldb lld \
+            clang openssl libssl-dev gcc g++ wget curl jq
+          echo "/usr/local/nvidia/bin:/usr/local/cuda/bin" >> $GITHUB_PATH
+
+      - uses: actions/checkout@93ea575cb5d8a053eaa0ac8fa3b40d7e05a33cc8 # v3
+
+      - name: Setup rustup
+        run: |
+          wget -q -O - https://sh.rustup.rs | bash -s -- -y
+          echo "${HOME}/.cargo/bin" >> "${GITHUB_PATH}"
+          echo "export PATH=\"$HOME/.cargo/bin:\$PATH\"" >> "${HOME}/.bash_profile"
+
+      - name: Setup CMake
+        run: |
+          curl -LO https://github.com/Kitware/CMake/releases/download/v3.24.3/cmake-3.24.3-linux-x86_64.sh && \
+            chmod +x cmake-3.24.3-linux-x86_64.sh && \
+            ./cmake-3.24.3-linux-x86_64.sh --skip-license --prefix=/usr/local
+
+      - name: Prepare bellman-cuda directory
+        shell: bash
+        # Curl ugliness is required because gh can't work with authentication: https://github.com/cli/cli/issues/2680.
+        run: |
+          release=($(curl --silent https://api.github.com/repos/matter-labs/era-bellman-cuda/releases | jq -r '.[0] | .name, .tarball_url, .assets[0].browser_download_url'))
+          curl --silent -L "${release[1]}" --output bellman-cuda-source.tar.gz
+          curl --silent -L "${release[2]}" --output bellman-cuda.tar.gz
+          mkdir -p bellman-cuda
+          tar xvfz bellman-cuda.tar.gz -C ./bellman-cuda
+          tar xvfz bellman-cuda-source.tar.gz -C ./bellman-cuda --strip-components=1 --wildcards \*/src/
+
+      - name: Check CMake version
+        run: |
+          cmake --version
+
+      - name: Check CUDA version
+        run: |
+          nvcc --version
+
+      - name: Setup rust
+        run: |
+          rustup set profile minimal
+          rustup toolchain install nightly-2025-05-24
+          rustup default nightly-2025-05-24
+
+      - name: Build tests and copy binary to a separate dir
+        shell: bash
+        run: |
+          mkdir artifacts
+          CUDAARCHS=80 CARGO_TARGET_DIR=./build \
+          cargo +nightly-2025-05-24 test -p ${{ matrix.package }} --no-run --release --message-format=json -q \
+          | jq -r 'select(.executable != null) | .executable' \
+          | while read binary; do
+            cp "$binary" artifacts/${{ matrix.package }}
+          done
+
+      - name: Upload test binary
+        uses: actions/upload-artifact@v4
+        with:
+          name: zksync-airbender-${{ matrix.cuda }}-${{ matrix.package }}-test-binary
+          path: artifacts/${{ matrix.package }}
+          if-no-files-found: error
+
+  zksync-airbender-test:
+    runs-on: [ matterlabs-ci-gpu-runner ]
+    strategy:
+      matrix:
+        package: [ "zksmith" ]
+    needs: zksync-airbender-build
+    env:
+      BELLMAN_CUDA_DIR: ${{ github.workspace }}/bellman-cuda
+    steps:
+      - name: Prepare environment
+        run: |
+          echo "/usr/local/nvidia/bin:/usr/local/cuda/bin" >> $GITHUB_PATH
+
+      - uses: actions/checkout@93ea575cb5d8a053eaa0ac8fa3b40d7e05a33cc8 # v3
+
+      - name: Check Nvidia driver version
+        run: |
+          nvidia-smi
+
+      - name: Prepare bellman-cuda directory
+        shell: bash
+        # Curl ugliness is required because gh can't work with authentication: https://github.com/cli/cli/issues/2680.
+        run: |
+          release=($(curl --silent https://api.github.com/repos/matter-labs/era-bellman-cuda/releases | jq -r '.[0] | .name, .tarball_url, .assets[0].browser_download_url'))
+          curl --silent -L "${release[1]}" --output bellman-cuda-source.tar.gz
+          curl --silent -L "${release[2]}" --output bellman-cuda.tar.gz
+          mkdir -p bellman-cuda
+          tar xvfz bellman-cuda.tar.gz -C ./bellman-cuda
+          tar xvfz bellman-cuda-source.tar.gz -C ./bellman-cuda --strip-components=1 --wildcards \*/src/
+
+
+      - name: Download test binary built with CUDA 12.5
+        uses: actions/download-artifact@v4
+        with:
+          name: zksync-airbender-12.5.0-devel-ubuntu20.04-${{ matrix.package }}-test-binary
+          path: zksync-airbender-test-binary/12.5/
+
+      - name: Run test binary built with CUDA 12.5
+        id: test_cuda_12_5
+        continue-on-error: true
+        run: |
+          chmod +x zksync-airbender-test-binary/12.5/${{ matrix.package }}
+          zksync-airbender-test-binary/12.5/${{ matrix.package }} --nocapture

tools/zksmith/src/main.rs

@@ -67,11 +67,23 @@ struct LocalProver {
 impl LocalProver {
     fn new(zksync_os_bin_path: String) -> LocalProver {
         let binary = load_binary_from_path(&zksync_os_bin_path);
+        LocalProver::new_internal(binary)
+    }
+
+    #[cfg(test)]
+    fn new_with_binary(binary: &[u8]) -> LocalProver {
+        let padded_binary = execution_utils::get_padded_binary(&binary);
+        LocalProver::new_internal(padded_binary)
+    }
 
+    fn new_internal(padded_binary: Vec<u32>) -> LocalProver {
         let mut gpu_state = GpuSharedState::default();
-        gpu_state.preheat_for_universal_verifier(&binary);
+        gpu_state.preheat_for_universal_verifier(&padded_binary);
 
-        LocalProver { binary, gpu_state }
+        LocalProver {
+            binary: padded_binary,
+            gpu_state,
+        }
     }
 
     fn create_proof_for_data(
@@ -398,3 +410,71 @@ async fn main() {
     println!("Server running at http://127.0.0.1:3030");
     warp::serve(routes).run(([127, 0, 0, 1], 3030)).await;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Instant;
+
+    /// This test requires GPU.
+    #[tokio::test]
+    async fn test_local_prover_with_files() {
+        // These are generated by running some transactions on anvil-zksync and then calling anvil_zks_getBoojumWitness.
+        // We are inlining them here, due to CI (as on CI we build on separate device, and then only ship artifacts to the
+        // gpu device for execution).
+        let test_files = vec![
+            include_str!("../testdata/1.json"),
+            include_str!("../testdata/2.json"),
+            include_str!("../testdata/3.json"),
+            include_str!("../testdata/4.json"),
+        ];
+        let mut timings = Vec::new();
+        let binary = include_bytes!("../../../examples/zksync_os/app.bin");
+
+        let mut prover = LocalProver::new_with_binary(binary);
+
+        for (i, data) in test_files.iter().enumerate() {
+            let parsed_data: Value = serde_json::from_str(&data).expect("Failed to parse JSON");
+            let result_field = parsed_data["result"]
+                .as_str()
+                .expect("Missing or invalid 'result' field")
+                .to_string();
+            let tmp_data = result_field
+                .strip_prefix("0x")
+                .unwrap_or(&result_field)
+                .to_string();
+
+            let batch_id = (i + 1) as u64;
+
+            let start = Instant::now();
+            let (_proof, duration, basic_duration, basic_proof_count, delegation_proof_count) =
+                prover.create_proof_for_data(&tmp_data, batch_id);
+            let elapsed = start.elapsed().as_millis();
+
+            timings.push((
+                batch_id,
+                elapsed,
+                duration,
+                basic_duration,
+                basic_proof_count,
+                delegation_proof_count,
+            ));
+        }
+
+        println!("Timing results:");
+        for (
+            batch_id,
+            elapsed,
+            duration,
+            basic_duration,
+            basic_proof_count,
+            delegation_proof_count,
+        ) in timings
+        {
+            println!(
+                "Batch ID: {}, Total Time: {}ms, Proof Duration: {}ms, Basic Duration: {}ms, Basic Proofs: {}, Delegation Proofs: {:?}",
+                batch_id, elapsed, duration, basic_duration, basic_proof_count, delegation_proof_count
+            );
+        }
+    }
+}