iframe-proxy

Tixxx · 2025-01-28T18:57:58Z

This pr adds runtime support for nvshmem reduction on-stream api.
It adds a new backend config to instruction the emitter to lower to nvshmem with corresponding collective thunks that call nvshmem reduction api.

google-cla · 2025-01-28T18:58:03Z

jprabhas · 2025-01-29T20:58:50Z

+    is_nvshmem_collective = backend_config.backend() == CollectiveBackendConfig::NVSHMEM;
+  }
+  return (alias->instruction()->opcode() == HloOpcode::kCustomCall &&
+      alias->instruction()->custom_call_target() == "mosaic_gpu") || is_nvshmem_collective;


why are mosaic_gpu custom calls in the same category as nvshmem ops?

This is from the allocator branch. We use nvshmem allocator for mosaic gpu calls or native nvshmem collectives. I think instead of using mosaic_gpu custom call, the pallas kernel will use a more specific custom call name so this will change after it's changed in the other pr and rebase on that.

PatriosTheGreat · 2025-02-17T10:54:50Z

+                                      num_elements);
+      break;
+    }
+    case PrimitiveType::F16: {


Should F32 also be here?

I tried to add F32 here and it leads to the deadlock. Apparently this reproduces only with NVSHMEM_TEAM_WORLD, when I'm passing NVSHMEMX_TEAM_NODE everything works fine. Not sure what's the under hood issue.

I tried reproing with multihost runner and setting ncshmem api to use world team but for 4 ranks it went through. Do you have a repro that i can run?

I rebased this PR at XLA head (if it would be useful I can share a patch or rebased PR). And I don't see an F32 issue at least using nvcc and compiled with cmake nvshmem. I'll investigate if this is an internal issue.

Added more fixes on top of base branch, needs one last rebase consolidate nccl and nvshmem helper functions to avoid clash

move barrier to init

PatriosTheGreat

Could you please split this PR to 4 smaller PRs? This would simplify review process and potential rollbacks-relands

The following PRs can be extracted

Nvshmem.symbols changes
Adding nvshmem backend_config property and runtime usage (compile_module_to_llvm_ir.cc/memory assignment changes).
Nvshmem runtime implementation (nvshmem_communicator with all_reduce logic)
GPU emitter NVSHMEM changes (the nvshmem thunks emitting logic)

PatriosTheGreat · 2025-04-28T13:05:05Z

+namespace xla {
+namespace gpu {
+
+absl::Status RunNvshmemAllReduce(GpuCollectives* collectives,


Could you please implement this similar to how currently nccl all reduce is implemented?
For this nvshmem_communicator should be created similar to nccl_communicator and all collective operations should be implemented there. This should be the only target which depends on nvshmem internal libraries.

Imported from GitHub PR #23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from #21973. Copybara import of the project: -- 08deddb by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82 by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fe by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45 by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f674569 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adb by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebb by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d06383 by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03 by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=#23829 from sfvaroglu:sevin/nvshmem_backend f682e03 PiperOrigin-RevId: 764806040

Imported from GitHub PR openxla/xla#23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from openxla/xla#21973. Copybara import of the project: -- 08deddb9283d20fbc1c53bc6f5e19ca32f562154 by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a4f77b3eefe03733dee9ccfd1aabc31f47 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82e0ab2ab2ed78664dbd6cedcb1356e220b by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada179952d9b84a143a6c1d16200eedb96c9 by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fefacb0ddd1810435ea11baa9f09d040a37 by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45e1f7d0c7131762ec47f6dce790b60b50f by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f6745690b6cbd82f48486ef58b43165d9cb926e1 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a765bb35fb77c762806bf51f36f8031977 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adbb1879407854b7027d484ffdbfcbd5844d by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21539fb2329d080195d613da390abf30a78 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a4daed393221e1aa6deab0b347de57d64e by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebbd98594c2bb9daa791534f12afb55dc62d by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d0638349447c0670fe9abaa507b6706302f597f by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03bd790f337ab28c82bf805f1443b38683c by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#23829 from sfvaroglu:sevin/nvshmem_backend f682e03bd790f337ab28c82bf805f1443b38683c PiperOrigin-RevId: 764806040

Imported from GitHub PR #23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from #21973. Copybara import of the project: -- 08deddb by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82 by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fe by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45 by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f674569 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adb by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebb by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d06383 by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03 by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=#23829 from sfvaroglu:sevin/nvshmem_backend f682e03 PiperOrigin-RevId: 764806040

Imported from GitHub PR #23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from #21973. Copybara import of the project: -- 08deddb by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82 by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fe by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45 by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f674569 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adb by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebb by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d06383 by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03 by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test -- ae7b51c by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla/pjrt/gpu:gpu_topology Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=#23829 from sfvaroglu:sevin/nvshmem_backend ae7b51c PiperOrigin-RevId: 764806040

Imported from GitHub PR openxla/xla#23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from openxla/xla#21973. Copybara import of the project: -- 08deddb9283d20fbc1c53bc6f5e19ca32f562154 by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a4f77b3eefe03733dee9ccfd1aabc31f47 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82e0ab2ab2ed78664dbd6cedcb1356e220b by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada179952d9b84a143a6c1d16200eedb96c9 by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fefacb0ddd1810435ea11baa9f09d040a37 by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45e1f7d0c7131762ec47f6dce790b60b50f by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f6745690b6cbd82f48486ef58b43165d9cb926e1 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a765bb35fb77c762806bf51f36f8031977 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adbb1879407854b7027d484ffdbfcbd5844d by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21539fb2329d080195d613da390abf30a78 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a4daed393221e1aa6deab0b347de57d64e by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebbd98594c2bb9daa791534f12afb55dc62d by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d0638349447c0670fe9abaa507b6706302f597f by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03bd790f337ab28c82bf805f1443b38683c by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test -- ae7b51cf853db44553cbef0710e0791a25af8534 by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla/pjrt/gpu:gpu_topology Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#23829 from sfvaroglu:sevin/nvshmem_backend ae7b51cf853db44553cbef0710e0791a25af8534 PiperOrigin-RevId: 764806040

Imported from GitHub PR #23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from #21973. Copybara import of the project: -- 08deddb by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82 by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fe by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45 by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f674569 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adb by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebb by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d06383 by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03 by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test -- ae7b51c by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla/pjrt/gpu:gpu_topology Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=#23829 from sfvaroglu:sevin/nvshmem_backend ae7b51c PiperOrigin-RevId: 767677333

Imported from GitHub PR openxla/xla#23829 This PR adds `CollectiveBackendAssigner` pass, which assigns either NVSHMEM or NCCL as the backend for collective operations based on: 1. Communication pattern (intranode vs internode) 2. Message size (compared against threshold_in_bytes) Also, there is some utility code taken from openxla/xla#21973. Copybara import of the project: -- 08deddb9283d20fbc1c53bc6f5e19ca32f562154 by Sevin Varoglu <svaroglu@nvidia.com>: Add CollectiveBackendAssigner pass for NVSHMEM vs. NCCL -- d505d8a4f77b3eefe03733dee9ccfd1aabc31f47 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 8341b82e0ab2ab2ed78664dbd6cedcb1356e220b by Sevin Varoglu <svaroglu@nvidia.com>: Remove NvshmemCollectives-based device/process logic -- 5f2aada179952d9b84a143a6c1d16200eedb96c9 by Sevin Varoglu <svaroglu@nvidia.com>: Mark gpu_compiler.cc as cuda-only -- 6b7e6fefacb0ddd1810435ea11baa9f09d040a37 by Sevin Varoglu <svaroglu@nvidia.com>: Remove cuda-only -- 69bdb45e1f7d0c7131762ec47f6dce790b60b50f by Sevin Varoglu <svaroglu@nvidia.com>: Remove nvshmem_collectives from CollectiveBackendAssigner -- f6745690b6cbd82f48486ef58b43165d9cb926e1 by Sevin Varoglu <svaroglu@nvidia.com>: Remove unnecessary params -- 7cb298a765bb35fb77c762806bf51f36f8031977 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 5c39adbb1879407854b7027d484ffdbfcbd5844d by Sevin Varoglu <svaroglu@nvidia.com>: Add tags -- d31af21539fb2329d080195d613da390abf30a78 by Sevin Varoglu <svaroglu@nvidia.com>: Incorporate review feedback -- 1448d6a4daed393221e1aa6deab0b347de57d64e by Sevin Varoglu <svaroglu@nvidia.com>: Add unit tests -- 1f4bebbd98594c2bb9daa791534f12afb55dc62d by Sevin Varoglu <svaroglu@nvidia.com>: Add CM cases to collective_ops_utils_test.cc -- 5d0638349447c0670fe9abaa507b6706302f597f by Sevin Varoglu <svaroglu@nvidia.com>: Remove e2e test -- f682e03bd790f337ab28c82bf805f1443b38683c by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla_test -- ae7b51cf853db44553cbef0710e0791a25af8534 by Sevin Varoglu <svaroglu@nvidia.com>: Remove xla/pjrt/gpu:gpu_topology Merging this change closes #23829 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#23829 from sfvaroglu:sevin/nvshmem_backend ae7b51cf853db44553cbef0710e0791a25af8534 PiperOrigin-RevId: 767677333