Expose backend type to python

[samnordmann]

On top of

• [Multidevice] Add backend type to Communication and expose to API #3914
• Expose MultiDeviceExecutor and overlapped AG+matmul to python API #3923

[github-actions]

Description

• Expose backend type to Python API

• Add backend type to Communication and Communication::toString

• Modify HostIrLower to set backend type for Communications

• Update MultiDeviceExecutor to accept and use MultiDeviceExecutorParams

---

Changes walkthrough 📝

Relevant files

Enhancement

12 files executor.cpp Modify HostIrLower instantiation in HostIrExecutor              +4/-2      lower.cpp Set backend type for Communications in HostIrLower              +9/-1      lower.h Add HostIrLowerParams with communicator_backend                    +11/-3    communication.cpp Add backend type to Communication constructor and toString +8/-5      communication.h Add backend type to Communication class                                    +10/-1    executor.cpp Modify MultiDeviceExecutor to use MultiDeviceExecutorParams +5/-4      executor.h Add MultiDeviceExecutorParams struct                                          +7/-1      fusion_definition.cpp Pass backend_type to MultiDeviceExecutor in FusionDefinition +4/-2      fusion_definition.h Add backend_type to FusionDefinition                                          +1/-0      python_bindings.cpp Add CommunicatorBackend enum and set_backend method to FusionDefinition +16/-3    test_multidevice_pipeline.cpp Update MultiDeviceExecutor instantiation in tests                +3/-4      test_multidevice.py Update FusionDefinition instantiation with communication_backend in tests +3/-2

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Code Duplication The new code introduces a new instance of HostIrLower within the loop, which might lead to performance issues or unintended behavior. Consider reusing the HostIrLower instance outside the loop. HostIrLower lower; std::vector<Expr*> communications = lower.lower(cloner.clone(e)); Default Backend The default backend is set to kNccl in the constructor. Ensure that this is the intended behavior and that it does not lead to unexpected behavior when no backend is explicitly set. CommunicatorBackend backend) : Expr(passkey) { Backend Setting The set_backend method in FusionDefinition uses a static boolean backend_set to prevent setting the backend more than once. Ensure that this mechanism is appropriate and does not lead to issues in multi-threaded environments. [](FusionDefinition& self, CommunicatorBackend backend_type) { static bool backend_set = false; NVF_CHECK(!backend_set, "Backend type has already been set"); backend_set = true; self.backend_type = backend_type; });