diff --git a/CMakeLists.txt b/CMakeLists.txt index ef13839d8..ba8a148ad 100644 --- a/CMakeLists.txt +++ b/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") diff --git a/cmake/FindNCCL.cmake b/cmake/FindNCCL.cmake new file mode 100644 index 000000000..9e1088114 --- /dev/null +++ b/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) diff --git a/cmake/FindNVSHMEM.cmake b/cmake/FindNVSHMEM.cmake new file mode 100644 index 000000000..46c14ceb1 --- /dev/null +++ b/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) diff --git a/cmake/FindcuBLASMp.cmake b/cmake/FindcuBLASMp.cmake new file mode 100644 index 000000000..0a346b471 --- /dev/null +++ b/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) diff --git a/cmake/FindcuFFTMp.cmake b/cmake/FindcuFFTMp.cmake new file mode 100644 index 000000000..f84422df8 --- /dev/null +++ b/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) diff --git a/cmake/FindcuSOLVERMp.cmake b/cmake/FindcuSOLVERMp.cmake new file mode 100644 index 000000000..e20211f6a --- /dev/null +++ b/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) diff --git a/docs_input/build.rst b/docs_input/build.rst index 312c4246f..3d2ae0d87 100644 --- a/docs_input/build.rst +++ b/docs_input/build.rst @@ -39,6 +39,31 @@ Optional Third-party Dependencies - `cutensor `_ 2.3.1.0+ (Required when using `einsum`) - `cutensornet `_ 25.09.1.12+ (Required when using `einsum`) - `cuDSS `_ 0.7.0.20+ (Required when using `solve` on sparse matrices) +- `NCCL `_ (Only required for multi-GPU or multi-node cuBLASMp and cuSOLVERMp support) +- `cuBLASMp `_ (Only required for multi-GPU or multi-node block-cyclic ``matmul``) +- `cuSOLVERMp `_ (Only required for multi-GPU or multi-node block-cyclic ``chol``) +- `cuFFTMp `_ (Only required for multi-process, multi-GPU, or multi-node FFT support) +- `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 ```` 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 diff --git a/docs_input/developer_guide/distributed_tensors.rst b/docs_input/developer_guide/distributed_tensors.rst index c7598a0b2..4f1946db8 100644 --- a/docs_input/developer_guide/distributed_tensors.rst +++ b/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 + + 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(layout, context); + auto b = make_distributed_tensor(layout, context); + auto c = make_distributed_tensor(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``, and +``complex`` 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>(layout, context); + auto output = make_distributed_tensor>(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`` and ``complex`` 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. diff --git a/docs_input/examples/distributed_cholesky_mpi.rst b/docs_input/examples/distributed_cholesky_mpi.rst new file mode 100644 index 000000000..1737471cf --- /dev/null +++ b/docs_input/examples/distributed_cholesky_mpi.rst @@ -0,0 +1,58 @@ +Multi-Process Cholesky Decomposition +#################################### + +``distributed_cholesky_mpi`` demonstrates a collective Cholesky decomposition +across multiple processes and GPUs. Each MPI rank selects one node-local GPU, +owns one block-cyclic portion of the global matrix, and enters the MatX +``chol`` expression collectively. MatX dispatches the operation to cuSOLVERMp +using an NCCL communicator exchanged through MPI. + +Unlike the single-node multi-GPU FFT example, this program can span nodes. +MPI controls process placement and the total number of GPUs: + +* ``mpirun -n 2`` uses two ranks and two GPUs. +* ``mpirun -n 8`` uses eight ranks and eight GPUs. +* A multi-node host mapping can place those ranks across several machines. + +Prerequisites +============= + +cuSOLVERMp, NCCL, and an MPI implementation with C++ development headers are +required. The example target is available only when ``MATX_EN_CUSOLVERMP`` is +enabled. + +Build the example with:: + + cmake -S . -B build \ + -DMATX_BUILD_EXAMPLES=ON \ + -DMATX_EN_CUSOLVERMP=ON \ + -DCMAKE_BUILD_TYPE=Release + cmake --build build --target distributed_cholesky_mpi + +Use ``CUSOLVERMP_HOME`` and ``NCCL_HOME``, or the corresponding CMake package +directory variables, when those libraries are outside the default search path. + +Execution +========= + +Run four processes on four GPUs in one node:: + + mpirun -n 4 ./build/examples/distributed_cholesky_mpi + +Run eight processes split across two four-GPU nodes using the host syntax +provided by the local MPI implementation, for example:: + + mpirun -n 8 --host node0:4,node1:4 \ + ./build/examples/distributed_cholesky_mpi + +The optional arguments select the square matrix size and square block size:: + + mpirun -n 4 ./build/examples/distributed_cholesky_mpi 4096 256 + +When all node-local GPUs are visible, MPI local-rank discovery assigns ranks to +distinct device ordinals. Launchers may instead expose one distinct physical +GPU as device zero to each rank. + +The input is a deterministic positive-definite diagonal matrix. Every rank +verifies only its local Cholesky-factor fragment, and an MPI all-reduce reports +the maximum error without gathering the matrix onto one process. diff --git a/docs_input/examples/distributed_fft_multi_gpu.rst b/docs_input/examples/distributed_fft_multi_gpu.rst new file mode 100644 index 000000000..d8e4cd75f --- /dev/null +++ b/docs_input/examples/distributed_fft_multi_gpu.rst @@ -0,0 +1,27 @@ +Single-Node Multi-GPU FFT +######################### + +``distributed_fft_multi_gpu`` demonstrates one process using several GPUs on a +single node. It creates a block-distributed rank-1 tensor and runs a MatX +``fft`` expression through the CUDA Toolkit's cuFFT Xt/Mg backend. MPI, NCCL, +cuFFTMp, and NVSHMEM are not used by this example. + +Build the example with:: + + cmake -S . -B build \ + -DMATX_BUILD_EXAMPLES=ON \ + -DCMAKE_BUILD_TYPE=Release + cmake --build build --target distributed_fft_multi_gpu + +The first argument selects how many visible GPUs to use. The second optional +argument selects the global transform length:: + + ./build/examples/distributed_fft_multi_gpu 4 1048576 + +The GPU count defaults to two and must be at least two. The program uses device +ordinals starting at zero, so the requested number cannot exceed the GPUs +visible to the process. The transform length must be supported by cuFFT's +multi-GPU backend. + +The input is a unit impulse at global index zero. The example verifies that +every distributed output element is one and reports the maximum error. diff --git a/docs_input/examples/index.rst b/docs_input/examples/index.rst index 32a7e8f8b..5a7c76a47 100644 --- a/docs_input/examples/index.rst +++ b/docs_input/examples/index.rst @@ -11,3 +11,5 @@ the examples require compiling them into an executable. Instructions for buildin :maxdepth: 1 fftconv.rst + distributed_fft_multi_gpu.rst + distributed_cholesky_mpi.rst diff --git a/docs_input/executor_compatibility.rst b/docs_input/executor_compatibility.rst index 9c83f109d..9a425e98e 100644 --- a/docs_input/executor_compatibility.rst +++ b/docs_input/executor_compatibility.rst @@ -33,6 +33,15 @@ batch-sharded ``matmul``, ``chol``, and ``fft`` when their operation dimensions remain local. See :ref:`distributed-tensors` for its constraints. Other existing operators do not implicitly become distributed operations. +The regular transform APIs select collective backends from the distributed +input layout. Rank-2 block-cyclic ``matmul`` and ``chol`` expressions require a +communicator-configured ``distributedCUDAExecutor`` and use cuBLASMp and +cuSOLVERMp, respectively. A rank-1 complex ``fft`` or ``ifft`` whose transform +dimension spans multiple local GPUs uses cuFFT Xt/Mg with the same executor +type. The executor provides streams, communicator state, and backend handles; +it does not choose the mathematical operation. These paths are collective or +synchronizing boundaries and do not participate in JIT fusion. + .. csv-table:: Operator Executor Compatibility Matrix :header: "Operator", "HostExecutor", "CUDAExecutor", "CUDAJITExecutor", "distributedCUDAExecutor", "Notes" :widths: 22 12 10 12 18 46 @@ -80,7 +89,7 @@ existing operators do not implicitly become distributed operations. "cgsolve", "|no|", "|yes|", "|no|", "|no|", "CUDA iterative solver path." "channelize_poly", "|yes|", "|yes|", "|no|", "|no|", "Polyphase channelizer; host path directly computes the per-branch FIR and DFT stages." "chirp", "|yes|", "|yes|", "|yes|", "|no|", "Generator expression." - "chol", "|yes|", "|yes|", "|yes|", "|partial|", "Host support requires the CPU solver backend. CUDAJITExecutor support uses cuSolverDx through MathDx for supported rank 2-4 square float, double, complex-float, and complex-double matrices. Experimental distributedCUDAExecutor support is limited to aligned batch sharding with fully local matrix dimensions." + "chol", "|yes|", "|yes|", "|yes|", "|partial|", "Host support requires the CPU solver backend. CUDAJITExecutor support uses cuSolverDx through MathDx for supported rank 2-4 square float, double, complex-float, and complex-double matrices. Experimental distributedCUDAExecutor support covers aligned batch sharding and communicator-backed rank-2 block-cyclic cuSOLVERMp execution." "clone", "|yes|", "|yes|", "|yes|", "|no|", "View expression." "concat", "|yes|", "|yes|", "|yes|", "|no|", "View/expression composition." "conj", "|yes|", "|yes|", "|yes|", "|no|", "Element-wise expression." @@ -104,7 +113,7 @@ existing operators do not implicitly become distributed operations. "exp", "|yes|", "|yes|", "|yes|", "|no|", "Element-wise expression." "expj", "|yes|", "|yes|", "|yes|", "|no|", "Element-wise expression." "eye", "|yes|", "|yes|", "|yes|", "|no|", "Generator expression." - "fft", "|yes|", "|yes|", "|yes|", "|partial|", "Host support requires the CPU FFT backend. CUDAJITExecutor support uses cuFFTDx through MathDx for supported runtime shapes, precisions, and layouts. Experimental distributedCUDAExecutor support is limited to aligned batch sharding with a fully local transform dimension." + "fft", "|yes|", "|yes|", "|yes|", "|partial|", "Host support requires the CPU FFT backend. CUDAJITExecutor support uses cuFFTDx through MathDx for supported runtime shapes, precisions, and layouts. Experimental distributedCUDAExecutor support covers aligned batch sharding and rank-1 complex cuFFT Xt/Mg execution across local GPUs." "fft2", "|yes|", "|yes|", "|yes|", "|no|", "Host support requires the CPU FFT backend. CUDAJITExecutor support uses cuFFTDx through MathDx for supported 2D runtime shapes, precisions, and layouts." "fftfreq", "|yes|", "|yes|", "|yes|", "|no|", "Generator expression." "fftshift1D", "|yes|", "|yes|", "|yes|", "|no|", "View/reindex expression." @@ -151,7 +160,7 @@ existing operators do not implicitly become distributed operations. "make_channelize_poly_stream", "|yes|", "|yes|", "|no|", "|no|", "Streaming (segmented) polyphase channelizer object; feeds segments through the one-shot channelize_poly over a retained history. Host and CUDA executors; the streaming object is a stateful driver, not a JIT-fusable expression." "make_conv1d_stream", "|yes|", "|yes|", "|no|", "|no|", "Streaming (segmented) 1D convolution object built on the direct conv1d; host and CUDA executors. The direct method limits the filter to 1024 taps. The streaming object is a stateful driver, not a JIT-fusable expression." "make_resample_poly_stream", "|yes|", "|yes|", "|no|", "|no|", "Streaming (segmented) polyphase resampler object; feeds segments through the one-shot resample_poly over a retained history. Host and CUDA executors; the streaming object is a stateful driver, not a JIT-fusable expression." - "matmul", "|yes|", "|yes|", "|yes|", "|partial|", "Host support requires the CPU BLAS backend and supported floating or complex types. CUDAJITExecutor support uses cuBLASDx through MathDx for supported runtime shapes, precisions, layouts, and block-size intersections. Experimental distributedCUDAExecutor support is limited to aligned batch sharding with fully local matrix dimensions." + "matmul", "|yes|", "|yes|", "|yes|", "|partial|", "Host support requires the CPU BLAS backend and supported floating or complex types. CUDAJITExecutor support uses cuBLASDx through MathDx for supported runtime shapes, precisions, layouts, and block-size intersections. Experimental distributedCUDAExecutor support covers aligned batch sharding and communicator-backed rank-2 block-cyclic cuBLASMp execution." "matrix_norm", "|partial|", "|yes|", "|no|", "|no|", "Reduction transform; host execution is available but reductions are not generally parallelized across host threads." "matvec", "|yes|", "|yes|", "|yes|", "|no|", "Host support requires the CPU BLAS backend and supported floating or complex types. CUDAJITExecutor support follows cuBLASDx matmul constraints." "max", "|partial|", "|yes|", "|no|", "|no|", "Reduction transform; host execution is available but reductions are not generally parallelized across host threads. Element-wise maximum through binary operators remains JIT-compatible." diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 59e18fc9f..70fe40ed8 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -47,17 +47,27 @@ else() target_compile_options(example_lib INTERFACE ${MATX_CUDA_FLAGS}) endif() -if (MULTI_GPU) - set_target_properties(example_lib PROPERTIES CUDA_SEPARABLE_COMPILATION ON) -endif() - - foreach( example ${examples} ) string( CONCAT file ${example} ".cu" ) add_executable( ${example} ${file} ) target_link_libraries(${example} example_lib) endforeach() +add_executable(distributed_fft_multi_gpu distributed_fft_multi_gpu.cu) +target_link_libraries(distributed_fft_multi_gpu example_lib) + +# The collective Cholesky example is available with the opt-in cuSOLVERMp +# backend. +if(MATX_EN_CUSOLVERMP) + find_package(MPI QUIET COMPONENTS CXX) + if(MPI_CXX_FOUND) + add_executable(distributed_cholesky_mpi distributed_cholesky_mpi.cu) + target_link_libraries(distributed_cholesky_mpi example_lib MPI::MPI_CXX) + else() + message(STATUS "MPI C++ support not found; skipping distributed_cholesky_mpi") + endif() +endif() + # Examples in subdirectories add_executable(sarbp sarbp/sarbp.cu) target_link_libraries(sarbp example_lib) diff --git a/examples/distributed_cholesky_mpi.cu b/examples/distributed_cholesky_mpi.cu new file mode 100644 index 000000000..5524e6bee --- /dev/null +++ b/examples/distributed_cholesky_mpi.cu @@ -0,0 +1,311 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. +//////////////////////////////////////////////////////////////////////////////// + +#include "matx.h" +#include "matx/distributed.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using namespace matx; +using namespace matx::experimental; + +namespace { + +void CheckNccl(ncclResult_t status, const char *operation) { + if (status != ncclSuccess) { + throw std::runtime_error(std::string{operation} + ": " + + ncclGetErrorString(status)); + } +} + +void CheckMpi(int status, const char *operation) { + if (status == MPI_SUCCESS) { + return; + } + + char error[MPI_MAX_ERROR_STRING]{}; + int length = 0; + (void)MPI_Error_string(status, error, &length); + throw std::runtime_error(std::string{operation} + ": " + + std::string{error, static_cast(length)}); +} + +index_t ParsePositive(const char *text, const char *name) { + size_t parsed = 0; + const long long value = std::stoll(text, &parsed); + if (text[parsed] != '\0' || value <= 0) { + throw std::invalid_argument(std::string{name} + + " must be a positive integer"); + } + return static_cast(value); +} + +class NcclCommunicator { +public: + NcclCommunicator(int rank, int size) { + ncclUniqueId id{}; + if (rank == 0) { + CheckNccl(ncclGetUniqueId(&id), "ncclGetUniqueId"); + } + CheckMpi(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD), + "MPI_Bcast(NCCL ID)"); + CheckNccl(ncclCommInitRank(&communicator_, size, id, rank), + "ncclCommInitRank"); + } + + NcclCommunicator(const NcclCommunicator &) = delete; + NcclCommunicator &operator=(const NcclCommunicator &) = delete; + + ~NcclCommunicator() { + if (communicator_ != nullptr) { + (void)ncclCommDestroy(communicator_); + } + } + + ncclComm_t get() const noexcept { return communicator_; } + +private: + ncclComm_t communicator_ = nullptr; +}; + +distributed_index_t<2> ProcessGrid(int process_count) { + int rows = static_cast(std::sqrt(static_cast(process_count))); + while (process_count % rows != 0) { + --rows; + } + return {static_cast(rows), + static_cast(process_count / rows)}; +} + +float CholeskyDiagonal(index_t row) { + return static_cast(row % 17 + 2); +} + +template +void FillLocalMatrix(Tensor &tensor, Generator &&generator) { + if (tensor.LocalFragmentCount() != 1) { + throw std::runtime_error("Expected exactly one local matrix fragment"); + } + + const auto &fragment = tensor.LocalFragment(0); + const auto &distribution = tensor.DistributionDescriptor(); + const auto shape = distribution.LocalShape(fragment.distribution_index); + std::vector host(static_cast(shape[0] * shape[1])); + for (index_t row = 0; row < shape[0]; ++row) { + for (index_t column = 0; column < shape[1]; ++column) { + const auto global = distribution.LocalToGlobal( + fragment.distribution_index, {row, column}); + host[static_cast(row * shape[1] + column)] = + generator(global[0], global[1]); + } + } + MATX_CUDA_CHECK(cudaMemcpy(tensor.LocalView(0).Data(), host.data(), + host.size() * sizeof(float), + cudaMemcpyHostToDevice)); +} + +template +float LocalMaximumError(const Tensor &tensor) { + const auto &fragment = tensor.LocalFragment(0); + const auto &distribution = tensor.DistributionDescriptor(); + const auto shape = distribution.LocalShape(fragment.distribution_index); + std::vector host(static_cast(shape[0] * shape[1])); + MATX_CUDA_CHECK(cudaMemcpy(host.data(), tensor.LocalView(0).Data(), + host.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + + float maximum_error = 0.0F; + for (index_t row = 0; row < shape[0]; ++row) { + for (index_t column = 0; column < shape[1]; ++column) { + const auto global = distribution.LocalToGlobal( + fragment.distribution_index, {row, column}); + if (global[0] < global[1]) { + continue; + } + const float expected = + global[0] == global[1] ? CholeskyDiagonal(global[0]) : 0.0F; + const float actual = + host[static_cast(row * shape[1] + column)]; + maximum_error = + std::max(maximum_error, std::abs(actual - expected)); + } + } + return maximum_error; +} + +int RunExample(int argc, char **argv, int world_rank) { + int world_size = 0; + CheckMpi(MPI_Comm_size(MPI_COMM_WORLD, &world_size), "MPI_Comm_size"); + if (world_size < 2) { + if (world_rank == 0) { + std::cerr << "Launch at least two MPI ranks; each rank uses one GPU\n"; + } + return 2; + } + if (argc > 3) { + if (world_rank == 0) { + std::cerr << "Usage: distributed_cholesky_mpi [matrix_size " + "[block_size]]\n"; + } + return 2; + } + + const index_t matrix_size = + argc > 1 ? ParsePositive(argv[1], "matrix_size") : 1024; + const index_t block_size = + argc > 2 ? ParsePositive(argv[2], "block_size") : 128; + if (block_size > matrix_size) { + throw std::invalid_argument("block_size cannot exceed matrix_size"); + } + + MPI_Comm local_comm = MPI_COMM_NULL; + CheckMpi(MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, world_rank, + MPI_INFO_NULL, &local_comm), + "MPI_Comm_split_type"); + + int local_rank = 0; + int local_size = 0; + CheckMpi(MPI_Comm_rank(local_comm, &local_rank), "MPI_Comm_rank(local)"); + CheckMpi(MPI_Comm_size(local_comm, &local_size), "MPI_Comm_size(local)"); + + int visible_devices = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&visible_devices)); + // A launcher may expose every node-local GPU to every rank, or expose one + // distinct physical GPU as device zero to each rank. + const int local_device_error = + visible_devices == 1 || visible_devices >= local_size ? 0 : 1; + int any_device_error = 0; + CheckMpi(MPI_Allreduce(&local_device_error, &any_device_error, 1, MPI_INT, + MPI_MAX, MPI_COMM_WORLD), + "MPI_Allreduce(device availability)"); + if (any_device_error != 0) { + if (world_rank == 0) { + std::cerr << "Each node must expose at least one GPU per local MPI rank\n"; + } + CheckMpi(MPI_Comm_free(&local_comm), "MPI_Comm_free"); + return 2; + } + + const int device = visible_devices == 1 ? 0 : local_rank; + MATX_CUDA_CHECK(cudaSetDevice(device)); + + std::vector rank_devices(static_cast(world_size)); + CheckMpi(MPI_Allgather(&device, 1, MPI_INT, rank_devices.data(), 1, MPI_INT, + MPI_COMM_WORLD), + "MPI_Allgather(rank devices)"); + + std::vector endpoints; + endpoints.reserve(static_cast(world_size)); + for (int rank = 0; rank < world_size; ++rank) { + endpoints.push_back({rank, rank_devices[static_cast(rank)]}); + } + + const auto process_grid = ProcessGrid(world_size); + NcclCommunicator communicator{world_rank, world_size}; + distributed_context context{{device}, world_rank, world_size}; + distributedCUDAExecutor executor{ + context, communicator.get(), static_cast(process_grid[0]), + static_cast(process_grid[1])}; + + block_cyclic_distribution_t distribution{ + {matrix_size, matrix_size}, + {block_size, block_size}, + process_grid, + endpoints}; + auto input = make_distributed_tensor(distribution, context); + auto output = make_distributed_tensor(distribution, context); + + // A positive diagonal matrix makes the expected lower factor unambiguous. + FillLocalMatrix(input, [](index_t row, index_t column) { + if (row != column) { + return 0.0F; + } + const float diagonal = CholeskyDiagonal(row); + return diagonal * diagonal; + }); + + // Every MPI rank enters this expression collectively. The rank count, + // selected with mpirun -n, is also the number of GPUs used. + (output = chol(input, SolverFillMode::LOWER)).run(executor); + executor.sync(); + + const float local_error = LocalMaximumError(output); + float maximum_error = 0.0F; + CheckMpi(MPI_Allreduce(&local_error, &maximum_error, 1, MPI_FLOAT, MPI_MAX, + MPI_COMM_WORLD), + "MPI_Allreduce(maximum error)"); + + if (world_rank == 0) { + std::cout << "cuSOLVERMp Cholesky of a " << matrix_size << "x" + << matrix_size << " matrix across " << world_size + << " MPI ranks/GPUs in a " << process_grid[0] << "x" + << process_grid[1] << " process grid; maximum error " + << maximum_error << '\n'; + } + + CheckMpi(MPI_Comm_free(&local_comm), "MPI_Comm_free"); + return maximum_error <= 1.0e-4F ? 0 : 1; +} + +} // namespace + +int main(int argc, char **argv) { + if (MPI_Init(&argc, &argv) != MPI_SUCCESS) { + std::cerr << "MPI_Init failed\n"; + return 1; + } + + int world_rank = -1; + (void)MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + int result = 1; + try { + result = RunExample(argc, argv, world_rank); + } catch (const std::exception &error) { + std::fprintf(stderr, "Rank %d failed: %s\n", world_rank, error.what()); + (void)MPI_Abort(MPI_COMM_WORLD, 1); + } + + ClearCachesAndAllocations(); + (void)MPI_Finalize(); + return result; +} diff --git a/examples/distributed_fft_multi_gpu.cu b/examples/distributed_fft_multi_gpu.cu new file mode 100644 index 000000000..bb1cd24c9 --- /dev/null +++ b/examples/distributed_fft_multi_gpu.cu @@ -0,0 +1,163 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. +//////////////////////////////////////////////////////////////////////////////// + +#include "matx.h" +#include "matx/distributed.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace matx; +using namespace matx::experimental; + +namespace { + +index_t ParsePositive(const char *text, const char *name) { + size_t parsed = 0; + const long long value = std::stoll(text, &parsed); + if (text[parsed] != '\0' || value <= 0) { + throw std::invalid_argument(std::string{name} + + " must be a positive integer"); + } + return static_cast(value); +} + +int RunExample(int argc, char **argv) { + if (argc > 3) { + std::cerr << "Usage: distributed_fft_multi_gpu [gpu_count " + "[element_count]]\n"; + return 2; + } + + const index_t requested_gpu_count = + argc > 1 ? ParsePositive(argv[1], "gpu_count") : 2; + const index_t element_count = + argc > 2 ? ParsePositive(argv[2], "element_count") : 1 << 20; + if (requested_gpu_count > std::numeric_limits::max()) { + throw std::invalid_argument("gpu_count exceeds the supported range"); + } + const int gpu_count = static_cast(requested_gpu_count); + if (gpu_count < 2) { + throw std::invalid_argument( + "The cuFFT multi-GPU backend requires at least two GPUs"); + } + if (element_count < gpu_count) { + throw std::invalid_argument( + "element_count must be at least the requested GPU count"); + } + + int visible_gpu_count = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&visible_gpu_count)); + if (gpu_count > visible_gpu_count) { + throw std::invalid_argument("Requested " + std::to_string(gpu_count) + + " GPUs, but only " + + std::to_string(visible_gpu_count) + + " are visible"); + } + + std::vector devices; + std::vector endpoints; + devices.reserve(static_cast(gpu_count)); + endpoints.reserve(static_cast(gpu_count)); + for (int device = 0; device < gpu_count; ++device) { + devices.push_back(device); + endpoints.push_back({0, device}); + } + + using complex_type = cuda::std::complex; + distributed_context context{devices}; + distributedCUDAExecutor executor{context}; + auto distribution = + block_distribution_t<1>::Slab({element_count}, endpoints, 0); + auto input = make_distributed_tensor(distribution, context); + auto output = make_distributed_tensor(distribution, context); + + // An impulse at global index zero has a constant Fourier transform. + for (size_t local = 0; local < input.LocalFragmentCount(); ++local) { + const auto &fragment = input.LocalFragment(local); + const index_t local_count = input.LocalView(local).Size(0); + std::vector host(static_cast(local_count), + complex_type{0.0F, 0.0F}); + const auto first_global = + distribution.LocalToGlobal(fragment.distribution_index, {0}); + if (first_global[0] == 0) { + host[0] = complex_type{1.0F, 0.0F}; + } + matx::detail::distributed_device_guard guard{fragment.endpoint.device_id}; + MATX_CUDA_CHECK(cudaMemcpy(input.LocalView(local).Data(), host.data(), + host.size() * sizeof(complex_type), + cudaMemcpyHostToDevice)); + } + + // This is one process coordinating several GPUs on one node through cuFFT + // Xt/Mg. No MPI communicator is involved. + (output = fft(input)).run(executor); + executor.sync(); + + float maximum_error = 0.0F; + for (size_t local = 0; local < output.LocalFragmentCount(); ++local) { + const auto &fragment = output.LocalFragment(local); + const index_t local_count = output.LocalView(local).Size(0); + std::vector host(static_cast(local_count)); + matx::detail::distributed_device_guard guard{fragment.endpoint.device_id}; + MATX_CUDA_CHECK(cudaMemcpy(host.data(), output.LocalView(local).Data(), + host.size() * sizeof(complex_type), + cudaMemcpyDeviceToHost)); + for (const auto &value : host) { + maximum_error = + std::max(maximum_error, std::abs(value.real() - 1.0F)); + maximum_error = std::max(maximum_error, std::abs(value.imag())); + } + } + + std::cout << "Single-node cuFFT transform of " << element_count + << " elements across " << gpu_count + << " GPUs; maximum error " << maximum_error << '\n'; + return maximum_error <= 1.0e-5F ? 0 : 1; +} + +} // namespace + +int main(int argc, char **argv) { + try { + return RunExample(argc, argv); + } catch (const std::exception &error) { + std::cerr << "distributed_fft_multi_gpu failed: " << error.what() << '\n'; + return 1; + } +} diff --git a/include/matx.h b/include/matx.h index 59b888cd2..1b13f2cdf 100644 --- a/include/matx.h +++ b/include/matx.h @@ -62,7 +62,7 @@ #include "matx/operators/operators.h" #include "matx/transforms/transforms.h" #include "matx/streaming/streaming.h" -#include "matx/core/distributed_tensor.h" // distributed support is experimental +#include "matx/distributed.h" // distributed support is experimental #include namespace matx { diff --git a/include/matx/core/distributed_tensor.h b/include/matx/core/distributed_tensor.h index 47572df2f..c8ae8aa99 100644 --- a/include/matx/core/distributed_tensor.h +++ b/include/matx/core/distributed_tensor.h @@ -282,9 +282,12 @@ class block_cyclic_distribution_t { block_cyclic_distribution_t(distributed_index_t<2> global_shape, distributed_index_t<2> block_shape, distributed_index_t<2> process_grid, - std::vector endpoints) + std::vector endpoints, + distributed_grid_layout layout = + distributed_grid_layout::row_major) : global_shape_{global_shape}, block_shape_{block_shape}, - process_grid_{process_grid}, endpoints_{std::move(endpoints)} { + process_grid_{process_grid}, endpoints_{std::move(endpoints)}, + layout_{layout} { for (int dim = 0; dim < 2; ++dim) { matx::detail::DistributedCheck( global_shape_[dim] > 0 && block_shape_[dim] > 0 && @@ -308,16 +311,14 @@ class block_cyclic_distribution_t { const distributed_index_t<2> &ProcessGrid() const noexcept { return process_grid_; } + distributed_grid_layout GridLayout() const noexcept { return layout_; } size_t FragmentCount() const noexcept { return endpoints_.size(); } const distributed_endpoint_t &FragmentEndpoint(size_t fragment) const { return endpoints_.at(fragment); } distributed_index_t<2> LocalShape(size_t fragment) const { - const index_t process_row = - static_cast(fragment) / process_grid_[1]; - const index_t process_col = - static_cast(fragment) % process_grid_[1]; + const auto [process_row, process_col] = ProcessCoordinate(fragment); return {OwnedExtent(global_shape_[0], block_shape_[0], process_grid_[0], process_row), OwnedExtent(global_shape_[1], block_shape_[1], process_grid_[1], @@ -332,10 +333,7 @@ class block_cyclic_distribution_t { local_index[0] >= 0 && local_index[0] < local_shape[0] && local_index[1] >= 0 && local_index[1] < local_shape[1], matxInvalidSize, "Local block-cyclic index is out of bounds"); - const index_t process_row = - static_cast(fragment) / process_grid_[1]; - const index_t process_col = - static_cast(fragment) % process_grid_[1]; + const auto [process_row, process_col] = ProcessCoordinate(fragment); return {MapIndex(local_index[0], block_shape_[0], process_grid_[0], process_row), MapIndex(local_index[1], block_shape_[1], process_grid_[1], @@ -346,7 +344,7 @@ class block_cyclic_distribution_t { return global_shape_ == other.global_shape_ && block_shape_ == other.block_shape_ && process_grid_ == other.process_grid_ && - endpoints_ == other.endpoints_; + endpoints_ == other.endpoints_ && layout_ == other.layout_; } private: @@ -369,12 +367,32 @@ class block_cyclic_distribution_t { return (local_block * processes + coordinate) * block + within_block; } + std::pair ProcessCoordinate(size_t fragment) const { + const index_t rank = static_cast(fragment); + if (layout_ == distributed_grid_layout::column_major) { + return {rank % process_grid_[0], rank / process_grid_[0]}; + } + return {rank / process_grid_[1], rank % process_grid_[1]}; + } + distributed_index_t<2> global_shape_{}; distributed_index_t<2> block_shape_{}; distributed_index_t<2> process_grid_{}; std::vector endpoints_; + distributed_grid_layout layout_ = distributed_grid_layout::row_major; }; +template +inline constexpr bool is_block_cyclic_distributed_tensor_v = [] { + if constexpr (is_distributed_tensor_v) { + return std::is_same_v::distribution_type, + block_cyclic_distribution_t>; + } + else { + return false; + } +}(); + /** One locally addressable, homogeneous piece of a distributed tensor. */ template struct local_fragment_t { using value_type = T; diff --git a/include/matx/distributed.h b/include/matx/distributed.h new file mode 100644 index 000000000..92a3ac123 --- /dev/null +++ b/include/matx/distributed.h @@ -0,0 +1,16 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include "matx/core/distributed_tensor.h" +#include "matx/executors/distributed.h" +#include "matx/transforms/distributed/fft_mg.h" + +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) +#include "matx/transforms/distributed/distributed_mp.h" +#endif diff --git a/include/matx/executors/distributed.h b/include/matx/executors/distributed.h index ce00457f6..1f9ed06a3 100644 --- a/include/matx/executors/distributed.h +++ b/include/matx/executors/distributed.h @@ -36,9 +36,24 @@ #include #include #include +#include #include #include +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) +#include +#endif +#ifdef MATX_EN_CUBLASMP +#if __has_include() +#include +#else +#include +#endif +#endif +#ifdef MATX_EN_CUSOLVERMP +#include +#endif + #include "matx/core/error.h" #include "matx/executors/cuda.h" @@ -53,6 +68,33 @@ inline void DistributedCheck(bool condition, matxError_t error, } } +inline void DistributedCheck(bool condition, matxError_t error, + const std::string &message) { + if (!condition) { + MATX_THROW(error, message); + } +} + +#ifdef MATX_EN_CUBLASMP +inline void CublasMpCheck(cublasMpStatus_t status, const char *operation) { + if (status != CUBLASMP_STATUS_SUCCESS) { + MATX_THROW(matxMatMulError, std::string(operation) + + " failed with cuBLASMp status " + + std::to_string(static_cast(status))); + } +} +#endif + +#ifdef MATX_EN_CUSOLVERMP +inline void CusolverMpCheck(cusolverStatus_t status, const char *operation) { + if (status != CUSOLVER_STATUS_SUCCESS) { + MATX_THROW(matxSolverError, std::string(operation) + + " failed with cuSOLVERMp status " + + std::to_string(static_cast(status))); + } +} +#endif + } // namespace detail /** Identifies a CUDA device in a process participating in distributed work. */ @@ -64,12 +106,14 @@ struct distributed_endpoint_t { const distributed_endpoint_t &) = default; }; +/** Ordering used to map communicator ranks onto a two-dimensional grid. */ +enum class distributed_grid_layout { row_major, column_major }; + /** * Lightweight topology identity shared by distributed tensors and executors. * - * Communication-library handles intentionally do not live here. A later - * multi-process executor can attach MPI/NCCL state without changing the tensor - * representation. + * Communication-library handles intentionally do not live here. An executor + * can attach collective state without changing the tensor representation. */ class distributed_context { private: @@ -161,13 +205,60 @@ class distributed_device_guard { bool changed_ = false; }; +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) +struct distributed_collective_state { + ~distributed_collective_state() { + int previous_device = 0; + const bool restore_device = + cudaGetDevice(&previous_device) == cudaSuccess && device_id >= 0; + if (device_id >= 0) { + (void)cudaSetDevice(device_id); + } +#ifdef MATX_EN_CUSOLVERMP + if (cusolver_grid != nullptr) { + (void)cusolverMpDestroyGrid(cusolver_grid); + } + if (cusolver_handle != nullptr) { + (void)cusolverMpDestroy(cusolver_handle); + } +#endif +#ifdef MATX_EN_CUBLASMP + if (cublas_grid != nullptr) { + (void)cublasMpGridDestroy(cublas_grid); + } + if (cublas_handle != nullptr) { + (void)cublasMpDestroy(cublas_handle); + } +#endif + if (restore_device) { + (void)cudaSetDevice(previous_device); + } + } + + int device_id = -1; + ncclComm_t communicator = nullptr; + int process_rows = 0; + int process_columns = 0; + distributed_grid_layout grid_layout = distributed_grid_layout::row_major; +#ifdef MATX_EN_CUBLASMP + cublasMpHandle_t cublas_handle = nullptr; + cublasMpGrid_t cublas_grid = nullptr; +#endif +#ifdef MATX_EN_CUSOLVERMP + cusolverMpHandle_t cusolver_handle = nullptr; + cusolverMpGrid_t cusolver_grid = nullptr; +#endif +}; +#endif + } // namespace detail /** - * Experimental single-process, multi-GPU executor. + * Experimental distributed CUDA executor. * - * One non-blocking CUDA stream is owned per local endpoint. Multi-process - * collective support is deliberately left behind this executor boundary. + * One non-blocking CUDA stream is owned per local endpoint. When an NVIDIA MP + * backend is enabled, an optional constructor attaches a borrowed NCCL + * communicator and creates the backend handles needed for collective work. */ class distributedCUDAExecutor { public: @@ -192,23 +283,48 @@ class distributedCUDAExecutor { } explicit distributedCUDAExecutor(std::vector local_devices) - : distributedCUDAExecutor( - distributed_context{std::move(local_devices)}) {} + : distributedCUDAExecutor(distributed_context{std::move(local_devices)}) { + } + +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) + distributedCUDAExecutor( + distributed_context context, ncclComm_t communicator, int process_rows, + int process_columns, + distributed_grid_layout layout = distributed_grid_layout::row_major) + : distributedCUDAExecutor(std::move(context)) { + try { + InitializeCollectives(communicator, process_rows, process_columns, + layout); + } catch (...) { + DestroyStreams(); + throw; + } + } +#endif distributedCUDAExecutor(const distributedCUDAExecutor &) = delete; - distributedCUDAExecutor & - operator=(const distributedCUDAExecutor &) = delete; + distributedCUDAExecutor &operator=(const distributedCUDAExecutor &) = delete; distributedCUDAExecutor(distributedCUDAExecutor &&other) noexcept - : context_{std::move(other.context_)}, - streams_{std::move(other.streams_)}, - transfer_count_{other.transfer_count_.load(std::memory_order_relaxed)} { + : context_{std::move(other.context_)}, streams_{std::move( + other.streams_)}, + transfer_count_{other.transfer_count_.load(std::memory_order_relaxed)} +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) + , + collectives_{std::move(other.collectives_)} +#endif + { other.streams_.clear(); } distributedCUDAExecutor &operator=(distributedCUDAExecutor &&) = delete; - ~distributedCUDAExecutor() { DestroyStreams(); } + ~distributedCUDAExecutor() { +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) + collectives_.reset(); +#endif + DestroyStreams(); + } const distributed_context &Context() const noexcept { return context_; } uint64_t ContextId() const noexcept { return context_.Id(); } @@ -249,6 +365,36 @@ class distributedCUDAExecutor { return transfer_count_.load(std::memory_order_relaxed); } +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) + bool HasCollectiveResources() const noexcept { + return collectives_ != nullptr; + } + int ProcessRows() const noexcept { return collectives_->process_rows; } + int ProcessColumns() const noexcept { return collectives_->process_columns; } + distributed_grid_layout GridLayout() const noexcept { + return collectives_->grid_layout; + } +#else + bool HasCollectiveResources() const noexcept { return false; } +#endif + +#ifdef MATX_EN_CUBLASMP + cublasMpHandle_t CublasHandle() const noexcept { + return collectives_->cublas_handle; + } + cublasMpGrid_t CublasGrid() const noexcept { + return collectives_->cublas_grid; + } +#endif +#ifdef MATX_EN_CUSOLVERMP + cusolverMpHandle_t CusolverHandle() const noexcept { + return collectives_->cusolver_handle; + } + cusolverMpGrid_t CusolverGrid() const noexcept { + return collectives_->cusolver_grid; + } +#endif + private: struct stream_entry_t { int device_id; @@ -269,9 +415,71 @@ class distributedCUDAExecutor { streams_.clear(); } +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) + void InitializeCollectives(ncclComm_t communicator, int process_rows, + int process_columns, + distributed_grid_layout layout) { + detail::DistributedCheck( + communicator != nullptr, matxInvalidExecutor, + "Collective execution requires an NCCL communicator"); + detail::DistributedCheck( + context_.LocalDevices().size() == 1, matxNotSupported, + "NVIDIA MP execution currently supports one GPU per process"); + detail::DistributedCheck( + process_rows > 0 && process_columns > 0 && + process_rows * process_columns == context_.ProcessCount(), + matxInvalidSize, + "The MP process grid must contain every distributed process"); + + auto state = std::make_unique(); + state->device_id = context_.LocalDevices().front(); + state->communicator = communicator; + state->process_rows = process_rows; + state->process_columns = process_columns; + state->grid_layout = layout; + const distributed_endpoint_t endpoint{context_.ProcessRank(), + context_.LocalDevices().front()}; + auto local = LocalExecutor(endpoint); + detail::distributed_device_guard guard{endpoint.device_id}; + +#ifdef MATX_EN_CUBLASMP + detail::CublasMpCheck( + cublasMpCreate(&state->cublas_handle, local.getStream()), + "cublasMpCreate"); + detail::CublasMpCheck( + cublasMpGridCreate(state->process_rows, state->process_columns, + state->grid_layout == + distributed_grid_layout::column_major + ? CUBLASMP_GRID_LAYOUT_COL_MAJOR + : CUBLASMP_GRID_LAYOUT_ROW_MAJOR, + state->communicator, &state->cublas_grid), + "cublasMpGridCreate"); +#endif + +#ifdef MATX_EN_CUSOLVERMP + detail::CusolverMpCheck(cusolverMpCreate(&state->cusolver_handle, + endpoint.device_id, + local.getStream()), + "cusolverMpCreate"); + detail::CusolverMpCheck( + cusolverMpCreateDeviceGrid( + state->cusolver_handle, &state->cusolver_grid, state->communicator, + state->process_rows, state->process_columns, + state->grid_layout == distributed_grid_layout::column_major + ? CUSOLVERMP_GRID_MAPPING_COL_MAJOR + : CUSOLVERMP_GRID_MAPPING_ROW_MAJOR), + "cusolverMpCreateDeviceGrid"); +#endif + collectives_ = std::move(state); + } +#endif + distributed_context context_; std::vector streams_; mutable std::atomic transfer_count_{0}; +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) + std::unique_ptr collectives_; +#endif }; } // namespace matx diff --git a/include/matx/operators/chol.h b/include/matx/operators/chol.h index b84e9a825..bdd095d94 100644 --- a/include/matx/operators/chol.h +++ b/include/matx/operators/chol.h @@ -39,6 +39,9 @@ #include "matx/operators/base_operator.h" #include "matx/core/operator_options.h" #include "matx/transforms/chol/chol_cuda.h" +#ifdef MATX_EN_CUSOLVERMP + #include "matx/transforms/distributed/distributed_mp.h" +#endif #if defined(MATX_EN_MATHDX) && defined(__CUDACC__) #include "matx/transforms/solver_cusolverdx.h" #endif @@ -298,6 +301,43 @@ namespace detail { }; } +namespace experimental::detail { + +template +class distributed_mp_chol_op { +public: + using distributed_expression = bool; + using value_type = typename remove_cvref_t::value_type; + + distributed_mp_chol_op(const OpA &a, SolverFillMode uplo) + : a_{a}, uplo_{uplo} {} + + template + void ExecuteTo(Out &out, Executor &executor) const { + static_assert(is_block_cyclic_distributed_tensor_v, + "Block-cyclic Cholesky requires a block-cyclic output"); + if constexpr (std::is_same_v, + distributedCUDAExecutor>) { +#ifdef MATX_EN_CUSOLVERMP + CholMpImpl(out, a_, executor, uplo_); +#else + MATX_THROW(matxNotSupported, + "Block-cyclic Cholesky requires MATX_EN_CUSOLVERMP"); +#endif + } + else { + MATX_THROW(matxInvalidExecutor, + "Block-cyclic Cholesky requires distributedCUDAExecutor"); + } + } + +private: + remove_cvref_t a_; + SolverFillMode uplo_; +}; + +} // namespace experimental::detail + /** * Performs a Cholesky factorization, saving the result in either the upper or * lower triangle of the output. @@ -319,15 +359,21 @@ namespace detail { template __MATX_INLINE__ auto chol(const OpA &a, SolverFillMode uplo = SolverFillMode::UPPER) { if constexpr (is_distributed_tensor_v) { - static_assert(remove_cvref_t::Rank() >= 3, - "Distributed Cholesky factorization requires a batch " - "dimension followed by matrix dimensions"); - auto local_chol = [uplo](const auto &local_a) { - return detail::CholOp(local_a, uplo); - }; - return experimental::detail::make_distributed_local_transform< - typename remove_cvref_t::value_type, 2>( - std::move(local_chol), a); + if constexpr ( + experimental::is_block_cyclic_distributed_tensor_v) { + return experimental::detail::distributed_mp_chol_op{a, uplo}; + } + else { + static_assert(remove_cvref_t::Rank() >= 3, + "Distributed Cholesky needs a block-cyclic rank-2 matrix " + "or a batch dimension followed by local matrix dimensions"); + auto local_chol = [uplo](const auto &local_a) { + return detail::CholOp(local_a, uplo); + }; + return experimental::detail::make_distributed_local_transform< + typename remove_cvref_t::value_type, 2>( + std::move(local_chol), a); + } } else { return detail::CholOp(a, uplo); diff --git a/include/matx/operators/fft.h b/include/matx/operators/fft.h index 6ad227e71..f967506a8 100644 --- a/include/matx/operators/fft.h +++ b/include/matx/operators/fft.h @@ -41,6 +41,7 @@ #include "matx/core/operator_options.h" #include "matx/core/log.h" +#include "matx/transforms/distributed/fft_mg.h" #include "matx/transforms/fft/fft_cuda.h" #ifdef MATX_EN_CPU_FFT #include "matx/transforms/fft/fft_fftw.h" @@ -519,6 +520,120 @@ namespace matx }; } + namespace experimental::detail { + + template + class distributed_fft_op { + public: + using distributed_expression = bool; + using value_type = ValueType; + + distributed_fft_op(const OpA &a, index_t fft_size, FFTNorm norm) + : a_{a}, fft_size_{fft_size}, norm_{norm} {} + + template + void ExecuteTo(Out &out, Executor &executor) const { + static_assert(is_distributed_tensor_v, + "Distributed FFT requires a distributed output"); + static_assert(is_distributed_executor_v, + "Distributed FFT requires a distributed executor"); + static_assert(remove_cvref_t::Rank() == + remove_cvref_t::Rank(), + "Distributed FFT input and output ranks must match"); + static_assert( + std::is_same_v::value_type, + value_type>, + "Distributed FFT output type does not match the transform"); + + if (TransformDimensionIsLocal(a_.DistributionDescriptor())) { + auto local_fft = [fft_size = fft_size_, + norm = norm_](const auto &local_a) { + using local_type = remove_cvref_t; + if constexpr (Direction == matx::detail::FFTDirection::FORWARD) { + constexpr auto fft_type = + matx::detail::ComplexInType(); + return matx::detail::FFTOp< + local_type, matx::detail::no_permute_t, Direction, + fft_type>(local_a, fft_size, matx::detail::no_permute_t{}, + norm); + } + else { + return matx::detail::FFTOp< + local_type, matx::detail::no_permute_t, Direction, + matx::detail::FFTType::C2C>( + local_a, fft_size, matx::detail::no_permute_t{}, norm); + } + }; + auto local_transform = + make_distributed_local_transform( + std::move(local_fft), a_); + local_transform.ExecuteTo(out, executor); + return; + } + + if constexpr (remove_cvref_t::Rank() != 1) { + MATX_THROW( + matxNotSupported, + "A partitioned distributed FFT dimension currently requires " + "a rank-1 complex transform"); + } + else { + matx::detail::DistributedCheck( + fft_size_ == 0 || fft_size_ == a_.Size(0), matxNotSupported, + "cuFFT multi-GPU does not currently support FFT resizing"); + for (size_t fragment = 0; + fragment < a_.DistributionDescriptor().FragmentCount(); + ++fragment) { + matx::detail::DistributedCheck( + a_.DistributionDescriptor() + .FragmentEndpoint(fragment) + .process_rank == a_.ProcessRank(), + matxNotSupported, + "A transform dimension spanning processes requires cuFFTMp, " + "whose distributed allocation path is not yet available"); + } + if constexpr ( + std::is_same_v< + typename remove_cvref_t::value_type, ValueType> && + (std::is_same_v> || + std::is_same_v>)) { + FftMgImpl(out, a_, executor, norm_); + } + else { + MATX_THROW(matxNotSupported, + "cuFFT multi-GPU currently supports only complex-to-" + "complex float and double transforms"); + } + } + } + + private: + template + static bool TransformDimensionIsLocal( + const Distribution &distribution) { + constexpr int rank = Distribution::Rank(); + for (size_t fragment = 0; fragment < distribution.FragmentCount(); + ++fragment) { + distributed_index_t local_zero{}; + const auto origin = + distribution.LocalToGlobal(fragment, local_zero); + const auto local_shape = distribution.LocalShape(fragment); + if (origin[rank - 1] != 0 || + local_shape[rank - 1] != distribution.GlobalShape()[rank - 1]) { + return false; + } + } + return true; + } + + remove_cvref_t a_; + index_t fft_size_; + FFTNorm norm_; + }; + + } // namespace experimental::detail + /** * Run a 1D FFT with a cached plan @@ -543,23 +658,13 @@ namespace matx constexpr auto fft_type = detail::ComplexInType(); const index_t fft_size_ = static_cast(fft_size); if constexpr (is_distributed_tensor_v) { - static_assert(remove_cvref_t::Rank() >= 2, - "Distributed FFT requires a batch dimension followed by " - "the transform dimension"); using input_type = typename remove_cvref_t::value_type; using output_type = std::conditional_t< is_complex_v, input_type, typename detail::scalar_to_complex::ctype>; - auto local_fft = [fft_size_, norm](const auto &local_a) { - using local_type = remove_cvref_t; - constexpr auto local_fft_type = detail::ComplexInType(); - return detail::FFTOp( - local_a, fft_size_, detail::no_permute_t{}, norm); - }; - return experimental::detail::make_distributed_local_transform< - output_type, 1>( - std::move(local_fft), a); + return experimental::detail::distributed_fft_op< + output_type, detail::FFTDirection::FORWARD, OpA>{ + a, fft_size_, norm}; } else { return detail::FFTOp __MATX_INLINE__ auto ifft(const OpA &a, uint64_t fft_size = 0, FFTNorm norm = FFTNorm::BACKWARD) { const index_t fft_size_ = static_cast(fft_size); - return detail::FFTOp(a, fft_size_, detail::no_permute_t{} , norm); + if constexpr (is_distributed_tensor_v) { + using output_type = typename remove_cvref_t::value_type; + return experimental::detail::distributed_fft_op< + output_type, detail::FFTDirection::BACKWARD, OpA>{ + a, fft_size_, norm}; + } + else { + return detail::FFTOp( + a, fft_size_, detail::no_permute_t{}, norm); + } } /** @@ -703,6 +819,8 @@ namespace matx */ template __MATX_INLINE__ auto ifft(const OpA &a, const int32_t (&axis)[1], uint64_t fft_size = 0, FFTNorm norm = FFTNorm::BACKWARD) { + static_assert(!is_distributed_tensor_v, + "Axis-selecting distributed IFFT is not supported"); if constexpr (is_dynamic_rank_op_v>) { auto perm = detail::getPermuteDims(detail::get_dyn_rank(a), axis); const index_t fft_size_ = static_cast(fft_size); diff --git a/include/matx/operators/matmul.h b/include/matx/operators/matmul.h index 87608f138..6ebdb850a 100644 --- a/include/matx/operators/matmul.h +++ b/include/matx/operators/matmul.h @@ -41,6 +41,9 @@ #include "matx/core/log.h" #include "matx/transforms/matmul/matmul_cuda.h" #include "matx/transforms/matmul/matmul_cusparse.h" +#ifdef MATX_EN_CUBLASMP + #include "matx/transforms/distributed/distributed_mp.h" +#endif #ifdef MATX_EN_CPU_MATMUL #include "matx/transforms/matmul/matmul_cblas.h" #endif @@ -456,6 +459,48 @@ namespace matx }; } + namespace experimental::detail { + + template + class distributed_mp_matmul_op { + public: + using distributed_expression = bool; + using value_type = typename remove_cvref_t::value_type; + + distributed_mp_matmul_op(const OpA &a, const OpB &b, float alpha, + float beta) + : a_{a}, b_{b}, alpha_{alpha}, beta_{beta} {} + + template + void ExecuteTo(Out &out, Executor &executor) const { + static_assert(is_block_cyclic_distributed_tensor_v, + "Block-cyclic matmul requires a block-cyclic output"); + if constexpr (std::is_same_v, + distributedCUDAExecutor>) { +#ifdef MATX_EN_CUBLASMP + MatmulMpImpl(out, a_, b_, executor, static_cast(alpha_), + static_cast(beta_)); +#else + MATX_THROW(matxNotSupported, + "Block-cyclic matmul requires MATX_EN_CUBLASMP"); +#endif + } + else { + MATX_THROW( + matxInvalidExecutor, + "Block-cyclic matmul requires distributedCUDAExecutor"); + } + } + + private: + remove_cvref_t a_; + remove_cvref_t b_; + float alpha_; + float beta_; + }; + + } // namespace experimental::detail + /** * Run a GEMM (generic matrix multiply)) @@ -490,21 +535,34 @@ namespace matx static_assert(is_distributed_tensor_v && is_distributed_tensor_v, "matmul requires both inputs to be distributed"); - static_assert(remove_cvref_t::Rank() == - remove_cvref_t::Rank(), - "First-pass distributed matmul requires equal input ranks"); - static_assert(remove_cvref_t::Rank() >= 3, - "Distributed matmul requires a batch dimension followed " - "by matrix dimensions"); - - auto local_matmul = [alpha, beta](const auto &local_a, - const auto &local_b) { - return detail::MatMulOp(local_a, local_b, alpha, beta, - detail::no_permute_t{}); - }; - return experimental::detail::make_distributed_local_transform< - typename remove_cvref_t::value_type, 2>( - std::move(local_matmul), A, B); + if constexpr ( + experimental::is_block_cyclic_distributed_tensor_v || + experimental::is_block_cyclic_distributed_tensor_v) { + static_assert( + experimental::is_block_cyclic_distributed_tensor_v && + experimental::is_block_cyclic_distributed_tensor_v, + "Block-cyclic matmul requires both inputs to use " + "block_cyclic_distribution_t"); + return experimental::detail::distributed_mp_matmul_op{ + A, B, alpha, beta}; + } + else { + static_assert( + remove_cvref_t::Rank() == remove_cvref_t::Rank(), + "Batch-sharded distributed matmul requires equal input ranks"); + static_assert(remove_cvref_t::Rank() >= 3, + "Distributed matmul needs block-cyclic rank-2 matrices " + "or a batch dimension followed by local matrix dimensions"); + + auto local_matmul = [alpha, beta](const auto &local_a, + const auto &local_b) { + return detail::MatMulOp(local_a, local_b, alpha, beta, + detail::no_permute_t{}); + }; + return experimental::detail::make_distributed_local_transform< + typename remove_cvref_t::value_type, 2>( + std::move(local_matmul), A, B); + } } else { return detail::MatMulOp(A, B, alpha, beta, detail::no_permute_t{}); diff --git a/include/matx/transforms/distributed/distributed_mp.h b/include/matx/transforms/distributed/distributed_mp.h new file mode 100644 index 000000000..6e8e5165d --- /dev/null +++ b/include/matx/transforms/distributed/distributed_mp.h @@ -0,0 +1,383 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#if !defined(MATX_EN_CUBLASMP) && !defined(MATX_EN_CUSOLVERMP) +#error "distributed_mp.h requires MATX_EN_CUBLASMP or MATX_EN_CUSOLVERMP" +#endif + +#include +#include +#include +#include + +#include "matx/core/distributed_tensor.h" +#include "matx/core/make_tensor.h" +#include "matx/core/operator_options.h" +#include "matx/core/type_utils.h" +#include "matx/executors/distributed.h" + +namespace matx::experimental { +namespace detail { + +template +inline constexpr bool is_mp_matrix_v = + is_distributed_tensor_v &&remove_cvref_t::Rank() == 2 && + std::is_same_v::distribution_type, + block_cyclic_distribution_t>; + +template +void ValidateMpMatrix(const Tensor &tensor, + const distributedCUDAExecutor &executor, + const char *name) { + static_assert(is_mp_matrix_v, + "MP matrices must use block_cyclic_distribution_t"); + matx::detail::DistributedCheck( + executor.HasCollectiveResources(), matxInvalidExecutor, + "Block-cyclic execution requires collective resources on its " + "distributedCUDAExecutor"); + matx::detail::DistributedCheck( + tensor.ContextId() == executor.ContextId(), matxInvalidExecutor, + std::string(name) + " belongs to a different distributed context"); + matx::detail::DistributedCheck( + tensor.LocalFragmentCount() == 1, matxNotSupported, + std::string(name) + " must have one local fragment per process"); + const auto &distribution = tensor.DistributionDescriptor(); + matx::detail::DistributedCheck( + distribution.ProcessGrid()[0] == executor.ProcessRows() && + distribution.ProcessGrid()[1] == executor.ProcessColumns() && + distribution.GridLayout() == executor.GridLayout(), + matxInvalidSize, + std::string(name) + + " process grid or rank ordering does not match its MP executor"); + for (size_t rank = 0; rank < distribution.FragmentCount(); ++rank) { + matx::detail::DistributedCheck( + distribution.FragmentEndpoint(rank).process_rank == + static_cast(rank), + matxInvalidParameter, + std::string(name) + + " endpoints must be ordered by NCCL communicator rank"); + } +} + +template +void ValidateSameMpGrid(const A &a, const B &b) { + const auto &ad = a.DistributionDescriptor(); + const auto &bd = b.DistributionDescriptor(); + matx::detail::DistributedCheck( + ad.ProcessGrid() == bd.ProcessGrid() && + ad.GridLayout() == bd.GridLayout(), + matxInvalidSize, "MP matrix process grids and rank ordering must match"); + for (size_t rank = 0; rank < ad.FragmentCount(); ++rank) { + matx::detail::DistributedCheck( + ad.FragmentEndpoint(rank) == bd.FragmentEndpoint(rank), + matxInvalidParameter, "MP matrix rank mappings must match"); + } +} + +template struct column_major_buffer { + T *data = nullptr; + index_t rows = 0; + index_t columns = 0; +}; + +template +auto PackColumnMajor(const Tensor &tensor, cudaExecutor &executor) { + using value_type = typename remove_cvref_t::value_type; + const auto &local = tensor.LocalView(0); + column_major_buffer packed{nullptr, local.Size(0), local.Size(1)}; + const size_t bytes = + static_cast(packed.rows * packed.columns) * sizeof(value_type); + MATX_CUDA_CHECK(cudaMallocAsync(reinterpret_cast(&packed.data), + bytes, executor.getStream())); + index_t shape[2]{packed.rows, packed.columns}; + index_t strides[2]{1, packed.rows}; + auto packed_view = make_tensor(packed.data, shape, strides, false); + (packed_view = local).run(executor); + return packed; +} + +template +void UnpackColumnMajor(Tensor &tensor, const column_major_buffer &packed, + cudaExecutor &executor) { + index_t shape[2]{packed.rows, packed.columns}; + index_t strides[2]{1, packed.rows}; + auto packed_view = make_tensor(packed.data, shape, strides, false); + (tensor.LocalView(0) = packed_view).run(executor); +} + +template +void FreeColumnMajor(column_major_buffer &buffer, cudaExecutor &executor) { + if (buffer.data != nullptr) { + MATX_CUDA_CHECK(cudaFreeAsync(buffer.data, executor.getStream())); + buffer.data = nullptr; + } +} + +#ifdef MATX_EN_CUBLASMP +class cublas_mp_descriptor { +public: + cublas_mp_descriptor(const block_cyclic_distribution_t &distribution, + index_t leading_dimension, cudaDataType_t type, + cublasMpGrid_t grid) { + const auto shape = distribution.GlobalShape(); + const auto block = distribution.BlockShape(); + matx::detail::CublasMpCheck(cublasMpMatrixDescriptorCreate( + shape[0], shape[1], block[0], block[1], 0, + 0, std::max(1, leading_dimension), + type, grid, &descriptor_), + "cublasMpMatrixDescriptorCreate"); + } + cublas_mp_descriptor(const cublas_mp_descriptor &) = delete; + ~cublas_mp_descriptor() { + if (descriptor_ != nullptr) { + (void)cublasMpMatrixDescriptorDestroy(descriptor_); + } + } + operator cublasMpMatrixDescriptor_t() const noexcept { return descriptor_; } + +private: + cublasMpMatrixDescriptor_t descriptor_ = nullptr; +}; +#endif + +#ifdef MATX_EN_CUSOLVERMP +class cusolver_mp_descriptor { +public: + cusolver_mp_descriptor(const block_cyclic_distribution_t &distribution, + index_t leading_dimension, cudaDataType_t type, + cusolverMpGrid_t grid) { + const auto shape = distribution.GlobalShape(); + const auto block = distribution.BlockShape(); + matx::detail::CusolverMpCheck( + cusolverMpCreateMatrixDesc(&descriptor_, grid, type, shape[0], shape[1], + block[0], block[1], 0, 0, + std::max(1, leading_dimension)), + "cusolverMpCreateMatrixDesc"); + } + cusolver_mp_descriptor(const cusolver_mp_descriptor &) = delete; + ~cusolver_mp_descriptor() { + if (descriptor_ != nullptr) { + (void)cusolverMpDestroyMatrixDesc(descriptor_); + } + } + operator cusolverMpMatrixDescriptor_t() const noexcept { return descriptor_; } + +private: + cusolverMpMatrixDescriptor_t descriptor_ = nullptr; +}; +#endif + +} // namespace detail + +#ifdef MATX_EN_CUBLASMP +/** + * Collective block-cyclic matrix multiplication using cuBLASMp. + * + * Every process in the executor grid must enter this function in the same + * order. MatX local fragments remain ordinary row-major tensor views; this + * adapter packs them into the column-major local layout required by cuBLASMp. + */ +namespace detail { + +template +void MatmulMpImpl(Out &out, const A &a, const B &b, + distributedCUDAExecutor &executor, + typename remove_cvref_t::value_type alpha = + typename remove_cvref_t::value_type{1}, + typename remove_cvref_t::value_type beta = + typename remove_cvref_t::value_type{0}) { + using value_type = typename remove_cvref_t::value_type; + static_assert(detail::is_mp_matrix_v && detail::is_mp_matrix_v && + detail::is_mp_matrix_v, + "Block-cyclic matmul requires rank-2 distributed tensors"); + static_assert( + std::is_same_v::value_type> && + std::is_same_v::value_type>, + "Block-cyclic matmul requires identical input and output types"); + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v> || + std::is_same_v>, + "Block-cyclic matmul supports float, double, complex, and " + "complex"); + + detail::ValidateMpMatrix(out, executor, "block-cyclic matmul output"); + detail::ValidateMpMatrix(a, executor, "block-cyclic matmul left input"); + detail::ValidateMpMatrix(b, executor, "block-cyclic matmul right input"); + detail::ValidateSameMpGrid(out, a); + detail::ValidateSameMpGrid(out, b); + matx::detail::DistributedCheck( + a.Size(1) == b.Size(0) && out.Size(0) == a.Size(0) && + out.Size(1) == b.Size(1), + matxInvalidSize, + "Block-cyclic matmul matrix dimensions are incompatible"); + + const auto endpoint = out.LocalFragment(0).endpoint; + executor.ForEndpoint(endpoint, [&](cudaExecutor &local_executor) { + auto packed_a = detail::PackColumnMajor(a, local_executor); + auto packed_b = detail::PackColumnMajor(b, local_executor); + auto packed_c = detail::PackColumnMajor(out, local_executor); + + detail::cublas_mp_descriptor desc_a{ + a.DistributionDescriptor(), packed_a.rows, + matx::detail::MatXTypeToCudaType(), + executor.CublasGrid()}; + detail::cublas_mp_descriptor desc_b{ + b.DistributionDescriptor(), packed_b.rows, + matx::detail::MatXTypeToCudaType(), + executor.CublasGrid()}; + detail::cublas_mp_descriptor desc_c{ + out.DistributionDescriptor(), packed_c.rows, + matx::detail::MatXTypeToCudaType(), + executor.CublasGrid()}; + + size_t device_workspace_size = 0; + size_t host_workspace_size = 0; + matx::detail::CublasMpCheck( + cublasMpGemm_bufferSize( + executor.CublasHandle(), CUBLAS_OP_N, CUBLAS_OP_N, out.Size(0), + out.Size(1), a.Size(1), &alpha, packed_a.data, 1, 1, desc_a, + packed_b.data, 1, 1, desc_b, &beta, packed_c.data, 1, 1, desc_c, + matx::detail::MatXTypeToCudaComputeType(), + &device_workspace_size, &host_workspace_size), + "cublasMpGemm_bufferSize"); + + void *device_workspace = nullptr; + if (device_workspace_size != 0) { + MATX_CUDA_CHECK(cudaMallocAsync(&device_workspace, device_workspace_size, + local_executor.getStream())); + } + std::vector host_workspace(host_workspace_size); + matx::detail::CublasMpCheck( + cublasMpGemm( + executor.CublasHandle(), CUBLAS_OP_N, CUBLAS_OP_N, out.Size(0), + out.Size(1), a.Size(1), &alpha, packed_a.data, 1, 1, desc_a, + packed_b.data, 1, 1, desc_b, &beta, packed_c.data, 1, 1, desc_c, + matx::detail::MatXTypeToCudaComputeType(), + device_workspace, device_workspace_size, host_workspace.data(), + host_workspace_size), + "cublasMpGemm"); + + detail::UnpackColumnMajor(out, packed_c, local_executor); + detail::FreeColumnMajor(packed_a, local_executor); + detail::FreeColumnMajor(packed_b, local_executor); + detail::FreeColumnMajor(packed_c, local_executor); + if (device_workspace != nullptr) { + MATX_CUDA_CHECK( + cudaFreeAsync(device_workspace, local_executor.getStream())); + } + // cuBLASMp may retain the host workspace until its stream work completes. + MATX_CUDA_CHECK(cudaStreamSynchronize(local_executor.getStream())); + }); +} + +} // namespace detail +#endif + +#ifdef MATX_EN_CUSOLVERMP +/** + * Collective block-cyclic Cholesky factorization using cuSOLVERMp. + */ +namespace detail { + +template +void CholMpImpl(Out &out, const A &a, distributedCUDAExecutor &executor, + SolverFillMode uplo = SolverFillMode::UPPER) { + using value_type = typename remove_cvref_t::value_type; + static_assert(detail::is_mp_matrix_v && detail::is_mp_matrix_v, + "Block-cyclic Cholesky requires rank-2 distributed tensors"); + static_assert( + std::is_same_v::value_type>, + "Block-cyclic Cholesky input and output types must match"); + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v> || + std::is_same_v>, + "Block-cyclic Cholesky supports float, double, complex, and " + "complex"); + + detail::ValidateMpMatrix(out, executor, "block-cyclic Cholesky output"); + detail::ValidateMpMatrix(a, executor, "block-cyclic Cholesky input"); + detail::ValidateSameMpGrid(out, a); + matx::detail::DistributedCheck( + a.Size(0) == a.Size(1) && out.Size(0) == a.Size(0) && + out.Size(1) == a.Size(1), + matxInvalidSize, + "Block-cyclic Cholesky requires equal square input and output"); + matx::detail::DistributedCheck( + out.DistributionDescriptor().Compatible(a.DistributionDescriptor()), + matxInvalidParameter, + "Block-cyclic Cholesky input and output layouts must match"); + matx::detail::DistributedCheck(a.DistributionDescriptor().BlockShape()[0] == + a.DistributionDescriptor().BlockShape()[1], + matxInvalidSize, + "cuSOLVERMp Cholesky requires square blocks"); + + const auto endpoint = out.LocalFragment(0).endpoint; + executor.ForEndpoint(endpoint, [&](cudaExecutor &local_executor) { + auto packed = detail::PackColumnMajor(a, local_executor); + detail::cusolver_mp_descriptor descriptor{ + a.DistributionDescriptor(), packed.rows, + matx::detail::MatXTypeToCudaType(), + executor.CusolverGrid()}; + const cublasFillMode_t fill = uplo == SolverFillMode::UPPER + ? CUBLAS_FILL_MODE_UPPER + : CUBLAS_FILL_MODE_LOWER; + + size_t device_workspace_size = 0; + size_t host_workspace_size = 0; + matx::detail::CusolverMpCheck( + cusolverMpPotrf_bufferSize( + executor.CusolverHandle(), fill, a.Size(0), packed.data, 1, 1, + descriptor, matx::detail::MatXTypeToCudaType(), + &device_workspace_size, &host_workspace_size), + "cusolverMpPotrf_bufferSize"); + + void *device_workspace = nullptr; + int *info = nullptr; + if (device_workspace_size != 0) { + MATX_CUDA_CHECK(cudaMallocAsync(&device_workspace, device_workspace_size, + local_executor.getStream())); + } + MATX_CUDA_CHECK(cudaMallocAsync(reinterpret_cast(&info), + sizeof(int), local_executor.getStream())); + std::vector host_workspace(host_workspace_size); + matx::detail::CusolverMpCheck( + cusolverMpPotrf(executor.CusolverHandle(), fill, a.Size(0), packed.data, + 1, 1, descriptor, + matx::detail::MatXTypeToCudaType(), + device_workspace, device_workspace_size, + host_workspace.data(), host_workspace_size, info), + "cusolverMpPotrf"); + + detail::UnpackColumnMajor(out, packed, local_executor); + int host_info = 0; + MATX_CUDA_CHECK(cudaMemcpyAsync(&host_info, info, sizeof(int), + cudaMemcpyDeviceToHost, + local_executor.getStream())); + detail::FreeColumnMajor(packed, local_executor); + if (device_workspace != nullptr) { + MATX_CUDA_CHECK( + cudaFreeAsync(device_workspace, local_executor.getStream())); + } + MATX_CUDA_CHECK(cudaFreeAsync(info, local_executor.getStream())); + // cuSOLVERMp may retain the host workspace until its stream work completes. + MATX_CUDA_CHECK(cudaStreamSynchronize(local_executor.getStream())); + matx::detail::DistributedCheck( + host_info == 0, matxSolverError, + "cusolverMpPotrf reported a non-positive-definite matrix or invalid " + "argument (info=" + + std::to_string(host_info) + ")"); + }); +} + +} // namespace detail +#endif + +} // namespace matx::experimental diff --git a/include/matx/transforms/distributed/fft_mg.h b/include/matx/transforms/distributed/fft_mg.h new file mode 100644 index 000000000..037a3bdcc --- /dev/null +++ b/include/matx/transforms/distributed/fft_mg.h @@ -0,0 +1,248 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include + +#include "matx/core/distributed_tensor.h" +#include "matx/core/operator_options.h" + +namespace matx::experimental { +namespace detail { + +inline void CufftMgCheck(cufftResult status, const char *operation) { + if (status != CUFFT_SUCCESS) { + MATX_THROW(matxCufftError, std::string(operation) + + " failed with cuFFT status " + + std::to_string(static_cast(status))); + } +} + +class cufft_mg_plan { +public: + cufft_mg_plan() { CufftMgCheck(cufftCreate(&plan_), "cufftCreate"); } + cufft_mg_plan(const cufft_mg_plan &) = delete; + ~cufft_mg_plan() { + if (input_ != nullptr) { + (void)cufftXtFree(input_); + } + if (output_ != nullptr) { + (void)cufftXtFree(output_); + } + if (plan_ != 0) { + (void)cufftDestroy(plan_); + } + } + + cufftHandle Handle() const noexcept { return plan_; } + cudaLibXtDesc **InputAddress() noexcept { return &input_; } + cudaLibXtDesc **OutputAddress() noexcept { return &output_; } + cudaLibXtDesc *Input() const noexcept { return input_; } + cudaLibXtDesc *Output() const noexcept { return output_; } + +private: + cufftHandle plan_ = 0; + cudaLibXtDesc *input_ = nullptr; + cudaLibXtDesc *output_ = nullptr; +}; + +template class pinned_buffer { +public: + explicit pinned_buffer(size_t count) : count_{count} { + MATX_CUDA_CHECK( + cudaMallocHost(reinterpret_cast(&data_), count * sizeof(T))); + } + pinned_buffer(const pinned_buffer &) = delete; + ~pinned_buffer() { + if (data_ != nullptr) { + (void)cudaFreeHost(data_); + } + } + T *Data() noexcept { return data_; } + const T *Data() const noexcept { return data_; } + size_t Size() const noexcept { return count_; } + +private: + T *data_ = nullptr; + size_t count_ = 0; +}; + +template +void GatherMgInput(const Tensor &tensor, + pinned_buffer &host) { + const auto &distribution = tensor.DistributionDescriptor(); + for (size_t fragment = 0; fragment < distribution.FragmentCount(); + ++fragment) { + const auto local_zero = distributed_index_t<1>{0}; + const index_t origin = distribution.LocalToGlobal(fragment, local_zero)[0]; + const auto &local = tensor.LocalFragmentForDistributionIndex(fragment).view; + matx::detail::distributed_device_guard guard{ + distribution.FragmentEndpoint(fragment).device_id}; + MATX_CUDA_CHECK(cudaMemcpy(host.Data() + origin, local.Data(), + static_cast(local.Size(0)) * + sizeof(typename Tensor::value_type), + cudaMemcpyDeviceToHost)); + } +} + +template +void ScatterMgOutput(Tensor &tensor, + const pinned_buffer &host) { + const auto &distribution = tensor.DistributionDescriptor(); + for (size_t fragment = 0; fragment < distribution.FragmentCount(); + ++fragment) { + const auto local_zero = distributed_index_t<1>{0}; + const index_t origin = distribution.LocalToGlobal(fragment, local_zero)[0]; + auto &local = tensor.LocalFragmentForDistributionIndex(fragment).view; + matx::detail::distributed_device_guard guard{ + distribution.FragmentEndpoint(fragment).device_id}; + MATX_CUDA_CHECK(cudaMemcpy(local.Data(), host.Data() + origin, + static_cast(local.Size(0)) * + sizeof(typename Tensor::value_type), + cudaMemcpyHostToDevice)); + } +} + +template +void FftMgImpl(Out &out, const In &in, Executor &executor, FFTNorm norm) { + using value_type = typename remove_cvref_t::value_type; + static_assert( + is_distributed_tensor_v && is_distributed_tensor_v && + remove_cvref_t::Rank() == 1 && remove_cvref_t::Rank() == 1, + "cuFFT multi-GPU currently supports rank-1 distributed tensors"); + static_assert( + std::is_same_v::distribution_type, + block_distribution_t<1>> && + std::is_same_v::distribution_type, + block_distribution_t<1>>, + "cuFFT multi-GPU requires block-distributed tensors"); + static_assert( + std::is_same_v::value_type> && + (std::is_same_v> || + std::is_same_v>), + "cuFFT multi-GPU supports complex and complex"); + static_assert(sizeof(cuda::std::complex) == sizeof(cufftComplex)); + static_assert(sizeof(cuda::std::complex) == + sizeof(cufftDoubleComplex)); + + matx::detail::DistributedCheck( + executor.Context().ProcessCount() == 1, matxNotSupported, + "cuFFT multi-GPU (Mg/Xt) is a single-process backend"); + matx::detail::DistributedCheck( + in.ContextId() == executor.ContextId() && + out.ContextId() == executor.ContextId(), + matxInvalidExecutor, + "cuFFT multi-GPU operands and executor contexts do not match"); + matx::detail::DistributedCheck( + in.Size(0) == out.Size(0) && + in.DistributionDescriptor().Compatible(out.DistributionDescriptor()), + matxInvalidSize, + "cuFFT multi-GPU input and output distributions must match"); + matx::detail::DistributedCheck( + in.LocalFragmentCount() >= 2, matxInvalidSize, + "cuFFT multi-GPU requires at least two local GPUs"); + matx::detail::DistributedCheck( + in.Size(0) <= static_cast(std::numeric_limits::max()), + matxInvalidSize, "cuFFT multi-GPU transform size exceeds its API limit"); + matx::detail::DistributedCheck( + in.LocalFragmentCount() == in.DistributionDescriptor().FragmentCount() && + out.LocalFragmentCount() == + out.DistributionDescriptor().FragmentCount(), + matxInvalidParameter, + "cuFFT multi-GPU requires every fragment to be local to this process"); + + std::vector devices; + devices.reserve(in.LocalFragmentCount()); + for (size_t fragment = 0; + fragment < in.DistributionDescriptor().FragmentCount(); ++fragment) { + const int device = + in.DistributionDescriptor().FragmentEndpoint(fragment).device_id; + matx::detail::DistributedCheck( + std::find(devices.begin(), devices.end(), device) == devices.end(), + matxInvalidParameter, + "cuFFT multi-GPU requires one fragment per distinct CUDA device"); + devices.push_back(device); + matx::detail::DistributedCheck( + in.LocalFragmentForDistributionIndex(fragment).view.IsContiguous() && + out.LocalFragmentForDistributionIndex(fragment).view.IsContiguous(), + matxNotSupported, + "cuFFT multi-GPU requires contiguous local fragments"); + } + + executor.sync(); + cufft_mg_plan plan; + CufftMgCheck(cufftXtSetGPUs(plan.Handle(), static_cast(devices.size()), + devices.data()), + "cufftXtSetGPUs"); + std::vector workspace_sizes(devices.size()); + const cufftType transform_type = + std::is_same_v> ? CUFFT_C2C + : CUFFT_Z2Z; + CufftMgCheck(cufftMakePlan1d(plan.Handle(), static_cast(in.Size(0)), + transform_type, 1, workspace_sizes.data()), + "cufftMakePlan1d"); + CufftMgCheck(cufftXtMalloc(plan.Handle(), plan.InputAddress(), + CUFFT_XT_FORMAT_INPLACE), + "cufftXtMalloc(input)"); + CufftMgCheck(cufftXtMalloc(plan.Handle(), plan.OutputAddress(), + CUFFT_XT_FORMAT_INPLACE), + "cufftXtMalloc(output)"); + + pinned_buffer host_input(static_cast(in.Size(0))); + pinned_buffer host_output(static_cast(out.Size(0))); + GatherMgInput(in, host_input); + CufftMgCheck(cufftXtMemcpy(plan.Handle(), plan.Input(), host_input.Data(), + CUFFT_COPY_HOST_TO_DEVICE), + "cufftXtMemcpy(host-to-device)"); + if constexpr (std::is_same_v>) { + CufftMgCheck( + cufftXtExecDescriptorC2C( + plan.Handle(), plan.Input(), plan.Output(), + Direction == matx::detail::FFTDirection::FORWARD ? CUFFT_FORWARD + : CUFFT_INVERSE), + "cufftXtExecDescriptorC2C"); + } else { + CufftMgCheck( + cufftXtExecDescriptorZ2Z( + plan.Handle(), plan.Input(), plan.Output(), + Direction == matx::detail::FFTDirection::FORWARD ? CUFFT_FORWARD + : CUFFT_INVERSE), + "cufftXtExecDescriptorZ2Z"); + } + CufftMgCheck(cufftXtMemcpy(plan.Handle(), host_output.Data(), plan.Output(), + CUFFT_COPY_DEVICE_TO_HOST), + "cufftXtMemcpy(device-to-host)"); + + double scale = 1.0; + if (norm == FFTNorm::ORTHO) { + scale = 1.0 / std::sqrt(static_cast(in.Size(0))); + } else if ((norm == FFTNorm::FORWARD && + Direction == matx::detail::FFTDirection::FORWARD) || + (norm == FFTNorm::BACKWARD && + Direction == matx::detail::FFTDirection::BACKWARD)) { + scale = 1.0 / static_cast(in.Size(0)); + } + if (scale != 1.0) { + for (size_t i = 0; i < host_output.Size(); ++i) { + host_output.Data()[i] *= + static_cast(scale); + } + } + ScatterMgOutput(out, host_output); +} + +} // namespace detail +} // namespace matx::experimental diff --git a/test/00_tensor/DistributedMpApiTests.cu b/test/00_tensor/DistributedMpApiTests.cu new file mode 100644 index 000000000..8dec20f75 --- /dev/null +++ b/test/00_tensor/DistributedMpApiTests.cu @@ -0,0 +1,27 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +//////////////////////////////////////////////////////////////////////////////// + +#include "matx.h" +#include "matx/distributed.h" +#include "gtest/gtest.h" + +using namespace matx; +using namespace matx::experimental; + +#if defined(MATX_EN_CUBLASMP) || defined(MATX_EN_CUSOLVERMP) +TEST(DistributedMpApi, RejectsNullCommunicatorBeforeBackendSetup) { + int device_count = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&device_count)); + if (device_count == 0) { + GTEST_SKIP() << "This API validation test requires a CUDA device"; + } + + distributed_context context{{0}}; + EXPECT_THROW((distributedCUDAExecutor{context, nullptr, 1, 1}), + matx::detail::matxException); +} +#endif diff --git a/test/00_tensor/DistributedMpIntegrationTests.cu b/test/00_tensor/DistributedMpIntegrationTests.cu new file mode 100644 index 000000000..558c6404d --- /dev/null +++ b/test/00_tensor/DistributedMpIntegrationTests.cu @@ -0,0 +1,242 @@ +//////////////////////////////////////////////////////////////////////////////// +// BSD 3-Clause License +// +// Copyright (c) 2026, NVIDIA Corporation +// All rights reserved. +//////////////////////////////////////////////////////////////////////////////// + +#include "gtest/gtest.h" +#include "matx.h" +#include "matx/distributed.h" + +#include +#include + +#include +#include + +using namespace matx; +using namespace matx::experimental; + +namespace { + +class NcclCommunicator { +public: + NcclCommunicator(int rank, int size) { + ncclUniqueId id{}; + if (rank == 0) { + const ncclResult_t status = ncclGetUniqueId(&id); + if (status != ncclSuccess) { + throw std::runtime_error("ncclGetUniqueId failed"); + } + } + if (MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD) != + MPI_SUCCESS) { + throw std::runtime_error("MPI_Bcast failed"); + } + const ncclResult_t status = + ncclCommInitRank(&communicator_, size, id, rank); + if (status != ncclSuccess) { + throw std::runtime_error("ncclCommInitRank failed"); + } + } + + NcclCommunicator(const NcclCommunicator &) = delete; + NcclCommunicator &operator=(const NcclCommunicator &) = delete; + + ~NcclCommunicator() { + if (communicator_ != nullptr) { + (void)ncclCommDestroy(communicator_); + } + } + + ncclComm_t get() const noexcept { return communicator_; } + +private: + ncclComm_t communicator_ = nullptr; +}; + +struct MpEnvironment { + int rank = 0; + int size = 0; + int device = 0; + int local_size = 0; + int process_rows = 1; + int process_columns = 1; + std::vector endpoints; +}; + +MpEnvironment GetMpEnvironment() { + MpEnvironment environment; + EXPECT_EQ(MPI_Comm_rank(MPI_COMM_WORLD, &environment.rank), MPI_SUCCESS); + EXPECT_EQ(MPI_Comm_size(MPI_COMM_WORLD, &environment.size), MPI_SUCCESS); + + int device_count = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&device_count)); + MPI_Comm local_communicator = MPI_COMM_NULL; + EXPECT_EQ(MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, + environment.rank, MPI_INFO_NULL, + &local_communicator), + MPI_SUCCESS); + int local_rank = 0; + EXPECT_EQ(MPI_Comm_rank(local_communicator, &local_rank), MPI_SUCCESS); + EXPECT_EQ(MPI_Comm_size(local_communicator, &environment.local_size), + MPI_SUCCESS); + environment.device = local_rank; + EXPECT_EQ(MPI_Comm_free(&local_communicator), MPI_SUCCESS); + + environment.process_rows = environment.size % 2 == 0 ? 2 : 1; + environment.process_columns = + environment.size / environment.process_rows; + std::vector devices(static_cast(environment.size)); + EXPECT_EQ(MPI_Allgather(&environment.device, 1, MPI_INT, devices.data(), 1, + MPI_INT, MPI_COMM_WORLD), + MPI_SUCCESS); + environment.endpoints.reserve(devices.size()); + for (int rank = 0; rank < environment.size; ++rank) { + environment.endpoints.push_back( + {rank, devices[static_cast(rank)]}); + } + return environment; +} + +template +void FillLocalMatrix(Tensor &tensor, Generator &&generator) { + ASSERT_EQ(tensor.LocalFragmentCount(), 1U); + const auto &fragment = tensor.LocalFragment(0); + const auto &distribution = tensor.DistributionDescriptor(); + const auto shape = distribution.LocalShape(fragment.distribution_index); + std::vector host(static_cast(shape[0] * shape[1])); + for (index_t row = 0; row < shape[0]; ++row) { + for (index_t column = 0; column < shape[1]; ++column) { + const auto global = distribution.LocalToGlobal( + fragment.distribution_index, {row, column}); + host[static_cast(row * shape[1] + column)] = + generator(global[0], global[1]); + } + } + MATX_CUDA_CHECK(cudaMemcpy(tensor.LocalView(0).Data(), host.data(), + host.size() * sizeof(float), + cudaMemcpyHostToDevice)); +} + +template +void VerifyLocalMatrix(const Tensor &tensor, Verifier &&verifier) { + ASSERT_EQ(tensor.LocalFragmentCount(), 1U); + const auto &fragment = tensor.LocalFragment(0); + const auto &distribution = tensor.DistributionDescriptor(); + const auto shape = distribution.LocalShape(fragment.distribution_index); + std::vector host(static_cast(shape[0] * shape[1])); + MATX_CUDA_CHECK(cudaMemcpy(host.data(), tensor.LocalView(0).Data(), + host.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + for (index_t row = 0; row < shape[0]; ++row) { + for (index_t column = 0; column < shape[1]; ++column) { + const auto global = distribution.LocalToGlobal( + fragment.distribution_index, {row, column}); + verifier(global[0], global[1], + host[static_cast(row * shape[1] + column)]); + } + } +} + +} // namespace + +#ifdef MATX_EN_CUBLASMP +TEST(DistributedMpIntegration, CublasMpMatmul) { + const auto environment = GetMpEnvironment(); + int device_count = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&device_count)); + if (environment.size < 2 || device_count < environment.local_size) { + GTEST_SKIP() << "Requires at least two MPI ranks and one GPU per rank"; + } + + MATX_CUDA_CHECK(cudaSetDevice(environment.device)); + NcclCommunicator communicator{environment.rank, environment.size}; + distributed_context context{{environment.device}, environment.rank, + environment.size}; + distributedCUDAExecutor executor{ + context, communicator.get(), environment.process_rows, + environment.process_columns}; + + constexpr index_t matrix_size = 16; + block_cyclic_distribution_t distribution{ + {matrix_size, matrix_size}, + {2, 2}, + {environment.process_rows, environment.process_columns}, + environment.endpoints}; + auto a = make_distributed_tensor(distribution, context); + auto b = make_distributed_tensor(distribution, context); + auto c = make_distributed_tensor(distribution, context); + + FillLocalMatrix(a, [](index_t row, index_t column) { + return row == column ? static_cast(row + 1) : 0.0F; + }); + FillLocalMatrix(b, [](index_t row, index_t column) { + return static_cast(row * matrix_size + column + 1); + }); + FillLocalMatrix(c, [](index_t, index_t) { return 0.0F; }); + + (c = matmul(a, b)).run(executor); + executor.sync(); + + VerifyLocalMatrix(c, [](index_t row, index_t column, float value) { + const float expected = + static_cast((row + 1) * (row * matrix_size + column + 1)); + EXPECT_NEAR(value, expected, 2.0e-3F); + }); +} +#endif + +#ifdef MATX_EN_CUSOLVERMP +TEST(DistributedMpIntegration, CusolverMpCholesky) { + const auto environment = GetMpEnvironment(); + int device_count = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&device_count)); + if (environment.size < 2 || device_count < environment.local_size) { + GTEST_SKIP() << "Requires at least two MPI ranks and one GPU per rank"; + } + + MATX_CUDA_CHECK(cudaSetDevice(environment.device)); + NcclCommunicator communicator{environment.rank, environment.size}; + distributed_context context{{environment.device}, environment.rank, + environment.size}; + distributedCUDAExecutor executor{ + context, communicator.get(), environment.process_rows, + environment.process_columns}; + + constexpr index_t matrix_size = 16; + block_cyclic_distribution_t distribution{ + {matrix_size, matrix_size}, + {2, 2}, + {environment.process_rows, environment.process_columns}, + environment.endpoints}; + auto a = make_distributed_tensor(distribution, context); + auto output = make_distributed_tensor(distribution, context); + + FillLocalMatrix(a, [](index_t row, index_t column) { + const float diagonal = static_cast(row + 2); + return row == column ? diagonal * diagonal : 0.0F; + }); + + (output = chol(a, SolverFillMode::LOWER)).run(executor); + executor.sync(); + + VerifyLocalMatrix(output, [](index_t row, index_t column, float value) { + if (row >= column) { + const float expected = + row == column ? static_cast(row + 2) : 0.0F; + EXPECT_NEAR(value, expected, 2.0e-3F); + } + }); +} +#endif + +int main(int argc, char **argv) { + MPI_Init(&argc, &argv); + ::testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + matx::ClearCachesAndAllocations(); + MPI_Finalize(); + return result; +} diff --git a/test/00_tensor/DistributedTensorMultiGpuTests.cu b/test/00_tensor/DistributedTensorMultiGpuTests.cu index 6adaf8f1a..0962a3e62 100644 --- a/test/00_tensor/DistributedTensorMultiGpuTests.cu +++ b/test/00_tensor/DistributedTensorMultiGpuTests.cu @@ -32,6 +32,7 @@ //////////////////////////////////////////////////////////////////////////////// #include "matx.h" +#include "matx/distributed.h" #include "gtest/gtest.h" #include @@ -89,6 +90,53 @@ void FillMultiGpu(Distributed &tensor, Generator &&generator) { } // namespace +TEST(DistributedTensorMultiGpu, CufftMgOneDimensionalComplex) { + int device_count = 0; + MATX_CUDA_CHECK(cudaGetDeviceCount(&device_count)); + if (device_count < 2) { + GTEST_SKIP() << "This integration test requires two CUDA devices"; + } + + constexpr index_t count = 1024; + using complex_type = cuda::std::complex; + distributed_context context{{0, 1}}; + distributedCUDAExecutor executor{context}; + auto distribution = + block_distribution_t<1>::Slab({count}, {{0, 0}, {0, 1}}); + auto input = make_distributed_tensor(distribution, context); + auto output = make_distributed_tensor(distribution, context); + + for (size_t local = 0; local < input.LocalFragmentCount(); ++local) { + const auto &fragment = input.LocalFragment(local); + const index_t local_size = input.LocalView(local).Size(0); + std::vector host(static_cast(local_size), + complex_type{0.0F, 0.0F}); + if (distribution.LocalToGlobal(fragment.distribution_index, {0})[0] == 0) { + host[0] = complex_type{1.0F, 0.0F}; + } + matx::detail::distributed_device_guard guard{fragment.endpoint.device_id}; + MATX_CUDA_CHECK(cudaMemcpy(input.LocalView(local).Data(), host.data(), + host.size() * sizeof(complex_type), + cudaMemcpyHostToDevice)); + } + + (output = fft(input)).run(executor); + + for (size_t local = 0; local < output.LocalFragmentCount(); ++local) { + const auto &fragment = output.LocalFragment(local); + const index_t local_size = output.LocalView(local).Size(0); + std::vector host(static_cast(local_size)); + matx::detail::distributed_device_guard guard{fragment.endpoint.device_id}; + MATX_CUDA_CHECK(cudaMemcpy(host.data(), output.LocalView(local).Data(), + host.size() * sizeof(complex_type), + cudaMemcpyDeviceToHost)); + for (index_t i = 0; i < local_size; ++i) { + EXPECT_NEAR(host[static_cast(i)].real(), 1.0F, 1.0e-5F); + EXPECT_NEAR(host[static_cast(i)].imag(), 0.0F, 1.0e-5F); + } + } +} + TEST(DistributedTensorMultiGpu, BlackScholesUnevenFragments) { int device_count = 0; MATX_CUDA_CHECK(cudaGetDeviceCount(&device_count)); diff --git a/test/00_tensor/DistributedTensorTests.cu b/test/00_tensor/DistributedTensorTests.cu index f3c59514a..860428872 100644 --- a/test/00_tensor/DistributedTensorTests.cu +++ b/test/00_tensor/DistributedTensorTests.cu @@ -235,6 +235,36 @@ TEST(DistributedTensor, BlockCyclicMapping) { EXPECT_EQ(distribution.LocalToGlobal(3, {0, 0})[1], 2); EXPECT_EQ(distribution.LocalToGlobal(0, {2, 2})[0], 4); EXPECT_EQ(distribution.LocalToGlobal(0, {2, 2})[1], 4); + + block_cyclic_distribution_t column_major{ + {8, 12}, + {2, 2}, + {2, 3}, + {{0, 0}, {1, 0}, {2, 0}, {3, 0}, {4, 0}, {5, 0}}, + distributed_grid_layout::column_major}; + // Communicator rank 1 maps to process coordinate (1, 0), not (0, 1). + EXPECT_EQ(column_major.LocalToGlobal(1, {0, 0})[0], 2); + EXPECT_EQ(column_major.LocalToGlobal(1, {0, 0})[1], 0); +} + +TEST(DistributedTensor, RegularOperatorsSelectBlockCyclicBackend) { + distributed_context context{{0}}; + distributedCUDAExecutor executor{context}; + block_cyclic_distribution_t distribution{ + {4, 4}, {2, 2}, {1, 1}, {{0, 0}}}; + auto a = make_distributed_tensor(distribution, context); + auto b = make_distributed_tensor(distribution, context); + auto output = make_distributed_tensor(distribution, context); + + auto matmul_op = matmul(a, b); + auto chol_op = chol(a); + static_assert(is_distributed_expression_v); + static_assert(is_distributed_expression_v); + + EXPECT_THROW((output = matmul_op).run(executor), + matx::detail::matxException); + EXPECT_THROW((output = chol_op).run(executor), + matx::detail::matxException); } TEST(DistributedTensor, PointwiseCopyAndMaterialize) { diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 33f1d649c..3c69eaf6c 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -26,6 +26,7 @@ set (test_sources 00_tensor/DynamicTensorTests.cu 00_tensor/DistributedTensorTests.cu 00_tensor/DistributedTensorMultiGpuTests.cu + 00_tensor/DistributedMpApiTests.cu ${OPERATOR_TEST_FILES} 00_operators/GeneratorTests.cu 00_operators/PWelch.cu @@ -149,6 +150,39 @@ foreach(test_file ${test_sources}) create_test_executable(${test_file}) endforeach() +if(MATX_EN_CUBLASMP OR MATX_EN_CUSOLVERMP) + find_package(MPI REQUIRED COMPONENTS CXX) + add_executable( + test_00_tensor_DistributedMpIntegrationTests + 00_tensor/DistributedMpIntegrationTests.cu + ) + target_compile_options( + test_00_tensor_DistributedMpIntegrationTests + PRIVATE $<$:${MATX_CUDA_FLAGS}> + ) + target_include_directories( + test_00_tensor_DistributedMpIntegrationTests + PRIVATE "${target_inc}" + SYSTEM PRIVATE "${system_inc}" + ) + target_link_libraries( + test_00_tensor_DistributedMpIntegrationTests + PRIVATE matx::matx gtest MPI::MPI_CXX + ) + add_test( + NAME test_00_tensor_DistributedMpIntegrationTests + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 2 + ${MPIEXEC_PREFLAGS} + $ + ${MPIEXEC_POSTFLAGS} + ) + set_tests_properties( + test_00_tensor_DistributedMpIntegrationTests + PROPERTIES LABELS "multi_gpu;multi_node" RUN_SERIAL TRUE + ) +endif() + # Create individual executables for proprietary tests foreach(test_file ${proprietary_sources}) create_test_executable(${test_file})