[RFC] Add source_target_pairs attribute to send and recv ops in stableHLO #2784
Conversation
rfcs/20250421-source-target-pairs.md
Outdated
| %results0, %results1 = "stablehlo.recv"(%token) { | ||
| source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>, | ||
| channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>, | ||
| is_host_transfer = true |
There was a problem hiding this comment.
Can we keep it false to make it clear that the source_target_pairs will have some effect in this particular example.
rfcs/20250421-source-target-pairs.md
Outdated
| #### Semantics | ||
|
|
||
| Sends `inputs` to a channel `channel_id`. Inputs are then sent to other devices | ||
| in the order specified by `source_target_pairs`. The operation produces a |
There was a problem hiding this comment.
Similar to recv can we defer the mention of in the order specified by source_target_pairs to the second paragraph (as it is now), unless we somehow capture the event of is_host_transfer being true in the same paragraph.
My goal here to keep the wording consistent with recv.
|
@GleasonK is this RFC good to go now? |
| proposed in conjunction with exposing send/recv operations through the JAX | ||
| `shard_map` API. | ||
|
|
||
| ## Proposed Specification |
There was a problem hiding this comment.
Can you add a section for "A Note on Backward Compatibility", denoting that we're technically making the semantics more strict for send given that any instances of send is_host_transfer=false that are serialized will no longer be deserializable, but this is likely OK since this would have been undefined behavior as it is.
There was a problem hiding this comment.
cc @frgossen to confirm this as well
No description provided.