ARG IMAGE_NAME
FROM ${IMAGE_NAME}:11.1.1-base-centos8 as base

ENV NV_CUDA_LIB_VERSION 11.1.1-1

FROM base as base-amd64

ENV NV_NVTX_VERSION 11.1.74-1
ENV NV_LIBNPP_VERSION 11.1.2.301-1
ENV NV_LIBCUBLAS_VERSION 11.3.0.106-1
ENV NV_LIBNCCL_PACKAGE_NAME libnccl
ENV NV_LIBNCCL_PACKAGE_VERSION 2.8.4-1
ENV NV_LIBNCCL_VERSION 2.8.4
ENV NCCL_VERSION 2.8.4
ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}-${NV_LIBNCCL_PACKAGE_VERSION}+cuda11.1
FROM base as base-arm64

ENV NV_NVTX_VERSION 11.1.74-1
ENV NV_LIBNPP_VERSION 11.1.2.301-1
ENV NV_LIBCUBLAS_VERSION 11.3.0.106-1
ENV NV_LIBNCCL_PACKAGE_NAME libnccl
ENV NV_LIBNCCL_PACKAGE_VERSION 2.8.4-1
ENV NV_LIBNCCL_VERSION 2.8.4
ENV NCCL_VERSION 2.8.4
ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}-${NV_LIBNCCL_PACKAGE_VERSION}+cuda11.1
FROM base as base-ppc64le

ENV NV_NVTX_VERSION 11.1.74-1
ENV NV_LIBNPP_VERSION 11.1.2.301-1
ENV NV_LIBCUBLAS_VERSION 11.3.0.106-1
ENV NV_LIBNCCL_PACKAGE_NAME libnccl
ENV NV_LIBNCCL_PACKAGE_VERSION 2.8.4-1
ENV NV_LIBNCCL_VERSION 2.8.4
ENV NCCL_VERSION 2.8.4
ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}-${NV_LIBNCCL_PACKAGE_VERSION}+cuda11.1

FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <<EMAIL>>"

RUN yum install -y \
    cuda-libraries-11-1-${NV_CUDA_LIB_VERSION} \
    cuda-nvtx-11-1-${NV_NVTX_VERSION} \
    libnpp-11-1-${NV_LIBNPP_VERSION} \
    libcublas-11-1-${NV_LIBCUBLAS_VERSION} \
    ${NV_LIBNCCL_PACKAGE} \
    && yum clean all \
    && rm -rf /var/cache/yum/*