Distributed tensors with transforms

CMakeLists.txt

@@ -74,6 +74,9 @@ option(MATX_EN_VISUALIZATION "Enable visualization support" OFF)
 #option(MATX_EN_CUTLASS OFF)
 option(MATX_EN_CUTENSOR OFF)
 option(MATX_EN_CUDSS OFF)
+option(MATX_EN_CUBLASMP "Enable cuBLASMp distributed dense linear algebra" OFF)
+option(MATX_EN_CUSOLVERMP "Enable cuSOLVERMp distributed dense solvers" OFF)
+option(MATX_EN_CUFFTMP "Enable cuFFTMp multi-process FFT dependencies" OFF)
 option(MATX_EN_FILEIO OFF)
 option(MATX_EN_NVTIFF OFF "Enable nvTiff support")
 option(MATX_EN_X86_FFTW OFF "Enable x86 FFTW support")
@@ -94,6 +97,11 @@ option(MATX_DISABLE_EXCEPTIONS "Disable C++ exceptions and log errors instead" O
 set(MATX_EN_PYBIND11 OFF CACHE BOOL "Enable pybind11 support")
 
 set(cudss_DIR "" CACHE PATH "Directory where cuDSS is installed.")
+set(cublasmp_DIR "" CACHE PATH "Directory where cuBLASMp is installed.")
+set(cusolvermp_DIR "" CACHE PATH "Directory where cuSOLVERMp is installed.")
+set(cufftmp_DIR "" CACHE PATH "Directory where cuFFTMp is installed.")
+set(nccl_DIR "" CACHE PATH "Directory where NCCL is installed.")
+set(nvshmem_DIR "" CACHE PATH "Directory where NVSHMEM is installed.")
 set(cutensor_DIR "" CACHE PATH "Directory where cuTENSOR is installed.")
 set(cutensornet_DIR "" CACHE PATH "Directory where cuTensorNet is installed.")
 set(eigen_DIR "" CACHE PATH "Directory where Eigen is installed")
@@ -416,6 +424,52 @@ if (MATX_EN_CUDSS)
     target_link_libraries(matx INTERFACE cuDSS::cuDSS)
 endif()
 
+if (MATX_EN_CUBLASMP OR MATX_EN_CUSOLVERMP)
+    include(cmake/FindNCCL.cmake)
+    if (NOT TARGET NCCL::NCCL)
+        message(FATAL_ERROR "NCCL is required when an MP backend is enabled")
+    endif()
+    target_link_libraries(matx INTERFACE NCCL::NCCL)
+endif()
+
+if (MATX_EN_CUBLASMP)
+    include(cmake/FindcuBLASMp.cmake)
+    if (NOT TARGET cuBLASMp::cuBLASMp)
+        message(FATAL_ERROR "MATX_EN_CUBLASMP requires cuBLASMp")
+    endif()
+    target_compile_definitions(matx INTERFACE MATX_EN_CUBLASMP)
+    target_link_libraries(matx INTERFACE cuBLASMp::cuBLASMp)
+endif()
+
+if (MATX_EN_CUSOLVERMP)
+    include(cmake/FindcuSOLVERMp.cmake)
+    if (NOT TARGET cuSOLVERMp::cuSOLVERMp)
+        message(FATAL_ERROR "MATX_EN_CUSOLVERMP requires cuSOLVERMp")
+    endif()
+    target_compile_definitions(matx INTERFACE MATX_EN_CUSOLVERMP)
+    target_link_libraries(matx INTERFACE cuSOLVERMp::cuSOLVERMp)
+endif()
+
+if (MATX_EN_CUFFTMP)
+    include(cmake/FindcuFFTMp.cmake)
+    if (NOT TARGET cuFFTMp::cuFFTMp)
+        message(FATAL_ERROR "MATX_EN_CUFFTMP requires cuFFTMp")
+    endif()
+
+    include(cmake/FindNVSHMEM.cmake)
+    if (NOT TARGET nvshmem::nvshmem_host OR
+        NOT TARGET nvshmem::nvshmem_device)
+        message(FATAL_ERROR
+            "MATX_EN_CUFFTMP requires NVSHMEM host and device libraries")
+    endif()
+
+    target_compile_definitions(matx INTERFACE MATX_EN_CUFFTMP)
+    target_link_libraries(matx INTERFACE
+        cuFFTMp::cuFFTMp
+        nvshmem::nvshmem_host
+        nvshmem::nvshmem_device)
+endif()
+
 # Find python3 and pybind11 for generating unit tests and benchmarks
 if (MATX_EN_FILEIO OR MATX_EN_VISUALIZATION OR MATX_EN_PYBIND11 OR MATX_BUILD_EXAMPLES OR MATX_BUILD_TESTS OR MATX_BUILD_BENCHMARKS)
     message(STATUS "Enabling pybind11 support")

cmake/FindNCCL.cmake

@@ -0,0 +1,24 @@
+# Find NCCL for the optional multi-process backends.
+
+find_path(NCCL_INCLUDE_DIR
+  NAMES nccl.h
+  HINTS ${nccl_DIR} ENV NCCL_HOME
+  PATH_SUFFIXES include)
+
+find_library(NCCL_LIBRARY
+  NAMES nccl
+  HINTS ${nccl_DIR} ENV NCCL_HOME
+  PATH_SUFFIXES lib lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NCCL
+  REQUIRED_VARS NCCL_INCLUDE_DIR NCCL_LIBRARY)
+
+if(NCCL_FOUND AND NOT TARGET NCCL::NCCL)
+  add_library(NCCL::NCCL UNKNOWN IMPORTED)
+  set_target_properties(NCCL::NCCL PROPERTIES
+    IMPORTED_LOCATION "${NCCL_LIBRARY}"
+    INTERFACE_INCLUDE_DIRECTORIES "${NCCL_INCLUDE_DIR}")
+endif()
+
+mark_as_advanced(NCCL_INCLUDE_DIR NCCL_LIBRARY)

cmake/FindNVSHMEM.cmake

@@ -0,0 +1,51 @@
+# Find NVSHMEM host and device libraries for cuFFTMp.
+
+find_package(NVSHMEM CONFIG QUIET
+  HINTS ${nvshmem_DIR} ENV NVSHMEM_PREFIX ENV NVSHMEM_HOME)
+if(TARGET nvshmem::nvshmem_host AND TARGET nvshmem::nvshmem_device)
+  set(NVSHMEM_FOUND TRUE)
+  return()
+endif()
+
+find_path(NVSHMEM_INCLUDE_DIR
+  NAMES nvshmem.h
+  HINTS ${nvshmem_DIR} ENV NVSHMEM_PREFIX ENV NVSHMEM_HOME ENV NVSHMEM_INC
+  PATH_SUFFIXES include)
+
+find_library(NVSHMEM_HOST_LIBRARY
+  NAMES nvshmem_host
+  HINTS ${nvshmem_DIR} ENV NVSHMEM_PREFIX ENV NVSHMEM_HOME ENV NVSHMEM_LIB
+  PATH_SUFFIXES lib lib64)
+
+find_library(NVSHMEM_DEVICE_LIBRARY
+  NAMES nvshmem_device
+  HINTS ${nvshmem_DIR} ENV NVSHMEM_PREFIX ENV NVSHMEM_HOME ENV NVSHMEM_LIB
+  PATH_SUFFIXES lib lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NVSHMEM
+  REQUIRED_VARS
+    NVSHMEM_INCLUDE_DIR
+    NVSHMEM_HOST_LIBRARY
+    NVSHMEM_DEVICE_LIBRARY)
+
+if(NVSHMEM_FOUND)
+  if(NOT TARGET nvshmem::nvshmem_host)
+    add_library(nvshmem::nvshmem_host UNKNOWN IMPORTED)
+    set_target_properties(nvshmem::nvshmem_host PROPERTIES
+      IMPORTED_LOCATION "${NVSHMEM_HOST_LIBRARY}"
+      INTERFACE_INCLUDE_DIRECTORIES "${NVSHMEM_INCLUDE_DIR}")
+  endif()
+
+  if(NOT TARGET nvshmem::nvshmem_device)
+    add_library(nvshmem::nvshmem_device UNKNOWN IMPORTED)
+    set_target_properties(nvshmem::nvshmem_device PROPERTIES
+      IMPORTED_LOCATION "${NVSHMEM_DEVICE_LIBRARY}"
+      INTERFACE_INCLUDE_DIRECTORIES "${NVSHMEM_INCLUDE_DIR}")
+  endif()
+endif()
+
+mark_as_advanced(
+  NVSHMEM_INCLUDE_DIR
+  NVSHMEM_HOST_LIBRARY
+  NVSHMEM_DEVICE_LIBRARY)

cmake/FindcuBLASMp.cmake

@@ -0,0 +1,24 @@
+# Find the separately distributed cuBLASMp package.
+
+find_path(cuBLASMp_INCLUDE_DIR
+  NAMES cublasMp.h cublasmp.h
+  HINTS ${cublasmp_DIR} ENV CUBLASMP_HOME
+  PATH_SUFFIXES include)
+
+find_library(cuBLASMp_LIBRARY
+  NAMES cublasmp
+  HINTS ${cublasmp_DIR} ENV CUBLASMP_HOME
+  PATH_SUFFIXES lib lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(cuBLASMp
+  REQUIRED_VARS cuBLASMp_INCLUDE_DIR cuBLASMp_LIBRARY)
+
+if(cuBLASMp_FOUND AND NOT TARGET cuBLASMp::cuBLASMp)
+  add_library(cuBLASMp::cuBLASMp UNKNOWN IMPORTED)
+  set_target_properties(cuBLASMp::cuBLASMp PROPERTIES
+    IMPORTED_LOCATION "${cuBLASMp_LIBRARY}"
+    INTERFACE_INCLUDE_DIRECTORIES "${cuBLASMp_INCLUDE_DIR}")
+endif()
+
+mark_as_advanced(cuBLASMp_INCLUDE_DIR cuBLASMp_LIBRARY)

cmake/FindcuFFTMp.cmake

@@ -0,0 +1,24 @@
+# Find the separately distributed cuFFTMp package.
+
+find_path(cuFFTMp_INCLUDE_DIR
+  NAMES cufftMp.h
+  HINTS ${cufftmp_DIR} ENV CUFFTMP_HOME ENV CUFFT_INC
+  PATH_SUFFIXES include include/cufftmp math_libs/include/cufftmp)
+
+find_library(cuFFTMp_LIBRARY
+  NAMES cufftMp
+  HINTS ${cufftmp_DIR} ENV CUFFTMP_HOME ENV CUFFT_LIB
+  PATH_SUFFIXES lib lib64 math_libs/lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(cuFFTMp
+  REQUIRED_VARS cuFFTMp_INCLUDE_DIR cuFFTMp_LIBRARY)
+
+if(cuFFTMp_FOUND AND NOT TARGET cuFFTMp::cuFFTMp)
+  add_library(cuFFTMp::cuFFTMp UNKNOWN IMPORTED)
+  set_target_properties(cuFFTMp::cuFFTMp PROPERTIES
+    IMPORTED_LOCATION "${cuFFTMp_LIBRARY}"
+    INTERFACE_INCLUDE_DIRECTORIES "${cuFFTMp_INCLUDE_DIR}")
+endif()
+
+mark_as_advanced(cuFFTMp_INCLUDE_DIR cuFFTMp_LIBRARY)

cmake/FindcuSOLVERMp.cmake

@@ -0,0 +1,24 @@
+# Find the separately distributed cuSOLVERMp package.
+
+find_path(cuSOLVERMp_INCLUDE_DIR
+  NAMES cusolverMp.h
+  HINTS ${cusolvermp_DIR} ENV CUSOLVERMP_HOME
+  PATH_SUFFIXES include)
+
+find_library(cuSOLVERMp_LIBRARY
+  NAMES cusolverMp cusolvermp
+  HINTS ${cusolvermp_DIR} ENV CUSOLVERMP_HOME
+  PATH_SUFFIXES lib lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(cuSOLVERMp
+  REQUIRED_VARS cuSOLVERMp_INCLUDE_DIR cuSOLVERMp_LIBRARY)
+
+if(cuSOLVERMp_FOUND AND NOT TARGET cuSOLVERMp::cuSOLVERMp)
+  add_library(cuSOLVERMp::cuSOLVERMp UNKNOWN IMPORTED)
+  set_target_properties(cuSOLVERMp::cuSOLVERMp PROPERTIES
+    IMPORTED_LOCATION "${cuSOLVERMp_LIBRARY}"
+    INTERFACE_INCLUDE_DIRECTORIES "${cuSOLVERMp_INCLUDE_DIR}")
+endif()
+
+mark_as_advanced(cuSOLVERMp_INCLUDE_DIR cuSOLVERMp_LIBRARY)

docs_input/build.rst

@@ -39,6 +39,31 @@ Optional Third-party Dependencies
 - `cutensor <https://developer.nvidia.com/cutensor>`_ 2.3.1.0+ (Required when using `einsum`)
 - `cutensornet <https://docs.nvidia.com/cuda/cuquantum/cutensornet>`_ 25.09.1.12+ (Required when using `einsum`)
 - `cuDSS <https://developer.nvidia.com/cudss>`_ 0.7.0.20+ (Required when using `solve` on sparse matrices)
+- `NCCL <https://developer.nvidia.com/nccl>`_ (Only required for multi-GPU or multi-node cuBLASMp and cuSOLVERMp support)
+- `cuBLASMp <https://docs.nvidia.com/cuda/cublasmp/>`_ (Only required for multi-GPU or multi-node block-cyclic ``matmul``)
+- `cuSOLVERMp <https://docs.nvidia.com/cuda/cusolvermp/>`_ (Only required for multi-GPU or multi-node block-cyclic ``chol``)
+- `cuFFTMp <https://docs.nvidia.com/cuda/cufftmp/>`_ (Only required for multi-process, multi-GPU, or multi-node FFT support)
+- `NVSHMEM <https://docs.nvidia.com/nvshmem/>`_ (Only required for multi-process, multi-GPU, or multi-node FFT support through cuFFTMp)
+
+Distributed NVIDIA MP backends are opt-in because cuBLASMp and cuSOLVERMp are
+needed only for multi-GPU or multi-node execution and are distributed
+separately from the CUDA Toolkit. Enable them with
+``-DMATX_EN_CUBLASMP=ON`` and/or ``-DMATX_EN_CUSOLVERMP=ON``. Package prefixes
+may be supplied through ``cublasmp_DIR``, ``cusolvermp_DIR``, and ``nccl_DIR``
+or the corresponding ``CUBLASMP_HOME``, ``CUSOLVERMP_HOME``, and ``NCCL_HOME``
+environment variables.
+
+The single-process cuFFT multi-GPU Xt/Mg path is part of the CUDA Toolkit and
+does not require another CMake option. Include ``<matx/distributed.h>`` to use
+the experimental distributed APIs.
+
+Multi-process cuFFTMp support has a separate dependency check because cuFFTMp
+and its compatible NVSHMEM build are distributed outside the CUDA Toolkit.
+Enable it with ``-DMATX_EN_CUFFTMP=ON``. Package prefixes may be supplied
+through ``cufftmp_DIR`` and ``nvshmem_DIR``, or through ``CUFFTMP_HOME`` and
+``NVSHMEM_PREFIX``/``NVSHMEM_HOME``. The discovery modules also recognize the
+HPC SDK sample variables ``CUFFT_INC``, ``CUFFT_LIB``, ``NVSHMEM_INC``, and
+``NVSHMEM_LIB``.
 
 Host (CPU) Support
 ------------------
@@ -190,6 +215,12 @@ Unless otherwise noted, these options are OFF by default.
     - ``-DMATX_EN_CUTENSOR=ON``
   * - cuDSS Support
     - ``-DMATX_EN_CUDSS=ON``
+  * - cuBLASMp Support
+    - ``-DMATX_EN_CUBLASMP=ON``
+  * - cuSOLVERMp Support
+    - ``-DMATX_EN_CUSOLVERMP=ON``
+  * - cuFFTMp Dependency Support
+    - ``-DMATX_EN_CUFFTMP=ON``
   * - FFTW Support
     - ``-DMATX_EN_X86_FFTW=ON``
   * - NVPL Support

docs_input/developer_guide/distributed_tensors.rst

@@ -8,8 +8,10 @@ Status and goals
 
 ``experimental::distributed_tensor_t`` is a prototype for representing one
 logical tensor whose storage is split across CUDA devices and, eventually,
-processes. The prototype currently executes single-process, multi-GPU
-pointwise work, batch-local ``matmul``, ``chol``, and ``fft``, and gathers a
+processes. The prototype executes single-process, multi-GPU pointwise work,
+batch-local ``matmul``, ``chol``, and ``fft``, communicator-backed
+cuBLASMp and cuSOLVERMp execution selected by block-cyclic inputs, a cuFFT
+Xt/Mg transform selected when an FFT dimension spans local GPUs, and gathers a
 distributed tensor into a regular tensor. It is not yet a general distributed
 MatX operator system.
 
@@ -113,6 +115,80 @@ endpoints. The operation performs no communication, and each local call uses
 the same regular MatX accelerated path it would use for a non-distributed
 tensor.
 
+Collective MP linear algebra
+============================
+
+``block_cyclic_distribution_t`` describes the two-dimensional block-cyclic
+layout used by cuBLASMp and cuSOLVERMp. The endpoint vector is ordered by the
+selected process-grid layout and contains one endpoint per NCCL rank. The first
+collective configuration supports one local CUDA device per process:
+
+.. code-block:: cpp
+
+  #include <matx/distributed.h>
+
+  distributed_context context{{local_device}, mpi_rank, mpi_size};
+  block_cyclic_distribution_t layout{
+      {n, n}, {block_rows, block_columns}, {process_rows, process_columns},
+      endpoints};
+
+  // The application bootstraps and owns this NCCL communicator.
+  distributedCUDAExecutor exec{
+      context, nccl_communicator, process_rows, process_columns};
+  auto a = make_distributed_tensor<float>(layout, context);
+  auto b = make_distributed_tensor<float>(layout, context);
+  auto c = make_distributed_tensor<float>(layout, context);
+
+  (c = matmul(a, b)).run(exec);
+  (c = chol(a, SolverFillMode::LOWER)).run(exec);
+
+Both operations are collective: every rank in the process grid must enter them
+in the same order. The executor borrows the NCCL communicator, which must
+outlive the executor. MatX local views remain ordinary row-major tensors. The
+adapters pack them into the column-major local buffers required by the MP
+libraries and unpack the result, so the initial path favors correctness and
+interoperability over eliminating local layout conversions. The operations
+synchronize their local stream before returning because the libraries may
+retain host workspace during execution.
+
+Enable these paths with ``MATX_EN_CUBLASMP`` and ``MATX_EN_CUSOLVERMP``. They
+support rank-2 ``float``, ``double``, ``complex<float>``, and
+``complex<double>`` tensors. Batched MP operations, transpose modes, mixed
+precision, redistribution, and more solver factorizations remain future work.
+
+cuFFT multi-GPU
+===============
+
+The regular ``fft`` and ``ifft`` functions use the CUDA Toolkit's cuFFT Xt
+multi-GPU API when a single rank-1 complex transform is split across two or
+more GPUs in one process:
+
+.. code-block:: cpp
+
+  auto layout =
+      block_distribution_t<1>::Slab({fft_size}, {{0, 0}, {0, 1}});
+  auto input = make_distributed_tensor<complex<float>>(layout, context);
+  auto output = make_distributed_tensor<complex<float>>(layout, context);
+  (output = fft(input)).run(exec);
+
+The input distribution determines whether the trailing transform dimension is
+fully local. Fully local transforms keep using the batch-local path; a
+partitioned rank-1 transform selects Xt/Mg. The adapter stages through pinned
+host memory because the public cuFFT Xt copy API converts between a contiguous
+host array and its opaque multi-GPU descriptor. This also restores natural
+output order. It supports ``complex<float>`` and ``complex<double>`` and honors
+MatX ``FFTNorm`` modes. cuFFT itself determines which transform sizes and GPU
+counts its multi-GPU planner accepts.
+
+cuFFTMp is deliberately not treated as an interchangeable cuFFT Mg backend.
+It targets multi-process 2D/3D slab and pencil decompositions and requires
+NVSHMEM-compatible allocation, bootstrapping, and descriptor ownership.
+Ordinary ``make_distributed_tensor`` allocations do not satisfy that contract.
+Configure with ``MATX_EN_CUFFTMP`` to require and link compatible cuFFTMp and
+NVSHMEM installations. Transform execution still requires a cuFFTMp-owned
+tensor factory and reshape semantics rather than silently copying through a
+nominally distributed tensor.
+
 Materialization
 ===============
 
@@ -139,12 +215,13 @@ API rather than changing assignment semantics.
 Limitations and next steps
 ==========================
 
-* Only one process is executable today; no optional communication dependency is
-  introduced by the core type.
+* Pointwise execution and materialization remain single-process. Multi-process
+  execution is currently limited to block-cyclic ``matmul`` and ``chol``
+  collectives.
 * Pointwise operations require identical layouts. Batch-local transforms
   require aligned batch fragments and fully local operation dimensions.
   Redistribution and scattering will be explicit operations.
-* Other distributed BLAS, solver, reduction, DLPack, printing, and global
-  element-access paths are not provided yet.
+* Other distributed BLAS, solver, cuFFTMp, reduction, DLPack, printing, and
+  global element-access paths are not provided yet.
 * Materialization enqueues copies on the per-endpoint CUDA streams;
   ``distributedCUDAExecutor::sync`` is the completion boundary.