# 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 \
    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
RUN --mount=type=cache,target=/root/.cache pip install --break-system-packages torch

# Install PyTorch-MLIR
WORKDIR /opt
RUN git clone https://github.com/llvm/torch-mlir.git
WORKDIR /opt/torch-mlir
RUN python3 -m pip install -e . --user

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

# Verify installations
RUN python3 -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"]