feat: tracegen on GPU (#2034)

Co-authored-by: Grigorii Alekseev <[email protected]>
Co-authored-by: Christian Altamirano <[email protected]>
Co-authored-by: stephenh-axiom-xyz <[email protected]>
Co-authored-by: matejav <[email protected]>
Co-authored-by: Teo Kitanovski <[email protected]>
Co-authored-by: Arayi Khalatyan <[email protected]>
Co-authored-by: Arayi <[email protected]>
Co-authored-by: Golovanov399 <[email protected]>
Co-authored-by: Xinding Wei <[email protected]>

Author: 9 people

.github/workflows/benchmark-call.yml

@@ -10,8 +10,8 @@ on:
       instance_type:
         type: string
         required: false
-        description: The type of runner to start ({1,2,4,8,16,32,48,64}cpu-linux-{arm64,x64})
-        default: 64cpu-linux-arm64
+        description: The AWS instance type or family (e.g., m7g.16xlarge)
+        default: g6.2xlarge
       memory_allocator:
         type: string
         required: false
@@ -63,8 +63,8 @@ on:
       instance_type:
         type: string
         required: false
-        description: The type of runner to start ({1,2,4,8,16,32,48,64}cpu-linux-{arm64,x64})
-        default: 64cpu-linux-arm64
+        description: The AWS instance type or family (e.g., m7g.16xlarge)
+        default: g6.2xlarge
       memory_allocator:
         type: string
         required: false
@@ -107,7 +107,7 @@ on:
 env:
   S3_METRICS_PATH: s3://openvm-public-data-sandbox-us-east-1/benchmark/github/metrics
   S3_FLAMEGRAPHS_PATH: s3://openvm-public-data-sandbox-us-east-1/benchmark/github/flamegraphs
-  FEATURE_FLAGS: "metrics,parallel,nightly-features,tco"
+  FEATURE_FLAGS: "metrics,parallel,tco"
   INPUT_ARGS: ""
   CARGO_NET_GIT_FETCH_WITH_CLI: "true"
 
@@ -118,22 +118,24 @@ jobs:
   bench-new:
     name: Run benchmark on workflow ref/branch
     runs-on:
-      - runs-on
-      - runner=${{ inputs.instance_type }}
-      - run-id=${{ github.run_id }}
-      - family=m7
-      - tag=bench-${{ inputs.benchmark_name }}-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}
+      - runs-on=${{ github.run_id }}-bench-${{ inputs.benchmark_id ||  github.event.inputs.benchmark_name || inputs.benchmark_name }}-${{ github.run_number }}-${{ github.run_attempt }}
+      - family=${{ github.event.inputs.instance_type || inputs.instance_type }}
+      - image=${{ startsWith(github.event.inputs.instance_type || inputs.instance_type, 'g') && 'ubuntu24-gpu-x64' || contains(github.event.inputs.instance_type || inputs.instance_type, 'g.') && 'ubuntu24-full-arm64' || 'ubuntu24-full-x64' }}
+      - spot=capacity-optimized
+      - tag=bench-${{ inputs.benchmark_id || github.event.inputs.benchmark_name || inputs.benchmark_name }}-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}
       - extras=s3-cache
     steps:
-      - uses: runs-on/action@v1
+      - uses: runs-on/action@v2
       ##########################################################################
       # Environment setup                                                      #
       ##########################################################################
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
         with:
           ref: ${{ github.event.pull_request.head.sha || github.sha }}
           repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
-
+      - run: | # avoid cross-device link error
+          rustup component remove clippy || true
+          rm -rf ~/.rustup/toolchains/stable-x86_64-unknown-linux-gnu || true
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:
@@ -146,9 +148,17 @@ jobs:
 
       - name: Display workflow inputs
         run: echo "${{ toJSON(inputs) }}"
-      - name: "Feature flags: memory allocator"
+      - name: "Feature flags: memory allocator + e2e"
         run: |
           ALLOCATOR=${{ inputs.memory_allocator || github.event.inputs.memory_allocator }}
+          if [[ ! "${{ inputs.instance_type || github.event.inputs.instance_type }}" =~ ^g ]]; then
+            FEATURE_FLAGS=${FEATURE_FLAGS},nightly-features
+            if [[ "${{ inputs.e2e_bench }}" == "true" ]]; then
+              FEATURE_FLAGS=${FEATURE_FLAGS},evm
+            fi
+          else
+            FEATURE_FLAGS=${FEATURE_FLAGS},cuda
+          fi
           echo "FEATURE_FLAGS=$ALLOCATOR,${FEATURE_FLAGS}" >> $GITHUB_ENV
       - name: "Feature flags: aggregation"
         if: contains(github.event.pull_request.labels.*.name, 'run-benchmark') || (github.event_name == 'push' && github.ref == 'refs/heads/main')

.github/workflows/benchmarks-execute.yml

@@ -42,8 +42,8 @@ jobs:
       CODSPEED_RUNNER_MODE: walltime
 
     steps:
-      - uses: runs-on/action@v1
-      - uses: actions/checkout@v4
+      - uses: runs-on/action@v2
+      - uses: actions/checkout@v5
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:
@@ -87,8 +87,8 @@ jobs:
       CODSPEED_RUNNER_MODE: instrumentation
 
     steps:
-      - uses: runs-on/action@v1
-      - uses: actions/checkout@v4
+      - uses: runs-on/action@v2
+      - uses: actions/checkout@v5
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:

.github/workflows/benchmarks-upload-fixtures.yml

@@ -17,8 +17,8 @@ jobs:
       - extras=s3-cache
 
     steps:
-      - uses: runs-on/action@v1
-      - uses: actions/checkout@v4
+      - uses: runs-on/action@v2
+      - uses: actions/checkout@v5
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:

.github/workflows/benchmarks.yml

@@ -83,9 +83,6 @@ jobs:
           else
             RUN_E2E=${{ github.event.inputs.run-benchmark-e2e || 'false' }}
           fi
-          if [[ "$RUN_E2E" == "true" ]]; then
-            FEATURE_FLAGS="evm,${FEATURE_FLAGS}"
-          fi
           if [[ "${{ github.event.inputs.aggregation }}" == "true" ]]; then
             FEATURE_FLAGS="aggregation,${FEATURE_FLAGS}"
           fi

.github/workflows/build.yml

@@ -20,8 +20,8 @@ jobs:
       - image=ubuntu24-full-arm64
       - extras=s3-cache
     steps:
-      - uses: runs-on/action@v1
-      - uses: actions/checkout@v4
+      - uses: runs-on/action@v2
+      - uses: actions/checkout@v5
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:

.github/workflows/cli.yml

@@ -42,7 +42,7 @@ jobs:
       - extras=s3-cache
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:

.github/workflows/docs.yml

@@ -25,7 +25,7 @@ jobs:
       - runs-on=${{ github.run_id }}
       - runner=8cpu-linux-arm64
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - name: Set up Rust toolchain
         uses: dtolnay/rust-toolchain@nightly
 

.github/workflows/extension-tests.cuda.yml

@@ -0,0 +1,138 @@
+name: Extension Tests (CUDA)
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["**"]
+    paths:
+      - "crates/circuits/**"
+      - "crates/vm/**"
+      - "crates/toolchain/**"
+      - "extensions/**"
+      - "Cargo.toml"
+      - ".github/workflows/extension-tests.cuda.yml"
+  merge_group:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
+
+env:
+  CARGO_TERM_COLOR: always
+  OPENVM_FAST_TEST: "1"
+
+jobs:
+  extension-tests-cuda:
+    strategy:
+      matrix:
+        extensions: # group extensions on the same runner based on test time
+          - "rv32im native"
+          - "keccak256 sha256 bigint algebra ecc pairing"
+    runs-on:
+      - runs-on=${{ github.run_id }}-extension-tests-cuda-${{ github.run_attempt }}
+      - runner=test-gpu-nvidia
+      - image=ubuntu24-gpu-x64
+      - cpu=8+32 # more cores for faster compilation
+      - spot=capacity-optimized
+      - tag=extension-cuda-tests
+      - extras=s3-cache
+
+    steps:
+      - uses: runs-on/action@v2
+      - uses: actions/checkout@v5
+      - uses: dorny/paths-filter@v3
+        id: filter
+        with:
+          filters: |
+            common:
+              - "crates/circuits/**"
+              - "crates/vm/**"
+              - "crates/toolchain/**"
+              - ".github/workflows/extension-tests.cuda.yml"
+            rv32im:
+              - "extensions/rv32im/**"
+            native:
+              - "extensions/native/**"
+            keccak256:
+              - "extensions/keccak256/**"
+            sha256:
+              - "extensions/sha256/**"
+            bigint:
+              - "extensions/bigint/**"
+            algebra:
+              - "extensions/algebra/**"
+            ecc:
+              - "extensions/ecc/**"
+            pairing:
+              - "extensions/pairing/**"
+
+      - name: Skip if no changes
+        env:
+          FILTER_OUTPUTS: ${{ toJSON(steps.filter.outputs) }}
+        run: |
+          # Check if common files changed
+          COMMON_CHANGED=$(echo $FILTER_OUTPUTS | jq -r .common)
+
+          # Check if any extension in the current matrix group changed
+          ANY_EXT_CHANGED=false
+          for ext in ${{ matrix.extensions }}; do
+            ext_changed=$(echo $FILTER_OUTPUTS | jq -r .$ext)
+            if [[ "$ext_changed" == "true" ]]; then
+              ANY_EXT_CHANGED=true
+              break
+            fi
+          done
+
+          # Skip if neither common nor any extension in this matrix group changed
+          if [[ "$COMMON_CHANGED" != "true" && "$ANY_EXT_CHANGED" != "true" ]]; then
+            echo "No relevant changes for extensions [${{ matrix.extensions }}], skipping tests."
+            exit 0
+          fi
+
+      - run: | # avoid cross-device link error
+          rustup component remove clippy || true
+          rm -rf ~/.rustup/toolchains/stable-x86_64-unknown-linux-gnu || true
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+        with:
+          cache-on-failure: true
+      - uses: taiki-e/install-action@nextest
+
+      - name: Check CUDA status and set environment variables
+        run: |
+          nvcc --version
+          echo "NEXTEST_ENV=CUDA_OPT_LEVEL=1" >> $GITHUB_ENV
+          echo "FEATURE_ARGS=--features=cuda" >> $GITHUB_ENV
+          echo "CIRCUIT_TEST_FILTER=cuda" >> $GITHUB_ENV
+
+      - name: Run extension tests
+        env:
+          FILTER_OUTPUTS: ${{ toJSON(steps.filter.outputs) }}
+        run: |
+          COMMON_CHANGED=$(echo $FILTER_OUTPUTS | jq -r .common)
+
+          for ext in ${{ matrix.extensions }}; do
+            ext_changed=$(echo $FILTER_OUTPUTS | jq -r .$ext)
+            if [[ "$COMMON_CHANGED" == "true" || "$ext_changed" == "true" ]]; then
+              echo "::group::Running tests for $ext"
+
+              # Run circuit tests
+              if [ -d "extensions/$ext/circuit" ]; then
+                pushd extensions/$ext/circuit
+                ${{ env.NEXTEST_ENV }} cargo nextest run ${{ env.CIRCUIT_TEST_FILTER }} ${{ env.FEATURE_ARGS }} --profile=heavy --no-tests=pass
+                popd
+              fi
+
+              # Run integration tests
+              if [ -d "extensions/$ext/tests" ]; then
+                pushd extensions/$ext/tests
+                rustup component add rust-src --toolchain nightly-2025-02-14
+                ${{ env.NEXTEST_ENV }} cargo nextest run ${{ env.FEATURE_ARGS }} --profile=heavy --no-tests=pass
+                popd
+              fi
+              echo "::endgroup::"
+            else
+              echo "Skipping tests for $ext as no relevant files changed."
+            fi
+          done

.github/workflows/extension-tests.yml

@@ -35,27 +35,29 @@ jobs:
           - { name: "algebra", path: "algebra" }
           - { name: "ecc", path: "ecc" }
           - { name: "pairing", path: "pairing" }
+        platform:
+          - { runner: "64cpu-linux-arm64", image: "ubuntu24-full-arm64" }
       # Ensure tests run in parallel even if one fails
       fail-fast: false
 
     runs-on:
-      - runs-on=${{ github.run_id }}
-      - runner=64cpu-linux-arm64
-      - image=ubuntu24-full-arm64
-      - tag=extension-${{ matrix.extension.name }}
+      - runs-on=${{ github.run_id }}-extension-tests-${{ github.run_attempt }}-${{ strategy.job-index }}
+      - runner=${{ matrix.platform.runner }}
+      - image=${{ matrix.platform.image }}
+      - tag=extension-${{ matrix.extension.name }}-${{ matrix.platform.runner }}
       - extras=s3-cache
 
     steps:
-      - uses: runs-on/action@v1
-      - uses: actions/checkout@v4
+      - uses: runs-on/action@v2
+      - uses: actions/checkout@v5
       - uses: dorny/paths-filter@v3
         id: filter
         with:
           filters: |
             - "crates/circuits/**"
             - "crates/vm/**"
             - "crates/toolchain/**"
-            - "extensions/${{ matrix.extensions.path }}/**"
+            - "extensions/${{ matrix.extension.path }}/**"
             - ".github/workflows/extension-tests.yml"
       - name: Skip if no changes
         if: steps.filter.outputs.matched == 'false'
@@ -71,21 +73,27 @@ jobs:
 
       - name: Run ${{ matrix.extension.name }} circuit crate tests
         working-directory: extensions/${{ matrix.extension.path }}/circuit
-        run: cargo nextest run --cargo-profile=fast --test-threads=32
+        run: |
+          TEST_THREADS="--test-threads=32"
+          if [[ "${{ matrix.extension.name }}" == "algebra" ]]; then
+            TEST_THREADS="--test-threads=4"
+          fi
+          cargo nextest run ${{ env.CIRCUIT_TEST_FILTER }} --cargo-profile=fast ${{ env.FEATURE_ARGS }} $TEST_THREADS --no-tests=pass
 
+        # We only need to run guest crate tests on CPU since we are not testing prover functionality
       - name: Run ${{ matrix.extension.name }} guest crate tests
         if: hashFiles(format('extensions/{0}/guest', matrix.extension.path)) != ''
         working-directory: extensions/${{ matrix.extension.path }}/guest
         run: |
-          FEATURE_ARGS=""
+          GUEST_FEATURE_ARGS=""
           if [[ "${{ matrix.extension.name }}" == "pairing" ]]; then
-            FEATURE_ARGS="--features=bn254,bls12_381,halo2curves"
+            GUEST_FEATURE_ARGS="--features=bn254,bls12_381,halo2curves"
           fi
-          cargo nextest run --cargo-profile=fast $FEATURE_ARGS --no-tests=pass
+          cargo nextest run --cargo-profile=fast $GUEST_FEATURE_ARGS --no-tests=pass
 
       - name: Run ${{ matrix.extension.name }} integration tests
         if: hashFiles(format('extensions/{0}/tests', matrix.extension.path)) != ''
         working-directory: extensions/${{ matrix.extension.path }}/tests
         run: |
           rustup component add rust-src --toolchain nightly-2025-02-14
-          cargo nextest run --cargo-profile=fast --profile=heavy --no-tests=pass
+          cargo nextest run --cargo-profile=fast ${{ env.FEATURE_ARGS }} --profile=heavy --no-tests=pass