From 2d86eef9490313e9d3ab2299448192775305db86 Mon Sep 17 00:00:00 2001
From: Matthew Francis-Landau <matthew.francis-landau@gm.com>
Date: Fri, 25 Jul 2025 12:09:07 -0700
Subject: [PATCH] [mlir-tensorrt docker] Add docker build file to make it
 easier to build project

Signed-off-by: Matthew Francis-Landau <matthew.francis-landau@gm.com>
---
 mlir-tensorrt/docker/Dockerfile    | 62 ++++++++++++++++++++++++++++++
 mlir-tensorrt/docker/run_docker.sh | 18 +++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 mlir-tensorrt/docker/Dockerfile
 create mode 100755 mlir-tensorrt/docker/run_docker.sh

diff --git a/mlir-tensorrt/docker/Dockerfile b/mlir-tensorrt/docker/Dockerfile
new file mode 100644
index 000000000..e7535fc11
--- /dev/null
+++ b/mlir-tensorrt/docker/Dockerfile
@@ -0,0 +1,62 @@
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# Also available under a BSD-style license. See LICENSE.
+
+ARG BASE_IMG=ubuntu:24.04
+FROM ${BASE_IMG} AS dev-base
+
+# Specify user IDs
+ARG GROUP
+ARG GID
+ARG USER
+ARG UID
+
+# Run below commands as root
+USER root
+
+# Install packages
+RUN apt-get update && \
+    apt-get install -y wget && \
+    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb -P /tmp/ && \
+    dpkg -i /tmp/cuda-keyring_1.1-1_all.deb && \
+    rm /tmp/cuda-keyring_1.1-1_all.deb
+
+RUN apt-get update && \
+    apt-get install -y \
+    lld \
+    clang \
+    clang-format \
+    gdb \
+    black \
+    python3-dev \
+    cmake \
+    sudo \
+    vim \
+    ninja-build \
+    patch \
+    pybind11-dev \
+    git \
+    libnccl-dev \
+    libopenmpi-dev \
+    openmpi-bin \
+    python3-pip \
+    python3-pynvml \
+    python3-numpy \
+    python3-psutil \
+    cuda-toolkit-12 && \
+    pip install --break-system-packages nanobind cupy-cuda12x
+
+
+# Set workdir before launching container
+WORKDIR /opt/src/mlir-tensorrt
+
+# Add user permissions
+RUN groupadd -o -g ${GID} ${GROUP} && \
+    useradd -u ${UID} -g ${GROUP} -ms /bin/bash ${USER} && \
+    usermod -aG sudo ${USER} && \
+    chown -R ${USER}:${GROUP} /opt/src/mlir-tensorrt && \
+    echo "${USER} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers
+
+# Switch to user
+USER ${USER}
diff --git a/mlir-tensorrt/docker/run_docker.sh b/mlir-tensorrt/docker/run_docker.sh
new file mode 100755
index 000000000..70f784351
--- /dev/null
+++ b/mlir-tensorrt/docker/run_docker.sh
@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
+set -ex
+
+docker build -f ${SCRIPT_DIR}/Dockerfile \
+             -t mlir-tensorrt:dev \
+             --build-arg GROUP=$(id -gn) \
+             --build-arg GID=$(id -g) \
+             --build-arg USER=$(id -un) \
+             --build-arg UID=$(id -u) \
+             .
+
+docker run -it \
+           -v "${SCRIPT_DIR}/../":"/opt/src/mlir-tensorrt" \
+           --gpus all \
+           mlir-tensorrt:dev