Dockerfile

# Copyright 2018-2023 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Create basic runtime environment base on Ubuntu 22.04 (jammy)
# Define build arguments
# Create and activate runtime virtual environment
FROM ubuntu:jammy AS base-runtime
ARG AMD_ARCH=AMD_GFX90A
ARG CUDA_ARCH=AMPERE80
ARG CUDA_INSTALLER=https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run
ARG GCC_VERSION=11
ARG KOKKOS_VERSION=4.2.00
ARG LIGHTNING_VERSION=0.34.0
ARG ROCM_INSTALLER=https://repo.radeon.com/amdgpu-install/5.7/ubuntu/jammy/amdgpu-install_5.7.50700-1_all.deb
RUN apt-get update \
    && apt-get install --no-install-recommends -y \
    apt-utils \
    ca-certificates \
    libgomp1 \
    python3 \
    python3-pip \
    python3-venv \
    tzdata \
    wget \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"

# Create basic build environment with build tools and compilers
FROM base-runtime AS base-build
RUN apt-get update \
    && apt-get install --no-install-recommends -y \
    build-essential \
    ccache \
    cmake \
    curl \
    git \
    ninja-build \
    python3-dev \
    gcc-${GCC_VERSION} g++-${GCC_VERSION} cpp-${GCC_VERSION} \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
RUN update-alternatives \
    --install /usr/bin/gcc gcc /usr/bin/gcc-${GCC_VERSION} 100 \
    --slave /usr/bin/g++ g++ /usr/bin/g++-${GCC_VERSION} \
    --slave /usr/bin/gcov gcov /usr/bin/gcov-${GCC_VERSION}
RUN /usr/sbin/update-ccache-symlinks
RUN mkdir /opt/ccache
RUN ccache --set-config=cache_dir=/opt/ccache

# Create and activate build virtual environment
# Install Lightning dev requirements
FROM base-build AS base-build-python
WORKDIR /opt/pennylane-lightning
ENV VIRTUAL_ENV=/opt/venv-build
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN wget --progress=dot:giga https://github.com/PennyLaneAI/pennylane-lightning/archive/refs/tags/v${LIGHTNING_VERSION}.tar.gz
RUN tar -zxf v${LIGHTNING_VERSION}.tar.gz && mv pennylane-lightning-${LIGHTNING_VERSION}/* /opt/pennylane-lightning
RUN pip install --no-cache-dir wheel
RUN pip install --no-cache-dir -r requirements-dev.txt
RUN pip install --no-cache-dir -r doc/requirements.txt

# Download Lightning release and build lightning-qubit backend
FROM base-build-python AS build-wheel-lightning-qubit
WORKDIR /opt/pennylane-lightning
RUN pip uninstall -y pennylane-lightning
RUN python setup.py build_ext
RUN python setup.py bdist_wheel

# Install lightning-qubit backend
FROM base-runtime AS wheel-lightning-qubit
COPY --from=build-wheel-lightning-qubit /opt/pennylane-lightning/dist/ /
RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
RUN pip install --no-cache-dir -U pennylane

# Download Kokkos release and build OpenMP backend
FROM base-build as build-kokkos-openmp
WORKDIR /opt/kokkos-${KOKKOS_VERSION}
RUN wget --progress=dot:giga https://github.com/kokkos/kokkos/archive/refs/tags/${KOKKOS_VERSION}.tar.gz
RUN tar -zxf ${KOKKOS_VERSION}.tar.gz && mv /opt/kokkos-${KOKKOS_VERSION}/kokkos-${KOKKOS_VERSION}/* /opt/kokkos-${KOKKOS_VERSION}
RUN cmake -S /opt/kokkos-${KOKKOS_VERSION} -B /opt/kokkos-${KOKKOS_VERSION}/build \
    -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_INSTALL_PREFIX=/opt/kokkos \
    -DBUILD_SHARED_LIBS:BOOL=ON \
    -DBUILD_TESTING:BOOL=OFF \
    -DKokkos_ENABLE_SERIAL:BOOL=ON \
    -DKokkos_ENABLE_OPENMP:BOOL=ON \
    -DCMAKE_CXX_STANDARD=20 \
    -DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
    -DKokkos_ENABLE_TESTS:BOOL=OFF \
    -DKokkos_ENABLE_LIBDL:BOOL=OFF
RUN cmake --build /opt/kokkos-${KOKKOS_VERSION}/build
RUN cmake --install /opt/kokkos-${KOKKOS_VERSION}/build

# Download Lightning release and build lightning-kokkos backend with Kokkos-OpenMP
FROM base-build-python AS build-wheel-lightning-kokkos-openmp
WORKDIR /opt/pennylane-lightning
COPY --from=build-kokkos-openmp /opt/kokkos /opt/kokkos
ENV CMAKE_PREFIX_PATH=/opt/kokkos:$CMAKE_PREFIX_PATH
RUN pip uninstall -y pennylane-lightning
RUN PL_BACKEND=lightning_kokkos python setup.py build_ext
RUN PL_BACKEND=lightning_kokkos python setup.py bdist_wheel

# Install lightning-kokkos OpenMP backend
FROM base-runtime AS wheel-lightning-kokkos-openmp
COPY --from=build-kokkos-openmp /opt/kokkos /opt/kokkos
COPY --from=build-wheel-lightning-kokkos-openmp /opt/pennylane-lightning/dist/ /
RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
RUN pip install --no-cache-dir -U pennylane

# Install CUDA-12 in build venv image
FROM base-build-python as base-build-cuda
WORKDIR /opt/cuda-build
RUN curl -o cuda-install.run ${CUDA_INSTALLER}
RUN chmod a+x cuda-install.run
RUN ./cuda-install.run --silent --toolkit --toolkitpath=/usr/local/cuda-$(echo ${CUDA_INSTALLER} | grep -o -P '/cuda/.{0,4}' | cut -d / -f 3)
ENV PATH=/usr/local/cuda/bin:${PATH}
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}

# Download Kokkos release and build CUDA backend
FROM base-build-cuda as build-kokkos-cuda
WORKDIR /opt/kokkos-${KOKKOS_VERSION}
# RUN wget --progress=dot:giga https://github.com/kokkos/kokkos/archive/refs/tags/${KOKKOS_VERSION}.tar.gz
# RUN tar -zxf ${KOKKOS_VERSION}.tar.gz && mv /opt/kokkos-${KOKKOS_VERSION}/kokkos-${KOKKOS_VERSION}/* /opt/kokkos-${KOKKOS_VERSION}
RUN git clone --branch lightning https://github.com/vincentmr/kokkos.git /opt/kokkos-${KOKKOS_VERSION}
RUN cmake -S /opt/kokkos-${KOKKOS_VERSION} -B /opt/kokkos-${KOKKOS_VERSION}/build \
    -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_INSTALL_PREFIX=/opt/kokkos \
    -DBUILD_SHARED_LIBS:BOOL=ON \
    -DBUILD_TESTING:BOOL=OFF \
    -DKokkos_ENABLE_SERIAL:BOOL=ON \
    -DKokkos_ENABLE_OPENMP:BOOL=ON \
    -DCMAKE_CXX_STANDARD=20 \
    -DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
    -DKokkos_ENABLE_TESTS:BOOL=OFF \
    -DKokkos_ENABLE_LIBDL:BOOL=OFF \
    -DKokkos_ENABLE_CUDA:BOOL=ON \
    -DKokkos_ARCH_${CUDA_ARCH}=ON
RUN cmake --build /opt/kokkos-${KOKKOS_VERSION}/build
RUN cmake --install /opt/kokkos-${KOKKOS_VERSION}/build

# Download Lightning release and build lightning-kokkos backend with Kokkos-CUDA
FROM base-build-cuda AS build-wheel-lightning-kokkos-cuda
WORKDIR /opt/pennylane-lightning
COPY --from=build-kokkos-cuda /opt/kokkos /opt/kokkos
ENV CMAKE_PREFIX_PATH=/opt/kokkos:$CMAKE_PREFIX_PATH
RUN pip uninstall -y pennylane-lightning
RUN echo >> cmake/support_kokkos.cmake && echo "find_package(CUDAToolkit REQUIRED)" >> cmake/support_kokkos.cmake
RUN PL_BACKEND=lightning_kokkos python setup.py build_ext
RUN PL_BACKEND=lightning_kokkos python setup.py bdist_wheel

# Install python3 and setup runtime virtual env in CUDA-12-runtime image (includes CUDA runtime and math libraries)
# Install lightning-kokkos CUDA backend
FROM nvidia/cuda:12.2.0-base-ubuntu22.04 AS wheel-lightning-kokkos-cuda
RUN apt-get update \
    && apt-get install --no-install-recommends -y \
    libgomp1 \
    python3 \
    python3-pip \
    python3-venv \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
COPY --from=build-kokkos-cuda /opt/kokkos /opt/kokkos
COPY --from=build-wheel-lightning-kokkos-cuda /opt/pennylane-lightning/dist/ /
RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
RUN pip install --no-cache-dir -U pennylane

# Download and build Lightning-GPU release
FROM base-build-cuda AS build-wheel-lightning-gpu
WORKDIR /opt/pennylane-lightning
RUN pip install --no-cache-dir wheel custatevec-cu12
RUN pip uninstall -y pennylane-lightning
RUN CUQUANTUM_SDK=$(python -c "import site; print( f'{site.getsitepackages()[0]}/cuquantum/lib')") PL_BACKEND=lightning_gpu python setup.py build_ext
RUN PL_BACKEND=lightning_gpu python setup.py bdist_wheel

# Install python3 and setup runtime virtual env in CUDA-12-runtime image (includes CUDA runtime and math libraries)
# Install lightning-kokkos CUDA backend
FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS wheel-lightning-gpu
RUN apt-get update \
    && apt-get install --no-install-recommends -y \
    libgomp1 \
    python3 \
    python3-pip \
    python3-venv \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN pip install --no-cache-dir custatevec-cu12
ENV LD_LIBRARY_PATH="$VIRTUAL_ENV/lib/python3.10/site-packages/cuquantum/lib:$LD_LIBRARY_PATH"
COPY --from=build-wheel-lightning-gpu /opt/pennylane-lightning/dist/ /
RUN pip install --no-cache-dir --force-reinstall PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
RUN pip install --no-cache-dir -U pennylane

# Install ROCm in build venv image
FROM base-build-python as base-build-rocm
RUN wget --progress=dot:giga ${ROCM_INSTALLER}
RUN apt-get update \
    && apt-get install --no-install-recommends -y \
    "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)" \
    ./$(echo ${ROCM_INSTALLER} | xargs -I {} basename {}) \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
RUN amdgpu-install -y --accept-eula --usecase=rocm,hiplibsdk --no-dkms

# Download Kokkos release and build HIP backend
FROM base-build-rocm as build-kokkos-rocm
WORKDIR /opt/kokkos-${KOKKOS_VERSION}
ENV LD_LIBRARY_PATH=/opt/rocm/lib:$LD_LIBRARY_PATH
# RUN wget --progress=dot:giga https://github.com/kokkos/kokkos/archive/refs/tags/${KOKKOS_VERSION}.tar.gz
# RUN tar -zxf ${KOKKOS_VERSION}.tar.gz && mv /opt/kokkos-${KOKKOS_VERSION}/kokkos-${KOKKOS_VERSION}/* /opt/kokkos-${KOKKOS_VERSION}
RUN git clone --branch fix-rtld-deepbind https://github.com/cz4rs/kokkos.git /opt/kokkos-${KOKKOS_VERSION}
RUN CXX=hipcc cmake -S /opt/kokkos-${KOKKOS_VERSION} -B /opt/kokkos-${KOKKOS_VERSION}/build \
    -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_INSTALL_PREFIX=/opt/kokkos \
    -DBUILD_SHARED_LIBS:BOOL=ON \
    -DBUILD_TESTING:BOOL=OFF \
    -DKokkos_ENABLE_SERIAL:BOOL=ON \
    -DKokkos_ENABLE_OPENMP:BOOL=ON \
    -DCMAKE_CXX_STANDARD=20 \
    -DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
    -DKokkos_ENABLE_TESTS:BOOL=OFF \
    -DKokkos_ENABLE_LIBDL:BOOL=OFF \
    -DKokkos_ENABLE_HIP:BOOL=ON \
    -DKokkos_ARCH_${AMD_ARCH}=ON
RUN cmake --build /opt/kokkos-${KOKKOS_VERSION}/build
RUN cmake --install /opt/kokkos-${KOKKOS_VERSION}/build

# Download Lightning release and build lightning-kokkos backend with Kokkos-ROCm
FROM base-build-rocm AS build-wheel-lightning-kokkos-rocm
WORKDIR /opt/pennylane-lightning
COPY --from=build-kokkos-rocm /opt/kokkos /opt/kokkos
ENV CXX=hipcc
ENV CMAKE_PREFIX_PATH=/opt/kokkos:$CMAKE_PREFIX_PATH
RUN pip uninstall -y pennylane-lightning
RUN PL_BACKEND=lightning_kokkos python setup.py build_ext
RUN PL_BACKEND=lightning_kokkos python setup.py bdist_wheel

# Install lightning-kokkos HIP backend
FROM rocm/dev-ubuntu-22.04:5.7 AS wheel-lightning-kokkos-rocm
RUN apt-get update \
    && apt-get install --no-install-recommends -y \
    libgomp1 \
    python3 \
    python3-pip \
    python3-venv \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
COPY --from=build-kokkos-rocm /opt/kokkos /opt/kokkos
COPY --from=build-wheel-lightning-kokkos-rocm /opt/pennylane-lightning/dist/ /
RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
RUN pip install --no-cache-dir -U pennylane