FROM rocm/pytorch:rocm6.3.2_ubuntu22.04_py3.10_pytorch_release_2.4.0

WORKDIR /root

RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y --no-install-recommends \
  build-essential git wget \
  libgtest-dev libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev llvm-dev \
  rocm-dev rocm-libs hip-dev hipblas-dev rocblas-dev \
  && apt-get clean autoclean && rm -rf /var/lib/apt/lists/{cache,log} /tmp/* /var/tmp/*

ENV PATH="/opt/conda/bin:${PATH}"
ENV LIBGL_ALWAYS_INDIRECT=1
ENV USE_ROCM=1
ENV USE_CUDA=0
ENV ROCM_HOME=/opt/rocm
ENV HIP_PLATFORM=amd
ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942"


RUN conda run -n py_3.10 conda install pip cmake -y && \
    conda run -n py_3.10 conda install -c conda-forge libstdcxx-ng=12 -y && \
    conda clean --all

RUN apt-get update && apt-get install -y python3 python3-dev python3-setuptools gcc libtinfo-dev zlib1g-dev build-essential cmake libedit-dev libxml2-dev && \
    apt-get clean autoclean && rm -rf /var/lib/apt/lists/{cache,log} /tmp/* /var/tmp/*

# Copy local tilelang directory instead of cloning from git
# Build from tilelang root: docker build -f docker/Dockerfile.rocm -t mi300:latest .
COPY . /root/tilelang

RUN mv /opt/conda/envs/py_3.10/compiler_compat /opt/conda/envs/py_3.10/compiler_compat.bak || true && \
    conda run -n py_3.10 bash -c "export USE_ROCM=1 USE_CUDA=0 && pip install 'numpy<2.0' --force-reinstall" && \
    conda run -n py_3.10 bash -c "cd /root/tilelang && \
        # Backup and modify pyproject.toml to remove torch from dependencies \
        cp pyproject.toml pyproject.toml.bak && \
        sed -i '/^[[:space:]]*\"torch/d' pyproject.toml && \
        # Install tilelang with all dependencies except torch \
        USE_ROCM=1 USE_CUDA=0 pip install -e . -v && \
        # Restore original pyproject.toml \
        mv pyproject.toml.bak pyproject.toml"

RUN conda init bash && \
    echo "conda activate py_3.10" >> /root/.bashrc

SHELL ["/bin/bash", "-l", "-c"]

ENTRYPOINT ["/bin/bash", "--login", "-i"]