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 PYTHON_VERSION="3.11"
ARG HIP_ARCH="gfx803;gfx900"
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 typing-extensions wheel

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
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

# https://github.com/pytorch/pytorch/issues/125230
# https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4

WORKDIR /home/rocm/pytorch/dist
RUN cp $(ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/torch.whl

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