From ff628b15e2ca705288490b04ca2721ecedfcad90 Mon Sep 17 00:00:00 2001
From: chivatam <satyanarayana1999.sn@gmail.com>
Date: Sun, 7 Sep 2025 10:26:19 -0400
Subject: [PATCH] rocshmem dependencies

---
 docker/amd-docker.Dockerfile | 55 ++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/docker/amd-docker.Dockerfile b/docker/amd-docker.Dockerfile
index 7d2d3ae5..c46ba18c 100644
--- a/docker/amd-docker.Dockerfile
+++ b/docker/amd-docker.Dockerfile
@@ -57,3 +57,58 @@ RUN sudo pip install \
     packaging \
     wheel \
     tinygrad
+
+RUN sudo apt-get update -y \
+    && sudo apt-get install -y --no-install-recommends \
+    autoconf \
+    automake \
+    libtool \
+    pkg-config \
+    build-essential \
+    gfortran \
+    flex \
+    bison \
+    && sudo rm -rf /var/lib/apt/lists/*
+
+ENV UCX_INSTALL_DIR=/opt/ucx
+ENV OMPI_INSTALL_DIR=/opt/openmpi
+ENV ROCSHMEM_INSTALL_DIR=/opt/rocshmem
+ENV ROCM_PATH=/opt/rocm
+
+RUN cd /tmp \
+    && git clone https://github.com/openucx/ucx.git -b v1.17.x \
+    && cd ucx \
+    && ./autogen.sh \
+    && ./configure --prefix=${UCX_INSTALL_DIR} --with-rocm=${ROCM_PATH} --enable-mt \
+    && make -j$(nproc) \
+    && sudo make install \
+    && cd / \
+    && sudo rm -rf /tmp/ucx
+
+RUN cd /tmp \
+    && git clone --recursive https://github.com/open-mpi/ompi.git -b v5.0.x \
+    && cd ompi \
+    && ./autogen.pl \
+    && ./configure --prefix=${OMPI_INSTALL_DIR} --with-rocm=${ROCM_PATH} --with-ucx=${UCX_INSTALL_DIR} \
+    && make -j$(nproc) \
+    && sudo make install \
+    && cd / \
+    && sudo rm -rf /tmp/ompi
+
+ENV PATH="${OMPI_INSTALL_DIR}/bin:${PATH}"
+ENV LD_LIBRARY_PATH="${OMPI_INSTALL_DIR}/lib:${UCX_INSTALL_DIR}/lib:${LD_LIBRARY_PATH}"
+
+
+RUN cd /tmp \
+    && git clone https://github.com/ROCm/rocSHMEM.git \
+    && cd rocSHMEM \
+    && mkdir build \
+    && cd build \
+    && MPI_ROOT=${OMPI_INSTALL_DIR} UCX_ROOT=${UCX_INSTALL_DIR} CMAKE_PREFIX_PATH="${ROCM_PATH}:$CMAKE_PREFIX_PATH" \
+       ../scripts/build_configs/ipc_single ${ROCSHMEM_INSTALL_DIR} \
+    && cd / \
+    && sudo rm -rf /tmp/rocSHMEM
+
+
+ENV ROCSHMEM_INSTALL_DIR=${ROCSHMEM_INSTALL_DIR}
+ENV LD_LIBRARY_PATH="${ROCSHMEM_INSTALL_DIR}/lib:${LD_LIBRARY_PATH}"