monarch_rdma/backend/ibverbs/memory_region.rs
1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Registered memory regions returned by [`IbvDomainImpl::register_mr`].
10//!
11//! [`IbvMemoryRegionView`] is the cheap, cloneable handle this process addresses
12//! its own registered memory through: the keys and addresses for a slice of
13//! registered memory, plus an `Arc<dyn IbvMemoryRegionKeepalive>` that keeps the
14//! backing registration's resources alive until the last clone of the view
15//! drops.
16//!
17//! [`IbvRemoteMemoryRegionView`] is what a peer gets instead: the same region
18//! reduced to what the wire can carry and the far side can use.
19//!
20//! [`IbvDomainImpl::register_mr`]: super::domain::IbvDomainImpl::register_mr
21
22use std::sync::Arc;
23
24use serde::Deserialize;
25use serde::Serialize;
26use typeuri::Named;
27
28use super::primitives::IbvMr;
29
30/// Guards the resources behind a registered MR, releasing them when the last
31/// [`IbvMemoryRegionView`] over it drops. Each implementor frees whatever it
32/// owns in its own `Drop`; the trait carries no methods and exists only to
33/// type-erase the guards so a view can hold any of them behind an
34/// `Arc<dyn IbvMemoryRegionKeepalive>`.
35pub(super) trait IbvMemoryRegionKeepalive: std::fmt::Debug + Send + Sync {}
36
37/// A standalone [`IbvMr`] guards its own registration: its `Drop` runs
38/// `ibv_dereg_mr` against the PD it owns.
39impl IbvMemoryRegionKeepalive for IbvMr {}
40
41/// A cloneable handle to a slice of registered memory: the keys and addresses
42/// a peer needs, plus an `Arc<dyn IbvMemoryRegionKeepalive>` keepalive.
43///
44/// Cheap to clone; every clone shares the same guard, so the backing
45/// registration stays alive (and registered) until the last clone drops.
46#[derive(Debug, Clone)]
47pub struct IbvMemoryRegionView {
48 /// Virtual address in the local process address space.
49 pub virtual_addr: usize,
50 /// RDMA address, possibly offset from the region's base MR address.
51 pub rdma_addr: usize,
52 pub size: usize,
53 pub lkey: u32,
54 pub rkey: u32,
55 /// Name of the RDMA device the view's protection domain is on.
56 pub device_name: String,
57 /// Keeps the backing registration alive for every clone of this view; the
58 /// last drop releases its resources. Never read directly.
59 pub(super) _guard: Arc<dyn IbvMemoryRegionKeepalive>,
60}
61
62impl IbvMemoryRegionView {
63 pub(super) fn new(
64 virtual_addr: usize,
65 rdma_addr: usize,
66 size: usize,
67 lkey: u32,
68 rkey: u32,
69 device_name: String,
70 guard: Arc<dyn IbvMemoryRegionKeepalive>,
71 ) -> Self {
72 Self {
73 virtual_addr,
74 rdma_addr,
75 size,
76 lkey,
77 rkey,
78 device_name,
79 _guard: guard,
80 }
81 }
82
83 /// Returns a view of `size` bytes starting at `offset`.
84 ///
85 /// The slice shares the same registration, keys, and device, and keeps the
86 /// registration alive.
87 pub fn try_slice(&self, offset: usize, size: usize) -> anyhow::Result<Self> {
88 anyhow::ensure!(
89 offset.checked_add(size).is_some_and(|end| end <= self.size),
90 "slice [{offset}, {offset}+{size}) is out of bounds for a {}-byte registration of \
91 [{:#x}, {:#x}) on {}",
92 self.size,
93 self.rdma_addr,
94 self.rdma_addr + self.size,
95 self.device_name,
96 );
97 Ok(Self {
98 virtual_addr: self.virtual_addr + offset,
99 rdma_addr: self.rdma_addr + offset,
100 size,
101 ..self.clone()
102 })
103 }
104
105 /// Fabricate a zero-sized view naming `device_name` and carrying `key` as
106 /// both its `lkey` and its `rkey`, over a null MR keepalive whose `Drop`
107 /// is a no-op. For tests that care only about which device a registration
108 /// belongs to, and which of several registrations they are holding.
109 #[cfg(test)]
110 pub(crate) fn for_test(device_name: &str, key: u32) -> Self {
111 Self::new(
112 0,
113 0,
114 0,
115 key,
116 key,
117 device_name.to_string(),
118 Arc::new(IbvMr::null()),
119 )
120 }
121}
122
123/// What a peer needs in order to address one of our registered memory regions
124/// over RDMA: the region's `rkey`, its RDMA address, its size, and the device
125/// serving it.
126///
127/// This is the wire form of an [`IbvMemoryRegionView`]. It carries no `lkey` and
128/// no keepalive: an `lkey` only means anything to the protection domain that
129/// issued it, and the registration is kept alive by the views the owning side
130/// holds.
131#[derive(Debug, Clone, Serialize, Deserialize, Named)]
132pub struct IbvRemoteMemoryRegionView {
133 pub rkey: u32,
134 /// RDMA address (may differ from virtual address).
135 pub addr: usize,
136 pub size: usize,
137 /// Name of the RDMA device this region is registered on (e.g., "mlx5_0").
138 pub device_name: String,
139}
140
141impl IbvRemoteMemoryRegionView {
142 /// Returns a view of `size` bytes starting at `offset` in the peer's region.
143 /// The slice retains the remote key and device.
144 pub fn try_slice(&self, offset: usize, size: usize) -> anyhow::Result<Self> {
145 anyhow::ensure!(
146 offset.checked_add(size).is_some_and(|end| end <= self.size),
147 "slice [{offset}, {offset}+{size}) is out of bounds for a {}-byte peer registration \
148 of [{:#x}, {:#x}) on {}",
149 self.size,
150 self.addr,
151 self.addr + self.size,
152 self.device_name,
153 );
154 Ok(Self {
155 addr: self.addr + offset,
156 size,
157 ..self.clone()
158 })
159 }
160}
161
162impl From<&IbvMemoryRegionView> for IbvRemoteMemoryRegionView {
163 /// The wire transport details are fully derived from the registered MR
164 /// view: the remote key, the RDMA address, the size, and the device name.
165 fn from(view: &IbvMemoryRegionView) -> Self {
166 Self {
167 rkey: view.rkey,
168 addr: view.rdma_addr,
169 size: view.size,
170 device_name: view.device_name.clone(),
171 }
172 }
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 fn view(rdma_addr: usize, size: usize) -> IbvMemoryRegionView {
180 IbvMemoryRegionView::new(
181 0x1000,
182 rdma_addr,
183 size,
184 0x1234,
185 0x5678,
186 "mlx5_0".to_string(),
187 Arc::new(IbvMr::null()),
188 )
189 }
190
191 #[test]
192 fn a_slice_moves_both_addresses_and_keeps_the_keys() {
193 let sliced = view(0x4000, 4096)
194 .try_slice(1024, 512)
195 .expect("a slice inside the region");
196 assert_eq!(sliced.rdma_addr, 0x4400, "the RDMA address advances");
197 assert_eq!(sliced.virtual_addr, 0x1400, "so does the virtual address");
198 assert_eq!(sliced.size, 512);
199 assert_eq!((sliced.lkey, sliced.rkey), (0x1234, 0x5678));
200 assert_eq!(sliced.device_name, "mlx5_0");
201 }
202
203 #[test]
204 fn a_remote_slice_moves_the_address_and_keeps_the_rkey() {
205 let remote = IbvRemoteMemoryRegionView::from(&view(0x4000, 4096));
206 let sliced = remote.try_slice(2048, 2048).expect("a slice at the end");
207 assert_eq!((sliced.addr, sliced.size), (0x4800, 2048));
208 assert_eq!(sliced.rkey, 0x5678);
209 assert_eq!(sliced.device_name, "mlx5_0");
210 }
211
212 #[test]
213 fn slices_are_bounded_by_the_registration() {
214 let local = view(0x4000, 4096);
215 let remote = IbvRemoteMemoryRegionView::from(&local);
216 assert!(local.try_slice(0, 4096).is_ok());
217 assert!(local.try_slice(4096, 0).is_ok());
218 assert!(local.try_slice(0, 4097).is_err());
219 assert!(local.try_slice(4096, 1).is_err());
220 assert!(local.try_slice(1, usize::MAX).is_err());
221 assert!(remote.try_slice(2048, 2049).is_err());
222 }
223}