OpenMPI + CUDA Compatiblity

[JuanPedroGHM]

Introduces an incompatibility list for different communication backends, in order to safely enable GPU compatible MPI communication.

Changes

• Incompatibility list in configuration file.
• Incompatibility check in _moveToCompDevice

Does this affect performance

• Some operations (Bcast and Allgatherv in resplit_(None)) have shown massive slowdowns when running the tests. We might want to add certain functions to the incompatibility list for performance reasons.

[codecov]

Codecov Report

❌ Patch coverage is 61.22449% with 38 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.73%. Comparing base (5d42619) to head (bf8f590).
⚠️ Report is 5 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
heat/core/_config.py	48.48%	34 Missing ⚠️
heat/core/communication.py	87.50%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2279      +/-   ##
==========================================
- Coverage   91.74%   91.73%   -0.02%     
==========================================
  Files          87       87              
  Lines       14120    14148      +28     
==========================================
+ Hits        12955    12979      +24     
- Misses       1165     1169       +4     

Flag	Coverage Δ
unit	91.73% <61.22%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[brownbaerchen]

I have a bunch of small comments about clarifying stuff. I don't know how MPI datatypes work and therefore don't understand some of the changes.

Most changes look really good, though!

[brownbaerchen]

Can we rename this variable to something more expressive? Maybe cuda_is_available?

[brownbaerchen]

rocm_is_available?

[JuanPedroGHM]

Benchmarks results - Sponsored by perun

FTPx86-AMD-Milan-Mi250 - CPU - N4 - RUNTIME

Function	Value [s]	Change %	30 day avg. [s]	Last entrie [s]
lanczos	0.323632	-30.2891	0.464249	0.464249
randomized_svd_split0	0.0389764	74.4738	0.0223394	0.0223394
randomized_svd_split1	0.0412789	152.521	0.0163467	0.0163467
reshape	0.19887	-15.9597	0.236636	0.236636
concatenate	0.17693	21.3462	0.145806	0.145806
apply_inplace_standard_scaler_and_inverse	0.0353156	22.3109	0.0288736	0.0288736
apply_inplace_normalizer	0.00812416	16.6639	0.00696373	0.00696373
lasso_predict	0.00379181	-48.0994	0.00730591	0.00730591

---

FTPx86-AMD-Milan-Mi250 - GPU - N4 - RUNTIME

Function	Value [s]	Change %	30 day avg. [s]	Last entrie [s]
matmul_split_1	0.0300992	-41.1512	0.0511468	0.0511468
qr_split_0	0.280176	18.4659	0.236503	0.236503
lanczos	1.61933	-37.3592	2.5851	2.5851
svd_zolo_split1	2.85499	-10.8745	3.20333	3.20333
randomized_svd_split0	0.0348533	-26.5707	0.047465	0.047465
randomized_svd_split1	0.0460559	-34.7792	0.0706153	0.0706153
randomized_eigh_split0	0.032163	-49.2933	0.0634295	0.0634295
randomized_eigh_split1	0.0600281	-28.4496	0.0838963	0.0838963
spectralclustering	0.121409	-17.1642	0.146566	0.146566
concatenate	0.142342	22.6735	0.116033	0.116033
apply_inplace_min_max_scaler_and_inverse	0.00496435	106.012	0.00240974	0.00240974
apply_inplace_max_abs_scaler_and_inverse	0.00301123	161.08	0.00115337	0.00115337
apply_inplace_normalizer	0.0141697	55.9175	0.00908794	0.00908794
incremental_pca_split0	7.59182	-15.4847	8.98278	8.98278
convolution_array_distributed	0.0178364	92.7011	0.00925598	0.00925598
convolution_batch_processing	0.00818806	166.656	0.00307064	0.00307064
convolution2d_batch_processing	0.00122375	-11.4546	0.00138206	0.00138206
lasso_predict	0.00415668	48.961	0.00279045	0.00279045

---

Grafana Dashboard
Last updated: 2026-07-01T13:06:46Z

[JuanPedroGHM]

I've addressed the comments, and for now disabled rocm buffer compatibility with MPI. Should be addressed later in a different issue/pr. Should be ready to merge if after a review and if all tests pass.

[brownbaerchen]

I still think the incompatibilities determination is difficult to maintain because it's a bit confusing.. I would be open to merging this anyways if you don't have bandwidth to fix this, though. GPU-aware MPI compatibility will be very nice!

[brownbaerchen]

Hm? This deactivates GPU computation by default, yes? If no incompatibilities have been recorded, the default is device_incompatibilities=None, right? Maybe that's what we want, though..

[JuanPedroGHM]

GPU compatibility is set to true if the library and environment say indicate that it is compatible (from ompi_info flags, for example), and if the list of incompatiblities is a list. If it is False or any other object, it will disable gpu support.