[wip] ADD ishmem support

CMakeLists.txt

@@ -56,6 +56,16 @@ if(USE_XCCL)
     caffe2_update_option(USE_C10D_XCCL OFF)
     update_caffe2_macros_file()
   endif()
+  if(USE_ISHMEM)
+    include(${TORCH_XPU_OPS_ROOT}/cmake/ISHMEM.cmake)
+    if(NOT PYTORCH_FOUND_ISHMEM)
+      message(WARNING "ISHMEM not found, disabling ISHMEM support")
+      caffe2_update_option(USE_ISHMEM OFF)
+      update_caffe2_macros_file()
+    else()
+      message(STATUS "ISHMEM support enabled")
+    endif()
+  endif()
 endif()
 
 set(USE_SYCLTLA ON)

cmake/ISHMEM.cmake

@@ -0,0 +1,24 @@
+if(NOT __ISHMEM_INCLUDED)
+  set(__ISHMEM_INCLUDED TRUE)
+
+  # ISHMEM_ROOT, ISHMEM_LIBRARY_DIR, ISHMEM_INCLUDE_DIR are handled by FindISHMEM.cmake.
+  find_package(ISHMEM REQUIRED)
+  if(NOT ISHMEM_FOUND)
+    set(PYTORCH_FOUND_ISHMEM FALSE)
+    message(WARNING "${ISHMEM_NOT_FOUND_MESSAGE}")
+    return()
+  endif()
+
+  set(PYTORCH_FOUND_ISHMEM TRUE)
+  add_library(torch::ishmem INTERFACE IMPORTED)
+  set_property(
+    TARGET torch::ishmem PROPERTY INTERFACE_INCLUDE_DIRECTORIES
+    ${ISHMEM_INCLUDE_DIR})
+  set_property(
+    TARGET torch::ishmem PROPERTY INTERFACE_LINK_LIBRARIES
+    ${ISHMEM_LIBRARY})
+
+  message(STATUS "Found Intel SHMEM: ${ISHMEM_ROOT}")
+  message(STATUS "  ISHMEM include dir: ${ISHMEM_INCLUDE_DIR}")
+  message(STATUS "  ISHMEM library: ${ISHMEM_LIBRARY}")
+endif()

cmake/Modules/FindISHMEM.cmake

@@ -0,0 +1,65 @@
+# This will define the following variables:
+# ISHMEM_FOUND               : True if the system has the ISHMEM library.
+# ISHMEM_INCLUDE_DIR         : Include directories needed to use ISHMEM.
+# ISHMEM_LIBRARY_DIR         : The path to the ISHMEM library.
+# ISHMEM_LIBRARY             : ISHMEM library fullname.
+
+include(${CMAKE_ROOT}/Modules/FindPackageHandleStandardArgs.cmake)
+
+if(NOT CMAKE_SYSTEM_NAME MATCHES "Linux")
+  set(ISHMEM_FOUND False)
+  set(ISHMEM_NOT_FOUND_MESSAGE "Intel SHMEM library is only supported on Linux!")
+  return()
+endif()
+
+set(ISHMEM_ROOT $ENV{ISHMEM_ROOT})
+
+if(NOT ISHMEM_ROOT)
+  set(ISHMEM_FOUND False)
+  set(ISHMEM_NOT_FOUND_MESSAGE "ISHMEM_ROOT environment variable not set. Please set it to your ISHMEM installation directory.")
+  return()
+endif()
+
+# Find include path from binary.
+find_path(
+  ISHMEM_INCLUDE_DIR
+  NAMES ishmem.h
+  HINTS ${ISHMEM_ROOT}/include
+  NO_DEFAULT_PATH
+)
+
+# Find library directory from binary.
+find_path(
+  ISHMEM_LIBRARY_DIR
+  NAMES libishmem.a
+  HINTS ${ISHMEM_ROOT}/lib
+  NO_DEFAULT_PATH
+)
+
+# Find ISHMEM library fullname (static library).
+find_library(
+  ISHMEM_LIBRARY
+  NAMES ishmem
+  HINTS ${ISHMEM_LIBRARY_DIR}
+  NO_DEFAULT_PATH
+)
+
+if((NOT ISHMEM_INCLUDE_DIR) OR (NOT ISHMEM_LIBRARY_DIR) OR (NOT ISHMEM_LIBRARY))
+  set(ISHMEM_FOUND False)
+  set(ISHMEM_NOT_FOUND_MESSAGE "Intel SHMEM library not found! Please set ISHMEM_ROOT environment variable.")
+  return()
+endif()
+
+SET(CMAKE_INCLUDE_PATH ${CMAKE_INCLUDE_PATH}
+  "${ISHMEM_INCLUDE_DIR}")
+SET(CMAKE_LIBRARY_PATH ${CMAKE_LIBRARY_PATH}
+  "${ISHMEM_LIBRARY_DIR}")
+
+find_package_handle_standard_args(
+  ISHMEM
+  FOUND_VAR ISHMEM_FOUND
+  REQUIRED_VARS ISHMEM_INCLUDE_DIR ISHMEM_LIBRARY_DIR ISHMEM_LIBRARY
+  REASON_FAILURE_MESSAGE "${ISHMEM_NOT_FOUND_MESSAGE}"
+)
+
+mark_as_advanced(ISHMEM_INCLUDE_DIR ISHMEM_LIBRARY_DIR ISHMEM_LIBRARY)

src/BuildOnLinux.cmake

@@ -17,6 +17,9 @@ macro(setup_common_libraries)
     target_compile_definitions(torch_xpu_ops PRIVATE USE_C10D_XCCL)
     target_link_libraries(torch_xpu_ops PUBLIC torch::xccl)
     target_link_libraries(torch_xpu_ops PUBLIC fmt::fmt-header-only)
+    if(USE_ISHMEM AND PYTORCH_FOUND_ISHMEM)
+      target_link_libraries(torch_xpu_ops PUBLIC torch::ishmem)
+    endif()
   endif()
 
   if(USE_SYCLTLA)
@@ -50,6 +53,9 @@ else()
     target_compile_definitions(torch_xpu_ops PRIVATE USE_C10D_XCCL)
     target_link_libraries(torch_xpu_ops  PUBLIC torch::xccl)
     target_link_libraries(torch_xpu_ops PUBLIC fmt::fmt-header-only)
+    if(USE_ISHMEM AND PYTORCH_FOUND_ISHMEM)
+      target_link_libraries(torch_xpu_ops PUBLIC torch::ishmem)
+    endif()
   endif()
 
   if(USE_SYCLTLA)