ARG ROCM_VERSION="5.7.3"

FROM wtongze/rocm-base:${ROCM_VERSION}

ARG ROCM_VERSION="5.7.3"
ARG PYTORCH_BRANCH="v2.2.2"
ARG TORCHVISION_BRANCH="v0.17.2"
ARG PYTHON_VERSION="3.11"
ARG NUMPY_VERSION="1.26.4"
# "gfx803" or "gfx803;gfx900"
ARG HIP_ARCH="gfx803"
ARG DEBIAN_FRONTEND="noninteractive"

# Add GPU target
RUN echo "${HIP_ARCH}" | sed 's/;/\n/g' | sudo tee /opt/rocm/bin/target.lst

WORKDIR /home/rocm
RUN git clone --depth 1 https://github.com/pytorch/pytorch.git -b ${PYTORCH_BRANCH} pytorch
WORKDIR /home/rocm/pytorch
RUN git submodule update --init --recursive

RUN <<EOF
sudo apt-get update
sudo apt install -y software-properties-common
sudo add-apt-repository -y ppa:deadsnakes/ppa
sudo apt-get update
sudo apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF

RUN python${PYTHON_VERSION} -m pip install wheel
RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
RUN python${PYTHON_VERSION} -m pip install numpy==${NUMPY_VERSION}

RUN <<EOF
sudo apt-get update
sudo apt-get install -y ninja-build
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF

ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
ENV PYTORCH_BUILD_NUMBER=0

# Torch
# > https://github.com/pytorch/pytorch/issues/125230
# > https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \
USE_ROCM=ON \
USE_NUMPY=1 \
USE_FLASH_ATTENTION=OFF \
USE_MEM_EFF_ATTENTION=OFF \
PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") python${PYTHON_VERSION} setup.py bdist_wheel

RUN <<EOF
sudo apt-get update
sudo apt-get install -y ffmpeg libjpeg-dev libpng-dev
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF

RUN python${PYTHON_VERSION} -m pip install /home/rocm/pytorch/dist/*.whl

# Torch Vision
WORKDIR /home/rocm/
RUN git clone --depth 1 https://github.com/pytorch/vision.git -b ${TORCHVISION_BRANCH} vision

WORKDIR /home/rocm/vision
RUN USE_ROCM=ON \
BUILD_VERSION=$(echo "${TORCHVISION_BRANCH}" | tr -d "v") \
python${PYTHON_VERSION} setup.py bdist_wheel

# Torch Audio
WORKDIR /home/rocm/
RUN git clone --depth 1 https://github.com/pytorch/audio.git -b ${PYTORCH_BRANCH} audio

WORKDIR /home/rocm/audio
RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
RUN USE_ROCM=ON \
BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
python${PYTHON_VERSION} setup.py bdist_wheel

RUN <<EOF
mkdir /home/rocm/build
cp /home/rocm/pytorch/dist/*.whl /home/rocm/build
cp /home/rocm/vision/dist/*.whl /home/rocm/build
cp /home/rocm/audio/dist/*.whl /home/rocm/build
EOF

RUN python${PYTHON_VERSION} -m pip install /home/rocm/build/*.whl

WORKDIR /home/rocm
CMD [ "bash", "-l" ]