From 83bb38fff43a4712be88a7d22ec8f561eb8e7552 Mon Sep 17 00:00:00 2001 From: Marius Eriksen Date: Thu, 10 Jul 2025 14:09:31 -0700 Subject: [PATCH] [hyperactor] mesh: combine CAST_RANK and CAST_SHAPE into a single header ... as these must always be set together; thus their usage is correct by construction. Differential Revision: [D78119117](https://our.internmc.facebook.com/intern/diff/D78119117/) **NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D78119117/)! [ghstack-poisoned] --- hyperactor_mesh/src/comm/multicast.rs | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/hyperactor_mesh/src/comm/multicast.rs b/hyperactor_mesh/src/comm/multicast.rs index 09f19165..1ab8f16d 100644 --- a/hyperactor_mesh/src/comm/multicast.rs +++ b/hyperactor_mesh/src/comm/multicast.rs @@ -179,19 +179,20 @@ pub(crate) struct ForwardMessage { } declare_attrs! { - /// Used inside headers for cast messages to store - /// the rank of the receiver. - attr CAST_RANK: usize; - /// Used inside headers to store the shape of the - /// actor mesh that a message was cast to. - attr CAST_SHAPE: Shape; + /// Casting metadata comprising: + /// 1) the rank of the receiver in the mesh, + /// 2) the shape of the mesh on which the message is being cast. + /// + /// The mesh here is the root mesh -- i.e., the shape is the extents + /// of the mesh on which the actors were spawned. + attr CAST_INFO: (usize, Shape); + /// Used inside headers to store the originating sender of a cast. pub attr CAST_ORIGINATING_SENDER: ActorId; } pub fn set_cast_info_on_headers(headers: &mut Attrs, rank: usize, shape: Shape, sender: ActorId) { - headers.set(CAST_RANK, rank); - headers.set(CAST_SHAPE, shape); + headers.set(CAST_INFO, (rank, shape)); headers.set(CAST_ORIGINATING_SENDER, sender); } @@ -206,10 +207,9 @@ pub trait CastInfo { impl CastInfo for Context<'_, A> { fn cast_info(&self) -> (usize, Shape) { let headers = self.headers(); - match (headers.get(CAST_RANK), headers.get(CAST_SHAPE)) { - (Some(rank), Some(shape)) => (*rank, shape.clone()), - (None, None) => (0, Shape::unity()), - _ => panic!("Expected either both rank and shape or neither"), + match headers.get(CAST_INFO) { + Some((rank, shape)) => (*rank, shape.clone()), + None => (0, Shape::unity()), } } }