cleaning up torch and mori deps

pemeliya · pemeliya · commit cdc2c3071aff · 2026-02-18T15:28:42.000Z
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -22,13 +22,16 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 option(USE_ROCM "Whether to use rocm" ON)
 option(USE_BNXT "Whether to use BNXT NIC" OFF)
 option(USE_IONIC "Whether to use IONIC" OFF)
+option(FORCE_CODE_OBJECT_VERSION_5 "Force code object version to Triton" ON)
 option(BUILD_EXAMPLES "Whether to build examples" ON)
 option(BUILD_APPLICATION "Whether to build application library" ON)
 option(BUILD_SHMEM "Whether to build shmem library" ON)
 option(BUILD_OPS "Whether to build mori operation kernels" ON)
 option(BUILD_IO "Whether to build mori io library" ON)
 option(BUILD_PYBINDS "Whether to build mori python bindings" ON)
 option(BUILD_TESTS "Whether to build mori CPP tests" ON)
+option(ENABLE_TORCH "Whether to enable Torch bootstrap and python bindings" ON)
+option(ENABLE_MPI "Whether to enable MPI bootstrap" ON)
 option(ENABLE_PROFILER "Enable kernel profiling" OFF)
 option(ENABLE_DEBUG_PRINTF "Enable debug printf in device kernels" OFF)
 option(ENABLE_STANDARD_MOE_ADAPT "Enable standard moe adapt" OFF)
@@ -67,6 +70,16 @@ if(USE_IONIC)
   add_compile_definitions(ENABLE_IONIC)
 endif()
 
+message(STATUS "ENABLE_TORCH = ${ENABLE_TORCH}")
+if(ENABLE_TORCH)
+  add_compile_definitions(ENABLE_TORCH)
+endif()
+
+message(STATUS "ENABLE_MPI = ${ENABLE_MPI}")
+if(ENABLE_MPI)
+  add_compile_definitions(ENABLE_MPI)
+endif()
+
 message(STATUS "ENABLE_PROFILER = ${ENABLE_PROFILER}")
 if(ENABLE_PROFILER)
   add_compile_definitions(ENABLE_PROFILER)
@@ -88,6 +101,9 @@ endif()
 message(STATUS "WARP_ACCUM_UNROLL is set to: ${WARP_ACCUM_UNROLL}")
 add_definitions(-DWARP_ACCUM_UNROLL=${WARP_ACCUM_UNROLL})
 add_definitions(-DHIP_ENABLE_WARP_SYNC_BUILTINS)
+
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-deprecated-literal-operator")
+
 if(USE_ROCM)
   list(APPEND CMAKE_PREFIX_PATH "/opt/rocm")
   project(mori LANGUAGES HIP CXX C)
diff --git a/include/mori/application/bootstrap/bootstrap.hpp b/include/mori/application/bootstrap/bootstrap.hpp
@@ -22,6 +22,10 @@
 #pragma once
 
 #include "mori/application/bootstrap/base_bootstrap.hpp"
+#ifdef ENABLE_MPI
 #include "mori/application/bootstrap/mpi_bootstrap.hpp"
+#endif
+#ifdef ENABLE_TORCH
 #include "mori/application/bootstrap/torch_bootstrap.hpp"
+#endif
 #include "mori/application/bootstrap/socket_bootstrap.hpp"
diff --git a/include/mori/shmem/shmem_api.hpp b/include/mori/shmem/shmem_api.hpp
@@ -21,7 +21,9 @@
 // SOFTWARE.
 #pragma once
 
+#ifdef ENABLE_MPI
 #include <mpi.h>
+#endif
 
 #include <array>
 #include <cstddef>
@@ -51,10 +53,13 @@ constexpr unsigned int MORI_SHMEM_INIT_WITH_UNIQUEID = 1;
 
 // TODO: provide unified initialize / finalize APIs
 int ShmemInit(application::BootstrapNetwork* bootNet);
+#ifdef ENABLE_MPI
 int ShmemInit();  // Default initialization using MPI_COMM_WORLD
 int ShmemMpiInit(MPI_Comm);
+#endif // ENABLE_MPI
+#ifdef ENABLE_TORCH
 int ShmemTorchProcessGroupInit(const std::string& groupName);
-
+#endif // ENABLE_TORCH
 // UniqueId-based initialization APIs (nvshmem/rocshmem compatible)
 int ShmemGetUniqueId(mori_shmem_uniqueid_t* uid);
 int ShmemSetAttrUniqueIdArgs(int rank, int nranks, mori_shmem_uniqueid_t* uid,
diff --git a/include/mori/shmem/shmem_device_api.hpp b/include/mori/shmem/shmem_device_api.hpp
@@ -22,7 +22,6 @@
 #pragma once
 
 #include <assert.h>
-#include <mpi.h>
 
 #include "mori/application/application.hpp"
 #include "mori/core/core.hpp"
diff --git a/include/mori/shmem/shmem_ibgda_kernels.hpp b/include/mori/shmem/shmem_ibgda_kernels.hpp
@@ -22,7 +22,6 @@
 #pragma once
 
 #include <assert.h>
-#include <mpi.h>
 
 #include "mori/application/application.hpp"
 #include "mori/core/core.hpp"
diff --git a/include/mori/shmem/shmem_p2p_kernels.hpp b/include/mori/shmem/shmem_p2p_kernels.hpp
@@ -22,8 +22,6 @@
 #pragma once
 
 #include <assert.h>
-#include <mpi.h>
-
 #include <type_traits>
 
 #include "mori/application/application.hpp"
diff --git a/include/mori/shmem/shmem_sdma_kernels.hpp b/include/mori/shmem/shmem_sdma_kernels.hpp
@@ -22,8 +22,6 @@
 #pragma once
 
 #include <assert.h>
-#include <mpi.h>
-
 #include <type_traits>
 
 #include "mori/application/application.hpp"
diff --git a/src/application/CMakeLists.txt b/src/application/CMakeLists.txt
@@ -1,26 +1,28 @@
-find_package(MPI REQUIRED)
+if(ENABLE_MPI)
+  find_package(MPI REQUIRED)
+endif()
 find_package(hsa-runtime64 REQUIRED)
 find_package(hsakmt REQUIRED)
 #find_library(IONIC_LIBRARY 
 #  NAMES ionic 
 #  HINTS /lib/x86_64-linux-gnu
 #  REQUIRED
 #)
-execute_process(
-  COMMAND python -c "import torch; print(torch.utils.cmake_prefix_path)"
-  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-  OUTPUT_VARIABLE TORCH_DIR
-  OUTPUT_STRIP_TRAILING_WHITESPACE)
-cmake_path(SET TORCH_CMAKE_DIR NORMALIZE "${TORCH_DIR}/Torch")
-list(APPEND CMAKE_PREFIX_PATH ${TORCH_CMAKE_DIR})
-message(STATUS "Found LibTorch CMake Path: ${CMAKE_PREFIX_PATH}")
 
-find_package(Torch REQUIRED)
+if(ENABLE_TORCH)
+  execute_process(
+    COMMAND python -c "import torch; print(torch.utils.cmake_prefix_path)"
+    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+    OUTPUT_VARIABLE TORCH_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
+  cmake_path(SET TORCH_CMAKE_DIR NORMALIZE "${TORCH_DIR}/Torch")
+  list(APPEND CMAKE_PREFIX_PATH ${TORCH_CMAKE_DIR})
+  message(STATUS "Found LibTorch CMake Path: ${CMAKE_PREFIX_PATH}")
+  find_package(Torch REQUIRED)
+endif()
 
 add_library(
   mori_application SHARED
-  bootstrap/mpi_bootstrap.cpp
-  bootstrap/torch_bootstrap.cpp
   bootstrap/socket_bootstrap.cpp
   bootstrap/local_bootstrap.cpp
   transport/rdma/rdma.cpp
@@ -40,11 +42,18 @@ add_library(
   topology/pci.cpp
   topology/system.cpp)
 
+if(ENABLE_TORCH)
+  target_sources(mori_application PRIVATE bootstrap/torch_bootstrap.cpp)
+endif()
+
+if(ENABLE_MPI)
+  target_sources(mori_application PRIVATE bootstrap/mpi_bootstrap.cpp)
+endif()
+
 target_include_directories(mori_application PUBLIC ${CMAKE_SOURCE_DIR}/include)
 target_include_directories(mori_application PUBLIC ${CMAKE_SOURCE_DIR})
 target_link_libraries(
   mori_application
-  MPI::MPI_CXX
   ibverbs
   hip::host
   hip::device
@@ -61,8 +70,14 @@ if(USE_BNXT)
 endif()
 
 if(USE_IONIC)
-	target_link_libraries(mori_application ${IONIC_LIB})
+  target_link_libraries(mori_application ${IONIC_LIB})
+endif()
+
+if(ENABLE_MPI)
+  target_link_libraries(mori_application MPI::MPI_CXX)
 endif()
 
-target_include_directories(mori_application PUBLIC ${TORCH_INCLUDE_DIRS})
-target_link_libraries(mori_application c10 torch torch_cpu c10_hip torch_hip)
+if(ENABLE_TORCH)
+  target_include_directories(mori_application PUBLIC ${TORCH_INCLUDE_DIRS})
+  target_link_libraries(mori_application c10 torch torch_cpu c10_hip torch_hip)
+endif()
diff --git a/src/pybind/CMakeLists.txt b/src/pybind/CMakeLists.txt
@@ -1,4 +1,6 @@
-find_package(Torch REQUIRED)
+if (ENABLE_TORCH)
+  find_package(Torch REQUIRED)
+endif()
 
 # You can set PYTHON_EXECUTABLE via command line: -DPYTHON_EXECUTABLE=/path/to/python
 if(NOT DEFINED PYTHON_EXECUTABLE)
@@ -37,21 +39,23 @@ if(ENABLE_PROFILER)
 endif()
 
 target_include_directories(
-  mori_pybinds PUBLIC ${PYTHON_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS}
-                      ${CMAKE_BINARY_DIR}/generated/include)
-target_link_directories(mori_pybinds PUBLIC ${TORCH_INSTALL_PREFIX}/lib)
+  mori_pybinds PUBLIC ${PYTHON_INCLUDE_DIRS} ${CMAKE_BINARY_DIR}/generated/include)
 target_link_libraries(
   mori_pybinds
   mori_ops
   mori_io
-  ${TORCH_LIBRARIES}
-  torch_python
   hip::host
   hip::device)
 
-# For python packages to find dependent libraries
-set_target_properties(
-  mori_pybinds
-  PROPERTIES BUILD_RPATH "$ORIGIN;$ORIGIN/../torch/lib"
+if(ENABLE_TORCH)
+  target_include_directories(
+    mori_pybinds PUBLIC ${TORCH_INCLUDE_DIRS})
+  target_link_directories(mori_pybinds PUBLIC ${TORCH_INSTALL_PREFIX}/lib)
+  target_link_libraries(${TORCH_LIBRARIES} torch_python)
+  # For python packages to find dependent libraries
+  set_target_properties(
+    mori_pybinds
+    PROPERTIES BUILD_RPATH "$ORIGIN;$ORIGIN/../torch/lib"
              INSTALL_RPATH "$ORIGIN;$ORIGIN/../torch/lib"
              BUILD_WITH_INSTALL_RPATH TRUE)
+endif()
diff --git a/src/shmem/CMakeLists.txt b/src/shmem/CMakeLists.txt
@@ -1,4 +1,6 @@
-find_package(MPI REQUIRED)
+if(ENABLE_MPI)
+  find_package(MPI REQUIRED)
+endif()
 
 option(BUILD_SHMEM_DEVICE_WRAPPER "Build SHMEM device API wrapper" OFF)
 
@@ -8,21 +10,33 @@ else()
   add_library(mori_shmem init.cpp memory.cpp)
 endif()
 
-target_include_directories(mori_shmem PUBLIC ${CMAKE_SOURCE_DIR}/include)
-target_include_directories(mori_shmem PUBLIC ${CMAKE_SOURCE_DIR})
+target_include_directories(mori_shmem PUBLIC ${CMAKE_SOURCE_DIR} 
+                                             ${CMAKE_SOURCE_DIR}/include)
 target_link_libraries(
   mori_shmem
   mori_application
   mori_logging
-  MPI::MPI_CXX
   ibverbs
   hip::host
   hip::device
   mlx5)
 
-# Use code object version 5 for compatibility with Triton
-target_compile_options(mori_shmem PUBLIC "-fgpu-rdc" "-mcode-object-version=5")
-target_link_options(mori_shmem PUBLIC "-fgpu-rdc" "-mcode-object-version=5")
+if(ENABLE_MPI)
+  target_link_libraries(mori_shmem MPI::MPI_CXX)
+endif()
+
+target_compile_options(mori_shmem PUBLIC "-fgpu-rdc")
+target_link_options(mori_shmem PUBLIC "-fgpu-rdc")
+if(FORCE_CODE_OBJECT_VERSION_5) 
+  # Use code object version 5 for compatibility with Triton
+  target_compile_options(mori_shmem PUBLIC "-mcode-object-version=5")
+  target_link_options(mori_shmem PUBLIC "-mcode-object-version=5")
+endif()
+
+if(NOT ENABLE_TORCH)
+  # This seems to be required for compilation for JAX
+  set_property(TARGET mori_shmem PROPERTY POSITION_INDEPENDENT_CODE ON)
+endif()
 
 option(SAVE_TEMPS "Save intermediate compilation files" ON)
 if(SAVE_TEMPS)
diff --git a/src/shmem/init.cpp b/src/shmem/init.cpp
@@ -19,8 +19,9 @@
 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
+#ifdef ENABLE_MPI
 #include <mpi.h>
-
+#endif
 #include <cstdlib>
 #include <cstring>
 #include <iostream>
@@ -633,16 +634,17 @@ int ShmemFinalize() {
 /*                                      Other Initialization APIs                                */
 /* ---------------------------------------------------------------------------------------------- */
 
+#ifdef ENABLE_MPI
 int ShmemMpiInit(MPI_Comm mpiComm) {
   return ShmemInit(new application::MpiBootstrapNetwork(mpiComm));
 }
-
 int ShmemInit() { return ShmemMpiInit(MPI_COMM_WORLD); }
-
+#endif
+#ifdef ENABLE_TORCH
 int ShmemTorchProcessGroupInit(const std::string& groupName) {
   return ShmemInit(new application::TorchBootstrapNetwork(groupName));
 }
-
+#endif
 /* ---------------------------------------------------------------------------------------------- */
 /*                                      Query APIs                                               */
 /* ---------------------------------------------------------------------------------------------- */
@@ -795,7 +797,7 @@ int ShmemInitAttr(unsigned int flags, mori_shmem_init_attr_t* attr) {
     MORI_SHMEM_ERROR("Invalid arguments");
     return -1;
   }
-
+#ifdef ENABLE_MPI
   // MPI-based initialization
   if (flags == MORI_SHMEM_INIT_WITH_MPI_COMM) {
     if (attr->mpi_comm == nullptr) {
@@ -805,7 +807,7 @@ int ShmemInitAttr(unsigned int flags, mori_shmem_init_attr_t* attr) {
     int result = ShmemMpiInit(*reinterpret_cast<MPI_Comm*>(attr->mpi_comm));
     return (result == 0) ? 0 : -1;
   }
-
+#endif // ENABLE_MPI
   // UniqueId-based initialization
   if (flags == MORI_SHMEM_INIT_WITH_UNIQUEID) {
     // Validate rank parameters
diff --git a/src/shmem/memory.cpp b/src/shmem/memory.cpp
@@ -19,8 +19,6 @@
 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
-#include <mpi.h>
-
 #include "mori/application/memory/symmetric_memory.hpp"
 #include "mori/shmem/internal.hpp"
 #include "mori/shmem/shmem_api.hpp"
diff --git a/xbuild_mori.sh b/xbuild_mori.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+set -e
+
+SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
+
+# -DCMAKE_PREFIX_PATH=/usr/local/lib/python3.12/dist-packages/torch/share/cmake
+
+#pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm7.1
+
+full=${full:-0}
+
+mkdir -p build
+pushd build
+
+if [[ ${full} -eq 1 ]]; then
+
+#  apt-get install -y \
+#     git \
+#    ibverbs-utils \
+#     libpci-dev \
+#     libdw1 \
+#     cython3 
+
+# NOTE this would screw up hipcc installation!!!
+# better install MPI manually
+    #openmpi-bin \
+    #libopenmpi-dev \
+    #locales
+
+  rm -rf *
+  cmake -DUSE_ROCM=ON -DCMAKE_BUILD_TYPE=Release \
+      -DBUILD_EXAMPLES=OFF -DWARP_ACCUM_UNROLL=1 -DUSE_BNXT=OFF \
+      -DGPU_TARGETS=gfx942 \
+      -DFORCE_CODE_OBJECT_VERSION_5=OFF \
+      -DENABLE_MPI=OFF -DENABLE_TORCH=OFF \
+      .. 
+fi
+
+make VERBOSE=1 -j 2>&1 | tee ../yyybuild.log
+
+cp src/pybind/libmori_pybinds.so \
+   src/application/libmori_application.so \
+   src/io/libmori_io.so \
+   $SCRIPT_DIR/python/mori
+
+popd
+
+if [[ ${full} -eq 1 ]]; then
+  pip3 install -e .
+fi
