# Use the LTS Ubuntu repo
FROM ubuntu:26.04 AS reqs
# Set global build variables
ENV PIP_BREAK_SYSTEM_PACKAGES=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive

# My sad server with 96 Gbytes of ram keeps running out of ram ;(
# I would love to use the whole 80 core CPU ;(
ENV MAX_JOBS=5

# Install libraries
RUN apt update -y && \
    apt upgrade -y && \
    apt install -y --no-install-recommends \
    ca-certificates \
    wget \
    curl \
    git \
    gfortran \
    ninja-build \
    g++ \
    pkg-config \
    xxd \
    automake \
    libtool \
    libegl1-mesa-dev \
    texinfo \
    bison \
    flex \
    cmake \
    build-essential \
    libsimde-dev \
    zlib1g-dev \
    libncurses5-dev \
    libopenblas-dev \
    libpthreadpool-dev \
    libgoogle-glog-dev \
    libgtest-dev \
    libbenchmark-dev \
    libgdbm-dev \
    libnss3-dev \
    libssl-dev \
    libreadline-dev \
    libffi-dev \
    libsqlite3-dev \
    libelf-dev \
    libpciaccess-dev \
    protobuf-compiler \
    libprotobuf-dev \
    libbz2-dev \
    liblzma-dev \
    && rm -rf /var/lib/apt/lists/*
    # Tested Working

# Build python 3.11
FROM reqs AS python
WORKDIR /build/python
RUN wget https://www.python.org/ftp/python/3.11.15/Python-3.11.15.tgz && \
    tar -xf Python-3.11.15.tgz && \
    cd Python-3.11.15 && \
    ./configure --enable-optimizations --enable-shared LDFLAGS="-Wl,-rpath /usr/local/lib" && \
    make -j 80 && \
    make install && \
    ldconfig /usr/local/lib && \
    wget https://bootstrap.pypa.io/get-pip.py && \
    python3 get-pip.py && \
    ln -s /usr/local/bin/pip3 /bin/pip && \
    ln -s /usr/local/bin/pip3 /bin/pip3 && \
    ln -s /usr/local/bin/python3 /bin/python && \
    ln -s /usr/local/bin/python3 /bin/python3.11
    # Tested Working

# Build ROCm latest
FROM python AS rocm
WORKDIR /build/therock
COPY ./TheRock .
ENV CFLAGS="-Wno-error=discarded-qualifiers -fcommon -w"
ENV CXXFLAGS="-Wno-error=discarded-qualifiers -fcommon -w"
RUN env INSTALL_PREFIX=/usr/local ./dockerfiles/install_pinned_patchelf.sh && \
    pip install -r requirements.txt && \
    cmake -S . -B build \
    -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DTHEROCK_AMDGPU_TARGETS="gfx1100" \
    -DTHEROCK_DIST_AMDGPU_FAMILIES="gfx1100" \
    -DTHEROCK_AMDGPU_DIST_BUNDLE_NAME="gfx1100" \
    -DLLVM_HOST_TRIPLE="aarch64-unknown-linux-gnu" \
    -DLLVM_DEFAULT_TARGET_TRIPLE="aarch64-unknown-linux-gnu" \
    -DCMAKE_C_FLAGS="$CFLAGS" \
    -DCMAKE_CXX_FLAGS="$CXXFLAGS" \
    -DLLVM_ENABLE_RUNTIMES="compiler-rt;libcxx;libcxxabi" \
    -DLLVM_ENABLE_PROJECTS="clang;lld;flang" \
    -DCOMPILER_RT_BUILD_SANITIZERS=OFF \
    -DCMAKE_INSTALL_PREFIX=./install && \
    cmake --build build -- -j${MAX_JOBS}
    # Work in progress [ Breaking due to x86 intrinsics not ported to Aarch64 ]

# Build pytorch
FROM rocm AS pytorch
# Setup Envronment
ENV PYTORCH_ROCM_ARCH=1100
ENV USE_ROCM=1
ENV USE_CUDA=0
ENV USE_AOTRITON=0
ENV USE_SYSTEM_PTHREADPOOL=1
ENV USE_SYSTEM_PROTOC=1
ENV BUILD_CUSTOM_PROTOBUF=OFF
ENV USE_PYTORCH_QNNPACK=OFF
ENV USE_XNNPACK=OFF
ENV USE_ONNX=OFF
ENV USE_QNNPACK=OFF
ENV USE_NNPACK=OFF
ENV USE_KINETO=OFF
ENV USE_OPENMP=OFF
ENV USE_TENSORPIPE=OFF
ENV ROCM_PATH=/usr
ENV DEVICE_LIB_PATH=/usr/lib/llvm-21/lib/clang/21/amdgcn/bitcode
ENV HIPFLAGS="--rocm-device-lib-path=$DEVICE_LIB_PATH"
ENV CMAKE_ARGS="-DCMAKE_HIP_COMPILER_WORKS=1 \
-DCMAKE_CXX_COMPILER_WORKS=1 \
-DCMAKE_C_COMPILER_WORKS=1 \
-DCMAKE_HIP_COMPILER_ROCM_LIB_PATH=$DEVICE_LIB_PATH \
-DCMAKE_HIP_FLAGS=--rocm-device-lib-path=$DEVICE_LIB_PATH \
-DCMAKE_C_COMPILER=/usr/lib/llvm-21/bin/clang \
-DCMAKE_CXX_COMPILER=/usr/lib/llvm-21/bin/clang++ \
-DCMAKE_HIP_COMPILER=/usr/lib/llvm-21/bin/clang++ \
-DPython3_ROOT_DIR=/usr/local \
-DPYTHON_EXECUTABLE=/usr/local/bin/python3.11 \
-DPYTHON_LIBRARY=/usr/local/lib/libpython3.11.so \
-DPYTHON_INCLUDE_DIR=/usr/local/include/python3.11"
# Pull and build pytorch
WORKDIR /build/pytorch
RUN git clone https://github.com/pytorch/pytorch.git && \
    cd pytorch && \
    git submodule update --init --recursive && \
    make triton && \
    python tools/amd_build/build_amd.py && \
    pip install -r requirements.txt && \
    mkdir -p /opt/rocm/bin \
    ln -s /usr/bin/hipcc /opt/rocm/bin/hipcc \
    ln -s /usr/lib/llvm-21/bin/clang++ /usr/bin/clang++ \
    python3 setup.py build && \
    python3 setup.py install
    # Failing

# Build pytorch-audio
FROM pytorch AS torchaudio
WORKDIR /build/torchaudio

# Build pytorch-vision
FROM torchaudio AS torchvision
WORKDIR /build/torchaudio
RUN git clone https://github.com/pytorch/vision.git

# Build final stripped down image
FROM ubuntu:26.04
# Copy in rocm, pytorch, torchaudio, torchvision