monarch_rdma/backend/ibverbs/

device.rs

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Backend-agnostic ibverbs device abstraction.
//!
//! [`IbvDevice`] owns the per-process state for a single opened RDMA
//! device: an `Arc<IbvContext>`, an [`IbvConfig`], and a map of named
//! [`IbvDomain`]s allocated against the device.
//! Per-backend behavior — claiming a device as the backend's,
//! allocating a domain, and seeding config defaults — is provided by
//! an [`IbvDeviceImpl`].
//!
//! Each [`IbvDeviceImpl`] registers itself via the `inventory` crate
//! (using [`register_ibv_device_impl!`]). At first access,
//! [`DEVICE_NAMES_BY_IMPL`] walks the ibverbs device list once and
//! assigns each device to the first registered impl whose
//! [`IbvDeviceImpl::is_instance`] claims it, caching the resulting
//! `typename() → device-infos` map. [`IbvDevice::open`] consults this
//! map: it returns `None` when `name` is not advertised under
//! `I::typename()`, and otherwise opens the device using the cached
//! [`IbvDeviceInfo`].

use std::collections::HashMap;
use std::ffi::CStr;
use std::marker::PhantomData;
use std::sync::Arc;
use std::sync::LazyLock;

use typeuri::Named;

use super::cq_pool::CqLease;
use super::cq_pool::CqPool;
use super::domain::IbvDomain;
use super::domain::IbvDomainImpl;
use super::primitives::IbvConfig;
use super::primitives::IbvContext;
use super::primitives::IbvDeviceInfo;
use super::primitives::query_device_info;

/// Per-backend driver for [`IbvDevice`]. Concrete impls register
/// themselves via [`register_ibv_device_impl!`].
pub trait IbvDeviceImpl: Named + std::fmt::Debug + Send + Sync + 'static {
    /// The per-domain strategy used by this backend (how memory regions
    /// are registered and queue pairs built against a PD). [`IbvDomain`] is
    /// generic over this.
    type Domain: IbvDomainImpl;

    /// Human-readable display name for the backend this impl
    /// drives (e.g., `"mellanox"`, `"efa"`). Surfaced in
    /// diagnostics; not used as a registry key.
    fn backend_name() -> &'static str;

    /// Returns `true` if `ctx` belongs to this backend. Called
    /// transiently, once per ibverbs device, while
    /// [`DEVICE_NAMES_BY_IMPL`] is built.
    fn is_instance(ctx: Arc<IbvContext>) -> bool;

    /// Seeds an [`IbvConfig`] with backend-appropriate defaults
    /// (e.g., EFA caps `max_send_sge` at 1).
    fn apply_config_defaults(config: &mut IbvConfig);
}

/// Inventory entry submitted by [`register_ibv_device_impl!`] for
/// each concrete [`IbvDeviceImpl`]. Captures function pointers to
/// the impl's `typename()`, `backend_name()`, and `is_instance()`,
/// all erased of the concrete type.
pub struct IbvDeviceImplRegistration {
    typename: fn() -> &'static str,
    backend_name: fn() -> &'static str,
    is_instance: fn(Arc<IbvContext>) -> bool,
}

/// Per-impl entry stored in [`DEVICE_NAMES_BY_IMPL`]. Keyed by
/// `typename()`; carries the human-readable `backend_name` for
/// diagnostics alongside the impl's device infos.
#[derive(Debug)]
struct RegisteredBackend {
    #[expect(
        dead_code,
        reason = "read only via Debug for the IbvDevice::open diagnostic warning"
    )]
    backend_name: &'static str,
    devices: Vec<IbvDeviceInfo>,
}

inventory::collect!(IbvDeviceImplRegistration);

/// Submits an `inventory` entry for a concrete [`IbvDeviceImpl`].
///
/// Place one invocation per impl at module scope. The expansion
/// references `inventory` and `typeuri` by bare crate name, so the
/// calling crate must list both as direct dependencies.
///
/// ```ignore
/// register_ibv_device_impl!(StandardImpl);
/// ```
#[macro_export]
macro_rules! register_ibv_device_impl {
    ($impl_ty:ty) => {
        inventory::submit! {
            $crate::backend::ibverbs::device::IbvDeviceImplRegistration::new(
                <$impl_ty as typeuri::Named>::typename,
                <$impl_ty as $crate::backend::ibverbs::device::IbvDeviceImpl>::backend_name,
                <$impl_ty as $crate::backend::ibverbs::device::IbvDeviceImpl>::is_instance,
            )
        }
    };
}

impl IbvDeviceImplRegistration {
    /// Construct a registration. Visible to the rest of the crate
    /// so the [`register_ibv_device_impl!`] macro can call it from
    /// sibling call sites.
    pub const fn new(
        typename: fn() -> &'static str,
        backend_name: fn() -> &'static str,
        is_instance: fn(Arc<IbvContext>) -> bool,
    ) -> Self {
        Self {
            typename,
            backend_name,
            is_instance,
        }
    }
}

/// Process-wide map from `IbvDeviceImpl::typename()` to a
/// [`RegisteredBackend`] holding the impl's display name and the
/// device infos it claims. Built once on first access by a single
/// walk of the ibverbs device list, assigning each device to the
/// first registration whose `is_instance` claims it.
static DEVICE_NAMES_BY_IMPL: LazyLock<HashMap<&'static str, RegisteredBackend>> =
    LazyLock::new(|| {
        let registrations: Vec<&IbvDeviceImplRegistration> =
            inventory::iter::<IbvDeviceImplRegistration>().collect();
        let mut by_impl: HashMap<&'static str, RegisteredBackend> = registrations
            .iter()
            .map(|reg| {
                (
                    (reg.typename)(),
                    RegisteredBackend {
                        backend_name: (reg.backend_name)(),
                        devices: Vec::new(),
                    },
                )
            })
            .collect();

        let mut num_devices = 0i32;
        // SAFETY: `ibv_get_device_list` populates `num_devices` and
        // returns either null or a pointer to `num_devices` entries;
        // we free it before returning.
        let device_list = unsafe { rdmaxcel_sys::ibv_get_device_list(&mut num_devices) };
        // A non-null list must be freed even when empty, so only the
        // null case returns early; the loop below is a no-op at zero
        // devices and the `ibv_free_device_list` at the end still runs.
        if device_list.is_null() {
            return by_impl;
        }
        for i in 0..num_devices {
            // SAFETY: `device_list` is non-null with `num_devices`
            // valid entries (checked above).
            let device = unsafe { *device_list.add(i as usize) };
            if device.is_null() {
                continue;
            }
            // SAFETY: `device` is non-null per the check above.
            let raw_ctx = unsafe { rdmaxcel_sys::ibv_open_device(device) };
            if raw_ctx.is_null() {
                continue;
            }
            // SAFETY: `raw_ctx` was just returned by `ibv_open_device` and is
            // non-null (checked above); this `Arc<IbvContext>` is its sole owner.
            let ctx = Arc::new(unsafe { IbvContext::from_raw(raw_ctx) });
            // Assign the device to the first impl that claims it.
            if let Some(reg) = registrations
                .iter()
                .find(|reg| (reg.is_instance)(Arc::clone(&ctx)))
            {
                // SAFETY: `device` and `ctx.as_ptr()` are non-null and
                // `ctx.as_ptr()` was returned by `ibv_open_device(device)`.
                match unsafe { query_device_info(device, ctx.as_ptr()) } {
                    Ok(info) => by_impl
                        .get_mut((reg.typename)())
                        .expect("registration typename present in map")
                        .devices
                        .push(info),
                    Err(err) => tracing::warn!("failed to query RDMA device info: {err:#}"),
                }
            }
            // `ctx` drops here, closing the transient context.
        }
        // SAFETY: `device_list` was returned by `ibv_get_device_list`
        // above and has not been freed.
        unsafe { rdmaxcel_sys::ibv_free_device_list(device_list) };
        by_impl
    });

/// An opened RDMA device.
///
/// Owns an `Arc<IbvContext>`, the queried [`IbvDeviceInfo`]
/// metadata, a device-scoped [`IbvDeviceConfig`], and a map of
/// named [`IbvDomain`]s allocated against the device (one PD
/// per name, created lazily by [`Self::get_or_create_domain`]).
///
/// `I` is the backend driver type, parameterizing all per-backend
/// behavior via [`IbvDeviceImpl`].
#[derive(Debug)]
pub(crate) struct IbvDevice<I: IbvDeviceImpl> {
    domains: HashMap<String, IbvDomain<I::Domain>>,
    device_info: IbvDeviceInfo,
    config: IbvConfig,
    context: Arc<IbvContext>,
    /// The completion queues the device's queue pairs share, created on demand.
    cq_pool: CqPool,
    _marker: PhantomData<I>,
}

#[expect(dead_code, reason = "called by manager_actor in follow-up commits")]
impl<I: IbvDeviceImpl> IbvDevice<I> {
    /// Opens `name` under impl `I`, storing `config` as-is, and returns an
    /// `IbvDevice<I>` with the device's queried [`IbvDeviceInfo`] and the given
    /// [`IbvConfig`].
    ///
    /// Fails when `name` is not one of the devices registered for
    /// `I::typename()`, when a registered device cannot be reopened, or when the
    /// device cannot hold a completion queue of the configured size.
    pub fn try_open(name: &str, config: IbvConfig) -> anyhow::Result<Self> {
        let device_info = DEVICE_NAMES_BY_IMPL
            .get(I::typename())
            .and_then(|entry| entry.devices.iter().find(|d| d.name() == name).cloned());
        let Some(device_info) = device_info else {
            anyhow::bail!(
                "ibv device {} not found under backend {}; available: {:?}",
                name,
                I::backend_name(),
                *DEVICE_NAMES_BY_IMPL,
            );
        };

        // The registry already confirmed `name` belongs to `I`, so
        // reopen the device by name; a miss here is a broken invariant.
        let mut num_devices = 0i32;
        // SAFETY: `ibv_get_device_list` populates `num_devices` and
        // returns either null or a pointer to `num_devices` entries;
        // we free it before returning.
        let device_list = unsafe { rdmaxcel_sys::ibv_get_device_list(&mut num_devices) };
        assert!(
            !device_list.is_null(),
            "RDMA device list was null while opening registered device {}",
            name,
        );
        if num_devices == 0 {
            // A non-null list must be freed even when empty, per the
            // libibverbs contract.
            // SAFETY: `device_list` is non-null (asserted above) and was
            // returned by `ibv_get_device_list`.
            unsafe { rdmaxcel_sys::ibv_free_device_list(device_list) };
            panic!(
                "no RDMA devices found while opening registered device {}",
                name
            );
        }
        let mut target = std::ptr::null_mut();
        let mut context = None;
        for i in 0..num_devices {
            // SAFETY: `device_list` is non-null with `num_devices`
            // valid entries (checked above).
            let device = unsafe { *device_list.add(i as usize) };
            if device.is_null() {
                continue;
            }
            // SAFETY: `device` is non-null per the check above;
            // `ibv_get_device_name` returns a null-terminated C string
            // owned by the device list.
            let dev_name = unsafe { CStr::from_ptr(rdmaxcel_sys::ibv_get_device_name(device)) };
            if dev_name.to_bytes() == name.as_bytes() {
                target = device;
                break;
            }
        }
        let mut failure = None;
        if target.is_null() {
            failure = Some(format!(
                "ibv device with name {} not found by ibv_get_device_list",
                name
            ));
        } else {
            // Open while `target` still points into `device_list`, then free.
            // SAFETY: `target`, when non-null, is one of `device_list`'s
            // entries; `ibv_open_device` returns null on failure.
            let raw_context = unsafe { rdmaxcel_sys::ibv_open_device(target) };
            if raw_context.is_null() {
                failure = Some(format!(
                    "registered ibv device {} could not be opened: {}",
                    name,
                    std::io::Error::last_os_error(),
                ));
            } else {
                // SAFETY: `raw_context` was just returned by `ibv_open_device` and is
                // non-null (checked above); this `Arc<IbvContext>` is its sole owner.
                context = Some(Arc::new(unsafe { IbvContext::from_raw(raw_context) }));
            }
        }

        // SAFETY: `device_list` was returned by `ibv_get_device_list`
        // above and has not been freed.
        unsafe { rdmaxcel_sys::ibv_free_device_list(device_list) };

        if let Some(failure) = failure {
            anyhow::bail!(failure);
        }

        let context = context.expect("device context opened above");
        let cq_pool = CqPool::new(Arc::clone(&context), &device_info, config.max_send_wr)?;
        Ok(Self {
            domains: HashMap::new(),
            device_info,
            config,
            context,
            cq_pool,
            _marker: PhantomData,
        })
    }

    /// Returns the `Arc<IbvContext>` for this device.
    pub fn context(&self) -> Arc<IbvContext> {
        Arc::clone(&self.context)
    }

    /// Returns the queried [`IbvDeviceInfo`] metadata for this
    /// device.
    pub fn device_info(&self) -> &IbvDeviceInfo {
        &self.device_info
    }

    /// Returns the [`IbvConfig`] this device was opened with.
    pub fn config(&self) -> &IbvConfig {
        &self.config
    }

    /// Returns the [`IbvDomain`] registered under `name`, if one has already
    /// been created. Read-only lookup; use [`Self::get_or_create_domain`] to
    /// create on demand.
    pub fn domain(&self, name: &str) -> Option<&IbvDomain<I::Domain>> {
        self.domains.get(name)
    }

    /// Returns the [`IbvDomain`] registered under `name`, creating (and
    /// caching) a new one via [`IbvDomain::new`] on first access.
    pub fn get_or_create_domain(&mut self, name: &str) -> anyhow::Result<&IbvDomain<I::Domain>> {
        if !self.domains.contains_key(name) {
            // SAFETY: `self.context` wraps the live `ibv_context` opened by
            // `IbvDevice::open`, valid for this device's lifetime.
            let domain = unsafe {
                IbvDomain::<I::Domain>::new(
                    Arc::clone(&self.context),
                    self.device_info.clone(),
                    &self.config,
                )
            }?;
            self.domains.insert(name.to_string(), domain);
        }
        Ok(self
            .domains
            .get(name)
            .expect("domain just inserted or already present"))
    }

    /// Creates a queue pair in the domain named `domain_name`, with one of this
    /// device's shared completion queues.
    ///
    /// The returned [`CqLease`] is the queue pair's lease on that completion
    /// queue. It must be held separately from the queue pair, because it must
    /// live until all of that queue pair's CQEs have been drained, which could
    /// differ from the queue pair's lifetime on unclean teardown.
    pub fn create_queue_pair(
        &mut self,
        domain_name: &str,
        config: &IbvConfig,
    ) -> anyhow::Result<(<I::Domain as IbvDomainImpl>::QueuePair, CqLease)> {
        // Reserve first: the lease borrows nothing from `self`, so the mutable
        // borrow ends here and the domain lookup below can take its own.
        let lease = self.cq_pool.acquire_one()?;
        let domain = self.get_or_create_domain(domain_name)?;
        let qp = domain.create_queue_pair(config, Arc::clone(lease.cq()))?;
        Ok((qp, lease))
    }

    /// Creates a queue pair in the domain named `domain_name` that never posts,
    /// and so needs no lease on the completion queue it is created against; see
    /// [`CqPool::cq_without_lease`].
    pub fn create_non_posting_queue_pair(
        &mut self,
        domain_name: &str,
        config: &IbvConfig,
    ) -> anyhow::Result<<I::Domain as IbvDomainImpl>::QueuePair> {
        let cq = self.cq_pool.cq_without_lease()?;
        let domain = self.get_or_create_domain(domain_name)?;
        domain.create_queue_pair(config, cq)
    }

    /// Whether any device claimed by impl `I` is present on this
    /// host. Reads the [`DEVICE_NAMES_BY_IMPL`] registry, built on
    /// first access by a single walk of the ibverbs device list.
    pub fn available() -> bool {
        DEVICE_NAMES_BY_IMPL
            .get(I::typename())
            .is_some_and(|backend| !backend.devices.is_empty())
    }

    /// All devices claimed by impl `I` on this host. Reads the
    /// [`DEVICE_NAMES_BY_IMPL`] registry, built on first access by a
    /// single walk of the ibverbs device list.
    pub fn list() -> Vec<IbvDeviceInfo> {
        DEVICE_NAMES_BY_IMPL
            .get(I::typename())
            .map(|backend| backend.devices.clone())
            .unwrap_or_default()
    }
}

/// All RDMA devices on this host, across every registered backend
/// impl. Reads the [`DEVICE_NAMES_BY_IMPL`] registry, built on first
/// access by a single walk of the ibverbs device list.
///
/// A device claimed by no registered [`IbvDeviceImpl`] is omitted; in
/// practice every RDMA NIC is claimed by some backend.
pub fn list_all_devices() -> Vec<IbvDeviceInfo> {
    DEVICE_NAMES_BY_IMPL
        .values()
        .flat_map(|backend| backend.devices.iter().cloned())
        .collect()
}