Dockerfile

# Base image with CUDA for GPU support
#FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
FROM ubuntu:latest

# Install dependencies
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get update && apt-get install -y \
    git \
    build-essential \
    cmake \
    python-is-python3 \
    python3-pip \
    python3-pybind11 \
    python3-dev \
    wget

# Install llvm
RUN wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get -y update && apt-get install -y mlir-18-tools llvm-18

# Install pytorch-mlir

#RUN pip install --break-system-packages --pre torch-mlir torch \
#  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
#  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
RUN pip install --break-system-packages --upgrade pip
RUN GIT_SSL_NO_VERIFY=true pip install --break-system-packages git+https://github.com/llvm/torch-mlir.git

# Add PyTorch-MLIR tools to PATH
ENV PATH="/root/.local/bin:${PATH}"

# Verify installations
RUN python -c "import torch; print(torch.__version__)" && \
    mlir-opt --version && \
    llvm-as --version

# Working directory for MLIR operations
WORKDIR /workspace

# Command to keep the container running
CMD ["bash"]