Dockerfile - Add support for arm64 build

.github/workflows/build-image.yml

@@ -30,6 +30,12 @@ jobs:
           tags: superbench/main:cuda12.4
           runner: [self-hosted, rocm-build]
           build_args: "NUM_MAKE_JOBS=8"
+        # # TODO: Enable ARM build, check for hosted aarch64 runner availability
+        # - name: cuda12.4-arm
+        #   dockerfile: cuda12.4
+        #   tags: superbench/main:cuda12.4
+        #   runner: [self-hosted, aarch64-build]
+        #   build_args: "NUM_MAKE_JOBS=8"
         - name: cuda12.2
           dockerfile: cuda12.2
           tags: superbench/main:cuda12.2

dockerfile/cuda12.4-arm.dockerfile

@@ -0,0 +1,139 @@
+FROM nvcr.io/nvidia/pytorch:24.05-py3
+
+# OS:
+#   - Ubuntu: 22.04
+#   - OpenMPI: 4.1.5rc2
+#   - Docker Client: 20.10.8
+# NVIDIA:
+#   - CUDA: 12.4.1
+#   - CUDA Driver: 550.54.15
+#   - cuBLAS: 12.4.5.8
+#   - cuDNN: 9.1.0.70
+#   - NCCL: 2.21.5
+# Mellanox:
+#   - OFED: 24.04-0.6.6.0
+#   - HPC-X: v2.19
+# Intel:
+#   - mlc: v3.11a
+
+LABEL maintainer="SuperBench"
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    autoconf \
+    automake \
+    bc \
+    build-essential \
+    curl \
+    dmidecode \
+    ffmpeg \
+    git \
+    iproute2 \
+    jq \
+    libaio-dev \
+    libavcodec-dev \
+    libavformat-dev \
+    libavutil-dev \
+    libboost-program-options-dev \
+    libcap2 \
+    libcurl4-openssl-dev \
+    libnuma-dev \
+    libpci-dev \
+    libswresample-dev \
+    libtinfo5 \
+    libtool \
+    lshw \
+    python3-mpi4py \
+    net-tools \
+    nlohmann-json3-dev \
+    openssh-client \
+    openssh-server \
+    pciutils \
+    sudo \
+    util-linux \
+    vim \
+    wget \
+    && \
+    apt-get autoremove && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* /tmp/*
+
+ARG NUM_MAKE_JOBS=
+
+# Install Docker
+ENV DOCKER_VERSION=20.10.8
+RUN cd /tmp && \
+    wget -q https://download.docker.com/linux/static/stable/aarch64/docker-${DOCKER_VERSION}.tgz -O docker.tgz && \
+    tar --extract --file docker.tgz --strip-components 1 --directory /usr/local/bin/ && \
+    rm docker.tgz
+
+# Update system config
+RUN mkdir -p /root/.ssh && \
+    touch /root/.ssh/authorized_keys && \
+    mkdir -p /var/run/sshd && \
+    sed -i "s/[# ]*PermitRootLogin prohibit-password/PermitRootLogin yes/" /etc/ssh/sshd_config && \
+    sed -i "s/[# ]*PermitUserEnvironment no/PermitUserEnvironment yes/" /etc/ssh/sshd_config && \
+    sed -i "s/[# ]*Port.*/Port 22/" /etc/ssh/sshd_config && \
+    echo "* soft nofile 1048576\n* hard nofile 1048576" >> /etc/security/limits.conf && \
+    echo "root soft nofile 1048576\nroot hard nofile 1048576" >> /etc/security/limits.conf
+
+# Install OFED
+ENV OFED_VERSION=24.04-0.6.6.0
+RUN cd /tmp && \
+    wget -q https://content.mellanox.com/ofed/MLNX_OFED-${OFED_VERSION}/MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu22.04-aarch64.tgz && \
+    tar xzf MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu22.04-aarch64.tgz && \
+    MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu22.04-aarch64/mlnxofedinstall --user-space-only --without-fw-update --without-ucx-cuda --force --all && \
+    rm -rf /tmp/MLNX_OFED_LINUX-${OFED_VERSION}*
+
+# Install HPC-X
+ENV HPCX_VERSION=v2.19
+RUN cd /opt && \
+    rm -rf hpcx && \
+    wget -q https://content.mellanox.com/hpc/hpc-x/${HPCX_VERSION}/hpcx-${HPCX_VERSION}-gcc-mlnx_ofed-ubuntu22.04-cuda12-aarch64.tbz -O hpcx.tbz && \
+    tar xf hpcx.tbz && \
+    mv hpcx-${HPCX_VERSION}-gcc-mlnx_ofed-ubuntu22.04-cuda12-aarch64 hpcx && \
+    rm hpcx.tbz
+
+# Install Intel MLC
+RUN cd /tmp && \
+    wget -q https://downloadmirror.intel.com/822971/mlc_v3.11a.tgz -O mlc.tgz && \
+    tar xzf mlc.tgz Linux/mlc && \
+    cp ./Linux/mlc /usr/local/bin/ && \
+    rm -rf ./Linux mlc.tgz
+
+# Install NCCL 2.21.5
+RUN cd /tmp && \
+    git clone -b v2.21.5-1 https://github.com/NVIDIA/nccl.git && \
+    cd nccl && \
+    make -j src.build && \
+    make install && \
+    rm -rf /tmp/nccl
+
+ENV PATH="${PATH}" \
+    LD_LIBRARY_PATH="/usr/local/lib:/usr/local/mpi/lib:${LD_LIBRARY_PATH}" \
+    SB_HOME=/opt/superbench_dev \
+    SB_MICRO_PATH=/opt/superbench_dev \
+    ANSIBLE_DEPRECATION_WARNINGS=FALSE \
+    ANSIBLE_COLLECTIONS_PATH=/usr/share/ansible/collections
+
+RUN echo PATH="$PATH" > /etc/environment && \
+    echo LD_LIBRARY_PATH="$LD_LIBRARY_PATH" >> /etc/environment && \
+    echo SB_MICRO_PATH="$SB_MICRO_PATH" >> /etc/environment && \
+    echo "source /opt/hpcx/hpcx-init.sh && hpcx_load" >> /etc/bash.bashrc
+
+# Add config files
+ADD dockerfile/etc /opt/microsoft/
+
+WORKDIR ${SB_HOME}
+
+ADD third_party third_party
+RUN make -C third_party cuda_with_msccl
+
+ADD . .
+RUN python3 -m pip install --upgrade setuptools==65.7 && \
+    python3 -m pip install --no-cache-dir .[nvworker] && \
+    make cppbuild && \
+    make postinstall && \
+    rm -rf .git
+

setup.py

@@ -215,8 +215,8 @@ def run(self):
             ],
             'ort': [
                 'onnx>=1.10.2',
-                'onnxruntime-gpu==1.10.0; python_version<"3.10"',
-                'onnxruntime-gpu; python_version>="3.10"',
+                'onnxruntime-gpu==1.10.0; python_version<"3.10" and platform_machine != "aarch64"',
+                'onnxruntime-gpu; python_version>="3.10" and platform_machine != "aarch64"',
             ],
             'nvidia': ['py3nvml>=0.2.6'],
             'amd': ['amdsmi'],

superbench/benchmarks/micro_benchmarks/cpu_stream_performance.py

@@ -23,7 +23,7 @@ def __init__(self, name, parameters=''):
         super().__init__(name, parameters)
 
         self._bin_name = 'streamZen3.exe'
-        self.__cpu_arch = ['other', 'zen3', 'zen4']
+        self.__cpu_arch = ['other', 'zen3', 'zen4', 'neo2']
 
     def add_parser_arguments(self):
         """Add the specified arguments."""
@@ -80,6 +80,8 @@ def _preprocess(self):
             exe = 'streamZen3.exe'
         elif self._args.cpu_arch == 'zen4':
             exe = 'streamZen4.exe'
+        elif self._args.cpu_arch == 'neo2':
+            exe = 'streamNeo2.exe'
         else:
             exe = 'streamx86.exe'