ci: update CI to build cuda enabled images for cuda compute cap 89,90,120

.github/workflows/test-zkvm-airbender.yml

@@ -21,3 +21,4 @@ jobs:
     with:
       zkvm: airbender
       skip_prove_test: true
+      cuda: true

.github/workflows/test-zkvm-openvm.yml

@@ -20,3 +20,4 @@ jobs:
       packages: write
     with:
       zkvm: openvm
+      cuda: true

.github/workflows/test-zkvm-risc0.yml

@@ -20,3 +20,4 @@ jobs:
       packages: write
     with:
       zkvm: risc0
+      cuda: true

.github/workflows/test-zkvm-sp1.yml

@@ -20,3 +20,4 @@ jobs:
       packages: write
     with:
       zkvm: sp1
+      cuda: true

.github/workflows/test-zkvm-zisk.yml

@@ -20,4 +20,5 @@ jobs:
       packages: write
     with:
       zkvm: zisk
+      cuda: true
       skip_prove_test: true

.github/workflows/test-zkvm.yml

@@ -7,6 +7,11 @@ on:
         description: 'zkVM to test'
         required: true
         type: string
+      cuda:
+        description: 'Whether to build CUDA-enabled images'
+        required: false
+        type: boolean
+        default: false
       # Remove when we use larger runners, currently only needed to skip for zisk
       skip_prove_test:
         description: 'Whether to skip prove test and ere-dockerized test or not'
@@ -55,13 +60,31 @@ jobs:
               --tag "$IMAGE_REGISTRY/$IMAGE:$SEMVER_TAG"
           done
 
-  build_image:
-    name: Build image
-    if: github.event_name == 'pull_request' || (github.event_name == 'push' && github.ref == 'refs/heads/master')
+      - name: Add SemVer tag to CUDA images
+        if: inputs.cuda
+        run: |
+          GIT_SHA="${{ github.sha }}"
+          GIT_SHA_TAG="${GIT_SHA:0:7}"
+
+          GIT_TAG="${{ github.ref_name }}"
+          SEMVER_TAG="${GIT_TAG#v}"
+
+          IMAGE_REGISTRY="ghcr.io/${{ github.repository }}"
+
+          for IMAGE in \
+            "ere-base" \
+            "ere-base-${{ inputs.zkvm }}" \
+            "ere-server-${{ inputs.zkvm }}"
+          do
+            echo "Tagging $IMAGE_REGISTRY/$IMAGE:${GIT_SHA_TAG}-cuda as ${SEMVER_TAG}-cuda"
+            docker buildx imagetools create \
+              "$IMAGE_REGISTRY/$IMAGE:${GIT_SHA_TAG}-cuda" \
+              --tag "$IMAGE_REGISTRY/$IMAGE:${SEMVER_TAG}-cuda"
+          done
+
+  get_image_metadata:
+    name: Get image metadata
     runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
     outputs:
       image_registry: ${{ steps.image_meta.outputs.image_registry }}
       image_tag: ${{ steps.image_meta.outputs.image_tag }}
@@ -70,22 +93,13 @@ jobs:
       base_zkvm_image: ${{ steps.image_meta.outputs.base_zkvm_image }}
       compiler_zkvm_image: ${{ steps.image_meta.outputs.compiler_zkvm_image }}
       server_zkvm_image: ${{ steps.image_meta.outputs.server_zkvm_image }}
+      base_image_cuda: ${{ steps.image_meta.outputs.base_image_cuda }}
+      base_zkvm_image_cuda: ${{ steps.image_meta.outputs.base_zkvm_image_cuda }}
+      server_zkvm_image_cuda: ${{ steps.image_meta.outputs.server_zkvm_image_cuda }}
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
 
-      - name: Free up disk space
-        if: github.event_name == 'push'
-        run: bash .github/scripts/free-up-disk-space.sh
-
-      - name: Log in to GitHub Container Registry
-        if: github.event_name == 'push'
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
       - name: Check Dockerfile changes
         id: changed_files
         uses: tj-actions/changed-files@v46
@@ -115,45 +129,123 @@ jobs:
           echo "image_registry=$IMAGE_REGISTRY" >> $GITHUB_OUTPUT
           echo "image_tag=$IMAGE_TAG" >> $GITHUB_OUTPUT
           echo "cached_image_tag=$CACHED_IMAGE_TAG" >> $GITHUB_OUTPUT
+
           echo "base_image=$BASE_IMAGE" >> $GITHUB_OUTPUT
           echo "base_zkvm_image=$BASE_ZKVM_IMAGE" >> $GITHUB_OUTPUT
           echo "compiler_zkvm_image=$COMPILER_ZKVM_IMAGE" >> $GITHUB_OUTPUT
           echo "server_zkvm_image=$SERVER_ZKVM_IMAGE" >> $GITHUB_OUTPUT
 
-      - name: Build and push ere-base and ere-base-${{ inputs.zkvm }} images
-        if: github.event_name == 'push'
+          echo "base_image_cuda=$IMAGE_REGISTRY/ere-base:${IMAGE_TAG}-cuda" >> $GITHUB_OUTPUT
+          echo "base_zkvm_image_cuda=$IMAGE_REGISTRY/ere-base-${{ inputs.zkvm }}:${IMAGE_TAG}-cuda" >> $GITHUB_OUTPUT
+          echo "server_zkvm_image_cuda=$IMAGE_REGISTRY/ere-server-${{ inputs.zkvm }}:${IMAGE_TAG}-cuda" >> $GITHUB_OUTPUT
+
+  build_image:
+    name: Build image
+    if: github.event_name == 'pull_request' || (github.event_name == 'push' && github.ref == 'refs/heads/master')
+    needs: get_image_metadata
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Free up disk space
+        run: bash .github/scripts/free-up-disk-space.sh
+
+      - name: Log in to GitHub Container Registry
+        if: github.event_name == 'push' && github.ref == 'refs/heads/master'
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build ere-base and ere-base-${{ inputs.zkvm }} images
         run: |
           bash .github/scripts/build-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ steps.image_meta.outputs.image_registry }} \
-            --tag ${{ steps.image_meta.outputs.image_tag }} \
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
             --base
-          docker push ${{ steps.image_meta.outputs.base_image }}
-          docker push ${{ steps.image_meta.outputs.base_zkvm_image }}
 
-      - name: Build and push ere-compiler-${{ inputs.zkvm }} image
-        if: github.event_name == 'push'
+      - name: Push ere-base and ere-base-${{ inputs.zkvm }} images
+        if: github.event_name == 'push' && github.ref == 'refs/heads/master'
+        run: |
+          docker push ${{ needs.get_image_metadata.outputs.base_image }}
+          docker push ${{ needs.get_image_metadata.outputs.base_zkvm_image }}
+
+      - name: Build ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }} images
         run: |
           bash .github/scripts/build-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ steps.image_meta.outputs.image_registry }} \
-            --tag ${{ steps.image_meta.outputs.image_tag }} \
-            --compiler
-          docker push ${{ steps.image_meta.outputs.compiler_zkvm_image }}
-
-      - name: Build and push ere-server-${{ inputs.zkvm }} image
-        if: github.event_name == 'push'
-        run: |
-          bash .github/scripts/build-image.sh \
-            --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ steps.image_meta.outputs.image_registry }} \
-            --tag ${{ steps.image_meta.outputs.image_tag }} \
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
+            --compiler \
             --server
-          docker push ${{ steps.image_meta.outputs.server_zkvm_image }}
+
+      - name: Push ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }} images
+        if: github.event_name == 'push' && github.ref == 'refs/heads/master'
+        run: |
+          docker push ${{ needs.get_image_metadata.outputs.compiler_zkvm_image }}
+          docker push ${{ needs.get_image_metadata.outputs.server_zkvm_image }}
+
+  build_image_cuda:
+    name: Build image (CUDA)
+    if: inputs.cuda && (github.event_name == 'pull_request' || (github.event_name == 'push' && github.ref == 'refs/heads/master'))
+    needs: get_image_metadata
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Free up disk space
+        run: bash .github/scripts/free-up-disk-space.sh
+
+      - name: Log in to GitHub Container Registry
+        if: github.event_name == 'push' && github.ref == 'refs/heads/master'
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build ere-base and ere-base-${{ inputs.zkvm }} CUDA images
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ inputs.zkvm }} \
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
+            --base \
+            --cuda-archs '89,120'
+
+      - name: Push ere-base and ere-base-${{ inputs.zkvm }} CUDA images
+        if: github.event_name == 'push' && github.ref == 'refs/heads/master'
+        run: |
+          docker push ${{ needs.get_image_metadata.outputs.base_image_cuda }}
+          docker push ${{ needs.get_image_metadata.outputs.base_zkvm_image_cuda }}
+
+      - name: Build ere-server-${{ inputs.zkvm }} CUDA image
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ inputs.zkvm }} \
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
+            --server \
+            --cuda-archs '89,120'
+
+      - name: Push ere-server-${{ inputs.zkvm }} CUDA image
+        if: github.event_name == 'push' && github.ref == 'refs/heads/master'
+        run: |
+          docker push ${{ needs.get_image_metadata.outputs.server_zkvm_image_cuda }}
 
   clippy_via_docker:
     name: Clippy via Docker
-    needs: build_image
+    needs: get_image_metadata
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -168,17 +260,12 @@ jobs:
           key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}
 
       - name: Pull base zkvm image or build locally
-        if: github.event_name == 'pull_request'
         run: |
           bash .github/scripts/pull-or-build-base-zkvm-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
-            --cached-tag "${{ needs.build_image.outputs.cached_image_tag }}"
-
-      - name: Pull base zkvm image
-        if: github.event_name == 'push'
-        run: docker pull ${{ needs.build_image.outputs.base_zkvm_image }}
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
+            --cached-tag "${{ needs.get_image_metadata.outputs.cached_image_tag }}"
 
       - name: Run cargo clippy for ere-${{ inputs.zkvm }} via Docker
         run: |
@@ -189,7 +276,7 @@ jobs:
             --volume $HOME/.cargo/registry:/usr/local/cargo/registry \
             --volume $HOME/.cargo/git:/usr/local/cargo/git \
             --workdir /ere \
-            ${{ needs.build_image.outputs.base_zkvm_image }} \
+            ${{ needs.get_image_metadata.outputs.base_zkvm_image }} \
             /bin/bash"
 
           cat <<EOF | $DOCKER_CMD
@@ -207,7 +294,7 @@ jobs:
 
   test_via_docker:
     name: Test via Docker
-    needs: build_image
+    needs: get_image_metadata
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -222,17 +309,12 @@ jobs:
           key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}
 
       - name: Pull base zkvm image or build locally
-        if: github.event_name == 'pull_request'
         run: |
           bash .github/scripts/pull-or-build-base-zkvm-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
-            --cached-tag "${{ needs.build_image.outputs.cached_image_tag }}"
-
-      - name: Pull base zkvm image
-        if: github.event_name == 'push'
-        run: docker pull ${{ needs.build_image.outputs.base_zkvm_image }}
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
+            --cached-tag "${{ needs.get_image_metadata.outputs.cached_image_tag }}"
 
       - name: Run cargo test for ere-${{ inputs.zkvm }} via Docker
         run: |
@@ -243,7 +325,7 @@ jobs:
             --volume $HOME/.cargo/registry:/usr/local/cargo/registry \
             --volume $HOME/.cargo/git:/usr/local/cargo/git \
             --workdir /ere \
-            ${{ needs.build_image.outputs.base_zkvm_image }} \
+            ${{ needs.get_image_metadata.outputs.base_zkvm_image }} \
             /bin/bash"
 
           cat <<EOF | $DOCKER_CMD
@@ -259,7 +341,7 @@ jobs:
 
   test_ere_dockerized:
     name: Test ere-dockerized with the selected zkVM
-    needs: build_image
+    needs: get_image_metadata
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -279,31 +361,24 @@ jobs:
           key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}
 
       - name: Pull images or build locally
-        if: github.event_name == 'pull_request'
         run: |
           bash .github/scripts/pull-or-build-base-zkvm-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
-            --cached-tag "${{ needs.build_image.outputs.cached_image_tag }}"
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
+            --cached-tag "${{ needs.get_image_metadata.outputs.cached_image_tag }}"
 
           # Build ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }}
           bash .github/scripts/build-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
+            --registry ${{ needs.get_image_metadata.outputs.image_registry }} \
+            --tag ${{ needs.get_image_metadata.outputs.image_tag }} \
             --compiler \
             --server
 
-      - name: Pull ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }} images
-        if: github.event_name == 'push'
-        run: |
-          docker pull ${{ needs.build_image.outputs.compiler_zkvm_image }}
-          docker pull ${{ needs.build_image.outputs.server_zkvm_image }}
-
       - name: Run cargo test for ere-${{ inputs.zkvm }} via ere-dockerized
         env:
-          ERE_IMAGE_REGISTRY: ${{ needs.build_image.outputs.image_registry }}
+          ERE_IMAGE_REGISTRY: ${{ needs.get_image_metadata.outputs.image_registry }}
         run: |
           cargo test --release --package ere-dockerized \
             -- ${{ inputs.zkvm }} ${{ inputs.skip_prove_test && '--skip prove' || ''  }} --test-threads=1