triton-therock.Dockerfile

# Use manylinux that ships with many Python versions for wheels. manylinux_2_28
# is AlmaLinux 8 based and is binary-compatible with Red Hat Enterprise Linux.
FROM quay.io/pypa/manylinux_2_28_x86_64@sha256:9042a22d33af2223ff7a3599f236aff1e4ffd07e1ed1ac93a58877638317515f

SHELL ["/bin/bash", "-o", "pipefail", "-c"]

######## Setup Python #######
# Choe our default Python version
ENV PATH="/opt/python/cp311-cp311/bin:${PATH}"

######## Install Build Packages #######
# First enable Extra Packages for Enterprise Linux (EPEL) support, then install
# common build packages, and finally dependency packages for building HIP.
RUN yum install -y epel-release && \
    yum install -y clang lld ninja-build && \
    yum install -y numactl-devel elfutils-libelf-devel vim-common git-lfs && \
    yum clean all && \
    rm -rf /var/cache/yum

######## Install HIP/TheRock Dependencies #######
RUN curl https://storage.googleapis.com/git-repo-downloads/repo > /usr/local/bin/repo && \
    chmod a+x /usr/local/bin/repo && \
    pip install CppHeaderParser

######## Build TheRock #######
WORKDIR /build-therock
ENV THEROCK_PY_VERSION=0.1

RUN git clone https://github.com/nod-ai/TheRock && cd TheRock && \
    build_tools/fetch_sources.py && rm -Rf sources/.repo && \
    cd python_projects/runtime && pip wheel -v . && \
    pip install ./TheRock_runtime-${THEROCK_PY_VERSION}-py3-none-linux_x86_64.whl && \
    cd / && rm -Rf /build-therock

######## Build Triton #######
WORKDIR /code

RUN git clone https://github.com/antiagainst/triton.git -b therock && \
    curl -L https://github.com/llvm/llvm-project/archive/$(cat triton/cmake/llvm-hash.txt).zip -o llvm.zip && \
    unzip -q llvm.zip && mv llvm-project-$(cat triton/cmake/llvm-hash.txt) llvm

RUN cd /code/llvm && \
    cmake -GNinja -S llvm/ -B build-mlir/ -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_C_COMPILER=$(which clang) -DCMAKE_CXX_COMPILER=$(which clang++) \
    -DLLVM_ENABLE_PROJECTS="llvm;mlir" -DLLVM_TARGETS_TO_BUILD="AMDGPU;NVPTX;X86" && \
    cmake --build build-mlir/

RUN pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm6.0 && \
    pip install numpy pytest

RUN cd /code/triton && \
    Python3_ROOT_DIR=/opt/python/cp311-cp311 \
    LLVM_INCLUDE_DIRS=/code/llvm/build-mlir/include \
    LLVM_LIBRARY_DIR=/code/llvm/build-mlir/lib \
    LLVM_SYSPATH=/code/llvm/build-mlir \
    TRITON_BUILD_WITH_CLANG_LLD=true pip install -v ./python