feat(prover): Add cpu-target for witness_generator compilation

.github/workflows/build-docker-from-tag.yml

@@ -94,6 +94,20 @@ jobs:
       DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
       DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
 
+  build-push-witness-generator-image-avx512:
+    name: Build and push image
+    needs: [setup]
+    uses: ./.github/workflows/build-witness-generator-template.yml
+    if: contains(github.ref_name, 'prover')
+    with:
+      image_tag_suffix: ${{ needs.setup.outputs.image_tag_suffix }}-avx512
+      ERA_BELLMAN_CUDA_RELEASE: ${{ vars.ERA_BELLMAN_CUDA_RELEASE }}
+      CUDA_ARCH: "60;70;75;89"
+      WITNESS_GENERATOR_RUST_FLAGS: "-Ctarget_feature=+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl"
+    secrets:
+      DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
+      DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
+
   build-gar-prover-fri-gpu:
     name: Build GAR prover FRI GPU
     needs: [setup, build-push-prover-images]

.github/workflows/build-witness-generator-template.yml

@@ -0,0 +1,206 @@
+name: Build witness generator image with custom compiler flags
+on:
+  workflow_call:
+    secrets:
+      DOCKERHUB_USER:
+        description: "DOCKERHUB_USER"
+        required: true
+      DOCKERHUB_TOKEN:
+        description: "DOCKERHUB_TOKEN"
+        required: true
+    inputs:
+      ERA_BELLMAN_CUDA_RELEASE:
+        description: "ERA_BELLMAN_CUDA_RELEASE"
+        type: string
+        required: true
+      image_tag_suffix:
+        description: "Optional suffix to override tag name generation"
+        type: string
+        required: false
+      action:
+        description: "Action with docker image"
+        type: string
+        default: "push"
+        required: false
+      is_pr_from_fork:
+        description: "Indicates whether the workflow is invoked from a PR created from fork"
+        type: boolean
+        default: false
+        required: false
+      CUDA_ARCH:
+        description: "CUDA Arch to build"
+        type: string
+        default: "89"
+        required: false
+      WITNESS_GENERATOR_RUST_FLAGS:
+        description: "Rust flags for witness_generator compilation"
+        type: string
+        default: ""
+        required: false
+    outputs:
+      protocol_version:
+        description: "Protocol version of the binary"
+        value: ${{ jobs.build-images.outputs.protocol_version }}
+
+jobs:
+  build-images:
+    name: Build and Push Docker Images
+    env:
+      IMAGE_TAG_SUFFIX: ${{ inputs.image_tag_suffix }}
+      RUNNER_COMPOSE_FILE: "docker-compose-runner-nightly.yml"
+      ERA_BELLMAN_CUDA_RELEASE: ${{ inputs.ERA_BELLMAN_CUDA_RELEASE }}
+      CUDA_ARCH: ${{ inputs.CUDA_ARCH }}
+      WITNESS_GENERATOR_RUST_FLAGS: ${{ inputs.WITNESS_GENERATOR_RUST_FLAGS }}
+    runs-on: [ matterlabs-ci-runner-c3d ]
+    strategy:
+      matrix:
+        component:
+          - witness-generator
+    outputs:
+      protocol_version: ${{ steps.protocolversion.outputs.protocol_version }}
+    steps:
+      - uses: actions/checkout@a5ac7e51b41094c92402da3b24376905380afc29 # v4
+        with:
+          submodules: "recursive"
+
+      - name: setup-env
+        run: |
+          echo ZKSYNC_HOME=$(pwd) >> $GITHUB_ENV
+          echo CI=1 >> $GITHUB_ENV
+          echo $(pwd)/bin >> $GITHUB_PATH
+          echo CI=1 >> .env
+          echo IN_DOCKER=1 >> .env
+
+      - name: start-services
+        run: |
+          echo "IMAGE_TAG_SUFFIX=${{ env.IMAGE_TAG_SUFFIX }}" >> .env
+          mkdir -p ./volumes/postgres
+          run_retried docker compose pull zk postgres
+          docker compose up -d zk postgres
+          ci_run sccache --start-server
+
+      - name: init
+        run: |
+          ci_run git config --global --add safe.directory /usr/src/zksync
+          ci_run git config --global --add safe.directory /usr/src/zksync/contracts/system-contracts
+          ci_run git config --global --add safe.directory /usr/src/zksync/contracts
+          ci_run zk
+
+      - name: download CRS for GPU compressor
+        if: matrix.component == 'proof-fri-gpu-compressor'
+        run: |
+          ci_run run_retried curl -LO https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2\^24.key
+
+
+      - name: login to Docker registries
+        if: github.event_name != 'pull_request' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
+        run: |
+          ci_run docker login -u ${{ secrets.DOCKERHUB_USER }} -p ${{ secrets.DOCKERHUB_TOKEN }}
+          ci_run gcloud auth configure-docker us-docker.pkg.dev -q
+
+      # We need to run this only when ERA_BELLMAN_CUDA_RELEASE is not available
+      # In our case it happens only when PR is created from fork
+      - name: Wait for runner IP to be not rate-limited against GH API
+        if: inputs.is_pr_from_fork == true
+        run: |
+          api_endpoint="https://api.github.com/users/zksync-era-bot"
+          wait_time=60
+          max_retries=60
+          retry_count=0
+
+          while [[ $retry_count -lt $max_retries ]]; do
+            response=$(run_retried curl -s -w "%{http_code}" -o temp.json "$api_endpoint")
+            http_code=$(echo "$response" | tail -n1)
+
+            if [[ "$http_code" == "200" ]]; then
+              echo "Request successful. Not rate-limited."
+              cat temp.json
+              rm temp.json
+              exit 0
+            elif [[ "$http_code" == "403" ]]; then
+              rate_limit_exceeded=$(jq -r '.message' temp.json | grep -i "API rate limit exceeded")
+              if [[ -n "$rate_limit_exceeded" ]]; then
+                retry_count=$((retry_count+1))
+                echo "API rate limit exceeded. Retry $retry_count of $max_retries. Retrying in $wait_time seconds..."
+                sleep $wait_time
+              else
+                echo "Request failed with HTTP status $http_code."
+                cat temp.json
+                rm temp.json
+                exit 1
+              fi
+            else
+              echo "Request failed with HTTP status $http_code."
+              cat temp.json
+              rm temp.json
+              exit 1
+            fi
+          done
+
+          echo "Reached the maximum number of retries ($max_retries). Exiting."
+          rm temp.json
+          exit 1
+
+      - name: protocol-version
+        id: protocolversion
+        # TODO: use -C flag, when it will become stable.
+        shell: bash
+        run: |
+          ci_run bash -c "cd prover && cargo build --release --bin prover_version"
+          PPV=$(ci_run prover/target/release/prover_version)
+          echo Protocol version is ${PPV}
+          echo "protocol_version=${PPV}" >> $GITHUB_OUTPUT
+          echo "PROTOCOL_VERSION=${PPV}" >> $GITHUB_ENV
+
+      - name: setup-rust-flags-env
+        if: matrix.component == 'witness-generator'
+        run: |
+          echo RUST_FLAGS="${{ env.WITNESS_GENERATOR_RUST_FLAGS }}" >> $GITHUB_ENV
+
+      - name: update-images
+        env:
+          DOCKER_ACTION: ${{ inputs.action }}
+          COMPONENT: ${{ matrix.component }}
+        run: |
+          PASSED_ENV_VARS="ERA_BELLMAN_CUDA_RELEASE,CUDA_ARCH,PROTOCOL_VERSION,RUST_FLAGS" \
+          ci_run zk docker $DOCKER_ACTION $COMPONENT
+
+      - name: Show sccache stats
+        if: always()
+        run: |
+          ci_run sccache --show-stats || true
+          ci_run cat /tmp/sccache_log.txt || true
+
+  copy-images:
+    name: Copy images between docker registries
+    needs: build-images
+    env:
+      IMAGE_TAG_SUFFIX: ${{ inputs.image_tag_suffix }}
+      PROTOCOL_VERSION: ${{ needs.build-images.outputs.protocol_version }}
+    runs-on: matterlabs-ci-runner
+    if: ${{ inputs.action == 'push' }}
+    strategy:
+      matrix:
+        component:
+          - witness-vector-generator
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to us-central1 GAR
+        run: |
+          gcloud auth print-access-token --lifetime=7200 --impersonate-service-account=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com | docker login -u oauth2accesstoken --password-stdin https://us-docker.pkg.dev
+
+      - name: Login and push to Asia GAR
+        run: |
+          gcloud auth print-access-token --lifetime=7200 --impersonate-service-account=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com | docker login -u oauth2accesstoken --password-stdin https://asia-docker.pkg.dev
+          docker buildx imagetools create \
+            --tag asia-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }} \
+            us-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }}
+
+      - name: Login and push to Europe GAR
+        run: |
+          gcloud auth print-access-token --lifetime=7200 --impersonate-service-account=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com | docker login -u oauth2accesstoken --password-stdin https://europe-docker.pkg.dev
+          docker buildx imagetools create \
+            --tag europe-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }} \
+            us-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }}

.github/workflows/ci.yml

@@ -166,6 +166,21 @@ jobs:
       DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
       DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
 
+  build-witness-generator-image-avx512:
+    name: Build prover images with avx512 instructions
+    needs: changed_files
+    if: ${{ (needs.changed_files.outputs.prover == 'true' || needs.changed_files.outputs.all == 'true') && !contains(github.ref_name, 'release-please--branches') }}
+    uses: ./.github/workflows/build-witness-generator-template.yml
+    with:
+      image_tag_suffix: ${{ needs.setup.outputs.image_tag_suffix }}-avx512
+      action: "build"
+      ERA_BELLMAN_CUDA_RELEASE: ${{ vars.ERA_BELLMAN_CUDA_RELEASE }}
+      is_pr_from_fork: ${{ github.event.pull_request.head.repo.fork == true }}
+      WITNESS_GENERATOR_RUST_FLAGS: "-Ctarget_feature=+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl"
+    secrets:
+      DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
+      DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
+
   ci-success:
     name: Github Status Check
     runs-on: ubuntu-latest

docker/witness-generator/Dockerfile

@@ -1,6 +1,8 @@
 FROM matterlabs/zksync-build-base:latest AS builder
 
 ARG DEBIAN_FRONTEND=noninteractive
+ARG RUST_FLAGS=""
+ENV RUSTFLAGS=${RUST_FLAGS}
 
 WORKDIR /usr/src/zksync
 COPY . .

infrastructure/zk/src/docker.ts

@@ -116,6 +116,12 @@ async function _build(image: string, tagList: string[], dockerOrg: string, platf
         const cudaArch = process.env.CUDA_ARCH;
         buildArgs += `--build-arg CUDA_ARCH='${cudaArch}' `;
     }
+    if (image === 'witness-generator') {
+        const rustFlags = process.env.RUST_FLAGS;
+        if (rustFlags) {
+            buildArgs += `--build-arg RUST_FLAGS='${rustFlags}' `;
+        }
+    }
     buildArgs += extraArgs;
 
     const buildCommand =