Distributed tensors with transforms

[cliffburdick]

Adds support for distributed tensor with math libraries that support it. cuFFTMg, cuBLASMp, and cuSOLVERMp are supported. Each of these libraries has different layout requirements that are outlined in the documentation. The layout requirements may chance with library updates, and the underlying tensor type should still be flexible enough to allow for this. For the initial commit the following transforms are supported: fft, matmul, cholesky. Ohter transforms may be added in the future.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cliffburdick]

/build

[greptile-apps]

Greptile Summary

Adds block-cyclic distributed matmul and Cholesky through cuBLASMp/cuSOLVERMp, plus single-process multi-GPU FFT execution through cuFFT Xt/Mg.

• Extends distributed layouts and executors with process-grid and collective-backend state.
• Adds dependency discovery, documentation, examples, and distributed integration tests.
• Introduces MP transform adapters that pack ordinary local views into backend-compatible column-major buffers.

Confidence Score: 4/5

The PR should not merge until MP transform temporaries are made exception-safe, because recoverable backend failures currently leak GPU memory.

Both new MP adapters allocate raw stream-ordered device buffers before checked backend calls, while their release logic is reachable only after those calls succeed.

Files Needing Attention: include/matx/transforms/distributed/distributed_mp.h

Important Files Changed

Filename	Overview
include/matx/transforms/distributed/distributed_mp.h	Adds cuBLASMp and cuSOLVERMp adapters, but temporary device allocations leak when a checked backend operation throws.
include/matx/transforms/distributed/fft_mg.h	Adds rank-1 complex cuFFT Xt/Mg execution with validated local fragments, natural-order host staging, and normalization.
include/matx/executors/distributed.h	Extends the distributed executor with borrowed NCCL communicator state and MP backend handles tied to its local stream.
include/matx/core/distributed_tensor.h	Adds row-major and column-major process-grid mapping to block-cyclic distributions.
include/matx/operators/fft.h	Selects local or cuFFT Mg execution based on whether the distributed FFT dimension is partitioned.
include/matx/operators/matmul.h	Routes rank-2 block-cyclic operands to the cuBLASMp distributed expression.
include/matx/operators/chol.h	Routes rank-2 block-cyclic inputs to the cuSOLVERMp distributed expression.
CMakeLists.txt	Adds opt-in discovery, compile definitions, and transitive links for distributed NVIDIA math backends.
test/CMakeLists.txt	Registers API tests and a two-rank MPI integration executable when MP backends are enabled.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Distributed tensor expression] --> B{Input distribution}
  B -->|Block-cyclic rank 2| C[Distributed MP adapter]
  C --> D[Pack local row-major fragment]
  D --> E[cuBLASMp or cuSOLVERMp collective]
  E --> F[Unpack local result]
  B -->|Partitioned rank-1 complex| G[cuFFT Xt/Mg adapter]
  G --> H[Gather natural-order host buffer]
  H --> I[Multi-GPU FFT]
  I --> J[Scatter distributed output]
  B -->|Operation dimensions local| K[Existing per-fragment transform]

Loading

_{Reviews (1): Last reviewed commit: "Add configurable multi-process Cholesky ..." | Re-trigger Greptile}

[greptile-apps]

Backend errors leak GPU buffers

When a cuBLASMp or cuSOLVERMp buffer-size or execution call fails, the checked call throws before the success-only cleanup runs, leaking the packed buffers and any allocated workspace. Repeated recoverable failures can therefore exhaust GPU memory and cause subsequent operations to fail.

Knowledge Base Used: Transforms