Dockerfile.onnx-mlir

ARG BASE_IMAGE
FROM ${BASE_IMAGE} AS onnx-mlir-test

# Label the base image so it can be cleaned up after becoming dangling
ARG ONNX_MLIR_PR_NUMBER
ARG ONNX_MLIR_PR_NUMBER2
LABEL onnx_mlir_pr_number=${ONNX_MLIR_PR_NUMBER}
LABEL onnx_mlir_pr_number2=${ONNX_MLIR_PR_NUMBER2}

ARG WORK_DIR=/workdir
WORKDIR ${WORK_DIR}

# Copy onnx-mlir from the pull request
COPY . onnx-mlir

# Setup onnx
# Onnx commit 9fdb3bbd7b8c99b7b6e356069193ba7150405fce switched from
# hardcoded lib to CMAKE_INSTALL_LIBDIR. On RHEL ubi8 (and variants)
# it fails to find protobuf because CMAKE_INSTALL_LIBDIR resolves to
# lib64 but protobuf is installed under /usr/local/lib. So we force
# CMAKE_INSTALL_LIBDIR to be lib.
RUN ONNX_ROOT=${WORK_DIR}/onnx-mlir/third_party/onnx \
    && cd ${ONNX_ROOT} \
    && CMAKE_ARGS="-DCMAKE_INSTALL_LIBDIR=lib" python3 -m pip install .

ARG NPROC=4
ARG ACCEL=NNPA
ARG TEST_MCPU
ARG KEEPSRC

RUN LLVM_PROJECT_ROOT=${WORK_DIR}/llvm-project \
    && ONNX_MLIR_ROOT=${WORK_DIR}/onnx-mlir \
# Install llc, opt, and libraries
    && cp ${LLVM_PROJECT_ROOT}/build/bin/llc /usr/local/bin \
    && cp ${LLVM_PROJECT_ROOT}/build/bin/opt /usr/local/bin \
    && cp -a ${LLVM_PROJECT_ROOT}/build/lib/* /usr/local/lib \
# Build onnx-mlir, run tests, and install
    && cd ${ONNX_MLIR_ROOT} \
    && rm -rf build && mkdir -p build && cd build \
# NNPA acclerator is built on all archs to enable lit tests
# (dependent libzdnn is built on s390x only)
    && cmake -DMLIR_DIR=${LLVM_PROJECT_ROOT}/build/lib/cmake/mlir \
             -DCMAKE_BUILD_TYPE=Release \
             -DCMAKE_INSTALL_MESSAGE=NEVER \
             -DONNX_MLIR_ACCELERATORS=${ACCEL} .. \
    && make -j${NPROC} \
    && make -j${NPROC} LIT_OPTS=-v check-onnx-lit \
# User image is built with SIMD (currently on s390x only)
    && TEST_MCPU=${TEST_MCPU:-$([ "$(uname -m)" = "s390x" ] && echo z14 || \
                                [ "$(uname -m)" = "x86_64" ] &&  echo || \
                                [ "$(uname -m)" = "ppc64le" ] && echo || echo)} \
    && TEST_ARGS="-mcpu=${TEST_MCPU}" \
    && make NPROC=${NPROC} \
            CTEST_PARALLEL_LEVEL=${NPROC} \
            TEST_MCPU=${TEST_MCPU} \
            TEST_ARGS="${TEST_ARGS}" \
            -j${NPROC} \
            check-onnx-backend-numerical \
    && make check-docs \
    && make -j${NPROC} install \
    && echo -e "/usr/local/lib" > \
            /etc/ld.so.conf.d/onnx-mlir.conf && ldconfig \
# Clean up
    && cd ${WORK_DIR} \
    && distro=$(cat /etc/os-release|grep -Po '(?<=^ID=).*'|sed 's/\"//g') \
    && if [ "${distro}" = "debian" ] || [ "${distro}" = "ubuntu" ]; then \
          apt-get purge --autoremove -qq -y \
                  autoconf automake cmake curl git libtool \
                  make ninja-build python3 && \
          rm -rf /var/lib/apt/lists/*; \
       elif [ "${distro}" = "rhel" ] || [ "${distro}" = "fedora" ]; then \
          pip3 uninstall -q -y Cython pybind11 pytest pytest-forked \
               pytest-xdist typing-extensions && \
          yum remove -q -y \
              adwaita-icon-theme autoconf automake cmake diffutils file \
              git libtool make python39 && \
          rm -rf /var/cache/dnf/* /usr/local/bin/ninja; \
       fi \
    && rm -rf /tmp/* /usr/bin/python \
    && ([ "${KEEPSRC}" = "yes" ] || rm -rf llvm-project onnx-mlir)

# Copy from onnx-mlir-test to squash removed source layer and
# remove dependency on base image (onnx-mlir-llvm-shared)
FROM scratch

# Label the images for various checking and cleanup
ARG ONNX_MLIR_SHA1
ARG ONNX_MLIR_SHA1_DATE
ARG ONNX_MLIR_DOCKERFILE_SHA1
ARG ONNX_MLIR_PR_NUMBER
ARG ONNX_MLIR_PR_NUMBER2
LABEL onnx_mlir_sha1=${ONNX_MLIR_SHA1}
LABEL onnx_mlir_sha1_date=${ONNX_MLIR_SHA1_DATE}
LABEL onnx_mlir_dockerfile_sha1=${ONNX_MLIR_DOCKERFILE_SHA1}
LABEL onnx_mlir_pr_number=${ONNX_MLIR_PR_NUMBER}
LABEL onnx_mlir_pr_number2=${ONNX_MLIR_PR_NUMBER2}

COPY --from=onnx-mlir-test / /
ENTRYPOINT [ "onnx-mlir" ]
CMD [ "--help" ]
LABEL onnx_mlir_successfully_built=yes