FROM ubuntu:22.04

# Update repo information
RUN apt update

# Set up a non-root user
RUN apt install sudo -y
RUN useradd -m tensorflow-rocm -g sudo
RUN echo "%sudo ALL=(ALL:ALL) NOPASSWD: ALL" >> /etc/sudoers
USER tensorflow-rocm
WORKDIR /home/tensorflow-rocm

# Install Bazel 5.3.0
RUN sudo apt update
RUN sudo apt install apt-transport-https curl gnupg -y
RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor >bazel-archive-keyring.gpg
RUN sudo mv bazel-archive-keyring.gpg /usr/share/keyrings
RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/bazel-archive-keyring.gpg] https://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
RUN sudo apt update && sudo apt install bazel-5.3.0 -y
RUN sudo ln -s /usr/bin/bazel-5.3.0 /usr/bin/bazel

# Install ROCm 5.5.0
RUN sudo apt install wget
RUN wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/jammy/amdgpu-install_5.5.50500-1_all.deb
RUN sudo apt-get install ./amdgpu-install_5.5.50500-1_all.deb -y
RUN sudo amdgpu-install --usecase=rocm --no-dkms -y
RUN echo gfx1100 | sudo tee -a /opt/rocm/bin/target.lst

# Install Tensorflow build dependencies
RUN sudo apt install python3-dev python3-pip -y
RUN pip install -U --user pip numpy==1.22.4 wheel packaging requests opt_einsum
RUN pip install -U --user keras_preprocessing --no-deps
RUN sudo ln -s /usr/bin/python3 /usr/bin/python
RUN sudo apt install patchelf -y

# Fetch the tensorflow-rocm repo on the 2.11 branch
RUN sudo apt install git -y
RUN git clone --depth 1 -b r2.11-rocm-enhanced https://github.com/ROCmSoftwarePlatform/tensorflow-upstream.git
WORKDIR /home/tensorflow-rocm/tensorflow-upstream

# Very hacky way of implementing gfx1100 support in tensorflow-rocm by replacing gfx1030 support with gfx1100 support.
# Writing a .patch or creating a tensorflow-rocm fork with gfx1100 support to pull from would be cleaner :)
RUN grep -rl gfx1030 . | xargs sed -i "s/gfx1030/gfx1100/g"

# Build & Install tensorflow-rocm
RUN ./build_rocm_python3