feat(prover): add GPU feature for compressor

.github/workflows/build-prover-template.yml

@@ -41,7 +41,7 @@ jobs:
       RUNNER_COMPOSE_FILE: "docker-compose-runner-nightly.yml"
       ERA_BELLMAN_CUDA_RELEASE: ${{ inputs.ERA_BELLMAN_CUDA_RELEASE }}
       CUDA_ARCH: ${{ inputs.CUDA_ARCH }}
-    runs-on: [matterlabs-ci-runner]
+    runs-on: [ matterlabs-ci-runner ]
     strategy:
       matrix:
         component:
@@ -82,6 +82,7 @@ jobs:
       - name: download CRS
         if: matrix.component == 'proof-fri-compressor'
         run: |
+          ci_run curl -LO https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2\^24.key
           ci_run curl -LO https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2\^26.key
 
       - name: login to Docker registries
@@ -137,7 +138,7 @@ jobs:
         env:
           DOCKER_ACTION: ${{ inputs.action }}
           COMPONENT: ${{ matrix.component }}
-        run: |          
+        run: |
           PASSED_ENV_VARS="ERA_BELLMAN_CUDA_RELEASE,CUDA_ARCH" \
           ci_run zk docker $DOCKER_ACTION $COMPONENT
 

docker/proof-fri-compressor/Dockerfile

@@ -1,24 +1,46 @@
 # Will work locally only after prior universal setup key download
-
-FROM matterlabs/zksync-build-base:latest as builder
+FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 as builder
 
 ARG DEBIAN_FRONTEND=noninteractive
 
+ARG CUDA_ARCH=89
+ENV CUDAARCHS=${CUDA_ARCH}
+
+RUN apt-get update && apt-get install -y curl clang openssl libssl-dev gcc g++ git \
+    pkg-config build-essential libclang-dev && \
+    rm -rf /var/lib/apt/lists/*
+
+ENV RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    PATH=/usr/local/cargo/bin:$PATH
+
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y && \
+    rustup install nightly-2023-08-21 && \
+    rustup default nightly-2023-08-21
+
+RUN curl -Lo cmake-3.24.2-linux-x86_64.sh https://github.com/Kitware/CMake/releases/download/v3.24.2/cmake-3.24.2-linux-x86_64.sh && \
+    chmod +x cmake-3.24.2-linux-x86_64.sh && \
+    ./cmake-3.24.2-linux-x86_64.sh --skip-license --prefix=/usr/local
+
 WORKDIR /usr/src/zksync
 COPY . .
 
-RUN cd prover && cargo build --release --bin zksync_proof_fri_compressor
+RUN mkdir keys && cd keys && mkdir setup && cd setup && curl -LO https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2^24.key
+
+RUN cd prover && \
+    git clone https://github.com/matter-labs/era-bellman-cuda.git --branch main bellman-cuda && \
+    cmake -Bbellman-cuda/build -Sbellman-cuda/ -DCMAKE_BUILD_TYPE=Release && \
+    cmake --build bellman-cuda/build/
 
-FROM debian:bookworm-slim
+RUN cd prover && BELLMAN_CUDA_DIR=$PWD/bellman-cuda cargo build --features "gpu" --release --bin zksync_proof_fri_compressor
+
+FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04
 
 RUN apt-get update && apt-get install -y curl libpq5 ca-certificates && rm -rf /var/lib/apt/lists/*
 
 # copy VK required for proof wrapping
 COPY prover/vk_setup_data_generator_server_fri/data/ /prover/vk_setup_data_generator_server_fri/data/
 
-# copy universal setup key required for proof compression
-COPY setup_2\^26.key /setup_2\^26.key
-
 COPY --from=builder /usr/src/zksync/prover/target/release/zksync_proof_fri_compressor /usr/bin/
 
 ENTRYPOINT ["zksync_proof_fri_compressor"]

docker/proof-gpu-fri-compressor/Dockerfile

@@ -0,0 +1,47 @@
+# Will work locally only after prior universal setup key download
+FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 as builder
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+ARG CUDA_ARCH=89
+ENV CUDAARCHS=${CUDA_ARCH}
+
+RUN apt-get update && apt-get install -y curl clang openssl libssl-dev gcc g++ \
+    pkg-config build-essential libclang-dev && \
+    rm -rf /var/lib/apt/lists/*
+
+ENV RUSTUP_HOME=/usr/local/rustup \
+    CARGO_HOME=/usr/local/cargo \
+    PATH=/usr/local/cargo/bin:$PATH
+
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y && \
+    rustup install nightly-2023-08-21 && \
+    rustup default nightly-2023-08-21
+
+RUN curl -Lo cmake-3.24.2-linux-x86_64.sh https://github.com/Kitware/CMake/releases/download/v3.24.2/cmake-3.24.2-linux-x86_64.sh && \
+    chmod +x cmake-3.24.2-linux-x86_64.sh && \
+    ./cmake-3.24.2-linux-x86_64.sh --skip-license --prefix=/usr/local
+
+WORKDIR /usr/src/zksync
+COPY . .
+
+RUN cd prover && \
+    git clone https://github.com/matter-labs/bellman-cuda.git --branch dev bellman-cuda && \
+    cmake -Bbellman-cuda/build -Sbellman-cuda/ -DCMAKE_BUILD_TYPE=Release && \
+    cmake --build bellman-cuda/build/ && \
+    export BELLMAN_CUDA_DIR=$PWD/bellman-cuda
+RUN cd prover && cargo build --features "gpu" --release --bin zksync_proof_fri_compressor
+
+FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04
+
+RUN apt-get update && apt-get install -y curl libpq5 ca-certificates && rm -rf /var/lib/apt/lists/*
+
+# copy VK required for proof wrapping
+COPY prover/vk_setup_data_generator_server_fri/data/ /prover/vk_setup_data_generator_server_fri/data/
+
+# copy universal setup key required for proof compression
+COPY setup_2\^24.key /setup_2\^24.key
+
+COPY --from=builder /usr/src/zksync/prover/target/release/zksync_proof_fri_compressor /usr/bin/
+
+ENTRYPOINT ["zksync_proof_fri_compressor"]

etc/env/base/fri_proof_compressor.toml

@@ -1,10 +1,10 @@
 [fri_proof_compressor]
-compression_mode=1
-prometheus_listener_port=3321
-prometheus_pushgateway_url="http://127.0.0.1:9091"
-prometheus_push_interval_ms=100
-generation_timeout_in_secs=3600
-max_attempts=5
-universal_setup_path="../keys/setup/setup_2^26.key"
-universal_setup_download_url="https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2^26.key"
-verify_wrapper_proof=true
+compression_mode = 1
+prometheus_listener_port = 3321
+prometheus_pushgateway_url = "http://127.0.0.1:9091"
+prometheus_push_interval_ms = 100
+generation_timeout_in_secs = 3600
+max_attempts = 5
+universal_setup_path = "../keys/setup/setup_2^24.key"
+universal_setup_download_url = "https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2^24.key"
+verify_wrapper_proof = true