Skip to content

Commit b205fa1

Browse files
danielhua23danielhua23
committed
add rocshmem support
1 parent 72bbe5f commit b205fa1

2 files changed

Lines changed: 68 additions & 2 deletions

File tree

‎docker/amd-docker.Dockerfile‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
FROM ghcr.io/actions/actions-runner:latest
22

33
ENV CXX=clang++
4+
ENV UCX_CXX=g++
5+
ENV UCX_CC=gcc
46

57
RUN sudo apt-get update -y \
68
&& sudo apt-get install -y software-properties-common \
@@ -40,7 +42,7 @@ RUN sudo apt update -y \
4042

4143
RUN sudo pip install --upgrade pip
4244

43-
RUN sudo pip install --no-cache-dir torch==2.10.0.dev20250916+rocm6.3 pytorch-triton-rocm --index-url https://download.pytorch.org/whl/nightly/rocm6.3
45+
RUN sudo pip install --no-cache-dir pytorch-triton-rocm torch --index-url https://download.pytorch.org/whl/nightly/rocm6.3
4446

4547
RUN git clone --recursive https://github.com/ROCm/aiter.git \
4648
&& cd aiter \
@@ -58,4 +60,60 @@ RUN sudo pip install \
5860
wheel \
5961
tinygrad
6062

61-
RUN sudo pip install git+https://github.com/ROCm/iris.git
63+
RUN sudo apt-get update -y \
64+
&& sudo apt-get install -y --no-install-recommends \
65+
autoconf \
66+
automake \
67+
libtool \
68+
pkg-config \
69+
build-essential \
70+
gfortran \
71+
flex \
72+
bison \
73+
libomp-dev \
74+
libhwloc-dev \
75+
libnuma-dev \
76+
&& sudo rm -rf /var/lib/apt/lists/*
77+
78+
ENV UCX_INSTALL_DIR=/opt/ucx
79+
ENV OMPI_INSTALL_DIR=/opt/openmpi
80+
ENV ROCSHMEM_INSTALL_DIR=/opt/rocshmem
81+
ENV ROCM_PATH=/opt/rocm
82+
83+
RUN cd /tmp \
84+
&& git clone https://github.com/openucx/ucx.git -b v1.17.x \
85+
&& cd ucx \
86+
&& ./autogen.sh \
87+
&& CC=gcc CXX=g++ ./configure --prefix=${UCX_INSTALL_DIR} --with-rocm=${ROCM_PATH} --enable-mt --disable-optimizations \
88+
&& make -j$(nproc) \
89+
&& sudo make install \
90+
&& cd / \
91+
&& sudo rm -rf /tmp/ucx
92+
93+
RUN cd /tmp \
94+
&& git clone --recursive https://github.com/open-mpi/ompi.git -b v5.0.x \
95+
&& cd ompi \
96+
&& ./autogen.pl \
97+
&& ./configure --prefix=${OMPI_INSTALL_DIR} --with-rocm=${ROCM_PATH} --with-ucx=${UCX_INSTALL_DIR} \
98+
&& make -j$(nproc) \
99+
&& sudo make install \
100+
&& cd / \
101+
&& sudo rm -rf /tmp/ompi
102+
103+
ENV PATH="${OMPI_INSTALL_DIR}/bin:${PATH}"
104+
ENV LD_LIBRARY_PATH="${OMPI_INSTALL_DIR}/lib:${UCX_INSTALL_DIR}/lib:/opt/rocm/lib"
105+
106+
107+
RUN cd /tmp \
108+
&& git clone https://github.com/ROCm/rocSHMEM.git \
109+
&& cd rocSHMEM \
110+
&& mkdir build \
111+
&& cd build \
112+
&& MPI_ROOT=${OMPI_INSTALL_DIR} UCX_ROOT=${UCX_INSTALL_DIR} CMAKE_PREFIX_PATH="${ROCM_PATH}:$CMAKE_PREFIX_PATH" \
113+
sudo ../scripts/build_configs/ipc_single -DCMAKE_INSTALL_PREFIX=/opt/rocshmem \
114+
&& cd / \
115+
&& sudo rm -rf /tmp/rocSHMEM
116+
117+
118+
ENV ROCSHMEM_INSTALL_DIR=${ROCSHMEM_INSTALL_DIR}
119+
ENV LD_LIBRARY_PATH="${ROCSHMEM_INSTALL_DIR}/lib:${LD_LIBRARY_PATH}"

‎scripts/rocshmem_test_payload.json‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"lang": "py",
3+
"sources": {
4+
"rocshmem_test.py": "import torch\nfrom torch.utils.cpp_extension import load_inline\nimport os\n\ndef test_rocshmem_compilation():\n \"\"\"Test ROCshmem compilation using PyTorch's load_inline\"\"\"\n \n print(\"=== ROCshmem PyTorch Inline Test ===\")\n \n # C++ source code for ROCshmem test\n cpp_source = \"\"\"\n #include <rocshmem.hpp>\n #include <iostream>\n #include <torch/extension.h>\n \n void test_rocshmem() {\n std::cout << \"Testing ROCshmem compilation...\" << std::endl;\n \n // Just test that we can compile and link with rocshmem\n // Don't actually initialize since we may not have proper MPI setup\n std::cout << \"ROCshmem headers included successfully!\" << std::endl;\n std::cout << \"Compilation test passed!\" << std::endl;\n }\n \n PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"test_rocshmem\", &test_rocshmem, \"Test ROCshmem compilation\");\n }\n \"\"\"\n \n # Set up include paths and libraries\n rocm_path = os.environ.get('ROCM_PATH', '/opt/rocm')\n rocshmem_path = os.environ.get('ROCSHMEM_INSTALL_DIR', '/home/runner/rocshmem')\n ompi_path = os.environ.get('OMPI_INSTALL_DIR', '/opt/openmpi')\n\n include_dirs = [\n f\"{rocm_path}/include\",\n f\"{rocshmem_path}/include/rocshmem\",\n f\"{ompi_path}/include\"\n ]\n\n library_dirs = [\n f\"{rocm_path}/lib\",\n f\"{rocshmem_path}/lib\",\n f\"{ompi_path}/lib\"\n ]\n\n libraries = [\n \"rocshmem\",\n \"mpi\", \n \"amdhip64\",\n \"hsa-runtime64\"\n ]\n\n ldflags = []\n for lib_dir in library_dirs:\n ldflags.append(f\"-L{lib_dir}\")\n\n for lib in libraries:\n ldflags.append(f\"-l{lib}\")\n\n extra_cflags = [f\"-I{include_dir}\" for include_dir in include_dirs]\n\n extra_ldflags = [\n \"--hip-link\"\n ] + ldflags\n \n try:\n # Use torch.utils.cpp_extension.load_inline to compile\n rocshmem_module = load_inline(\n name=\"rocshmem_test\",\n cpp_sources=cpp_source,\n extra_cflags=extra_cflags,\n extra_ldflags=extra_ldflags,\n verbose=True\n )\n \n print(\"Compilation successful!\")\n print(\"Linking successful!\")\n \n # Run the test\n rocshmem_module.test_rocshmem()\n \n print(\"ROCshmem test completed successfully!\")\n return True\n \n except Exception as e:\n print(f\"ROCshmem test failed: {e}\")\n return False\n\nif __name__ == \"__main__\":\n test_rocshmem_compilation()"
5+
},
6+
"main": "rocshmem_test.py",
7+
"mode": "test"
8+
}

0 commit comments

Comments
 (0)