[PATCH] gpu: nova-core: add ChannelIdPool

[Eliot Courtney <ecourtney@nvidia.com>]

Add `ChannelIdPool` which does allocation of channel IDs on top of a
maple tree. This is necessary for allotting ranges of channel IDs to be
used in e.g. vGPU and also later for allocating single channel IDs for
the host.

Pre-blackwell needs to allocate channel IDs with a particular alignment.
For this case, loop on `MapleTreeAlloc::alloc_range` while aligning up.
For unaligned cases, this always terminates on the first iteration.
Use a Mutex to ensure that the multi step insert-check-erase can't race.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
Add tracking of channel ID allocations, needed for allotting ranges of
channel IDs for vGPU guests and also (later) for regular host channel ID
reservation. Allocating channel IDs is not on a hot path (absolute peak
~100 per second, on average < 1 per second).

For pre-blackwell we need to be able to allocate a range of channel IDs
with a particular alignment. The approach I took here is to wrap
MapleTreeAlloc in a Mutex and loop on regular `alloc_range` lookups. An
alternative would be adding rust bindings for maple tree
`mas_empty_area` and doing a try+restart kind of thing.
---
 drivers/gpu/nova-core/gpu.rs         |   2 +
 drivers/gpu/nova-core/gpu/channel.rs | 174 +++++++++++++++++++++++++++++++++++
 2 files changed, 176 insertions(+)

diff --git a/drivers/gpu/nova-core/gpu.rs b/drivers/gpu/nova-core/gpu.rs
index 43c3f4f8df71..0c3b5de7d849 100644
--- a/drivers/gpu/nova-core/gpu.rs
+++ b/drivers/gpu/nova-core/gpu.rs
@@ -31,6 +31,8 @@
     regs,
 };
 
+#[cfg_attr(not(CONFIG_KUNIT = "y"), expect(dead_code))]
+mod channel;
 mod hal;
 
 macro_rules! define_chipset {
diff --git a/drivers/gpu/nova-core/gpu/channel.rs b/drivers/gpu/nova-core/gpu/channel.rs
new file mode 100644
index 000000000000..1f6299aa9f58
--- /dev/null
+++ b/drivers/gpu/nova-core/gpu/channel.rs
@@ -0,0 +1,174 @@
+// SPDX-License-Identifier: GPL-2.0
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+//! Channel related code.
+
+use core::{
+    num::NonZero,
+    ops::{
+        Deref,
+        Range, //
+    }, //
+};
+
+use kernel::{
+    maple_tree::MapleTreeAlloc,
+    prelude::*,
+    ptr::{
+        Alignable,
+        Alignment, //
+    },
+    sync::{
+        new_mutex,
+        Mutex, //
+    }, //
+};
+
+/// Pool for tracking reservations of channel IDs.
+#[pin_data]
+pub(crate) struct ChannelIdPool {
+    #[pin]
+    inner: Mutex<MapleTreeAlloc<()>>,
+    num_chids: usize,
+}
+
+impl ChannelIdPool {
+    /// Creates a pool managing `num_chids` channel IDs.
+    pub(crate) fn new(num_chids: usize) -> impl PinInit<Self> {
+        pin_init!(Self {
+            inner <- new_mutex!(MapleTreeAlloc::new()),
+            num_chids,
+        })
+    }
+
+    /// Reserves a contiguous area of `count` channel IDs starting at a multiple of `align`.
+    /// Returns a guard that releases the area on drop.
+    pub(crate) fn alloc_area(
+        &self,
+        count: NonZero<usize>,
+        align: Alignment,
+    ) -> Result<ChannelIdArea<'_>> {
+        // Hold the lock over the whole loop so the release and retry below cannot race with another
+        // allocation.
+        let tree = self.inner.lock();
+        let mut base = 0;
+        while base < self.num_chids {
+            let start = tree.alloc_range(count.get(), (), base..self.num_chids, GFP_KERNEL)?;
+            let aligned = start.align_up(align);
+
+            if aligned == Some(start) {
+                return Ok(ChannelIdArea {
+                    pool: self,
+                    range: start..start + count.get(),
+                });
+            }
+
+            // Release the unaligned area and retry from the next aligned channel ID.
+            tree.erase(start);
+            base = aligned.ok_or(EBUSY)?;
+        }
+        Err(EBUSY)
+    }
+}
+
+/// A reserved contiguous area of channel IDs.
+///
+/// Releases the whole area back to its [`ChannelIdPool`] when dropped. Releasing locks a sleeping
+/// [`Mutex`] and may allocate memory with `GFP_KERNEL`, so the area must be dropped in a context
+/// that is allowed to do so.
+#[must_use = "the channel ID area is released immediately when unused"]
+pub(crate) struct ChannelIdArea<'a> {
+    pool: &'a ChannelIdPool,
+    range: Range<usize>,
+}
+
+impl Drop for ChannelIdArea<'_> {
+    fn drop(&mut self) {
+        self.pool.inner.lock().erase(self.range.start);
+    }
+}
+
+impl Deref for ChannelIdArea<'_> {
+    type Target = Range<usize>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.range
+    }
+}
+
+#[kunit_tests(nova_core_channel)]
+mod tests {
+    use super::*;
+
+    const fn nz(count: usize) -> NonZero<usize> {
+        NonZero::new(count).unwrap()
+    }
+
+    #[test]
+    fn chid_area() -> Result {
+        let pool = KBox::pin_init(ChannelIdPool::new(2048), GFP_KERNEL)?;
+        let unaligned = Alignment::new::<1>();
+
+        let first = pool.alloc_area(nz(48), unaligned)?;
+        let second = pool.alloc_area(nz(48), unaligned)?;
+        assert_eq!(0, first.start);
+        assert_eq!(48, second.start);
+
+        drop(first);
+        let third = pool.alloc_area(nz(48), unaligned)?;
+        assert_eq!(0, third.start);
+        assert_eq!(96, pool.alloc_area(nz(48), unaligned)?.start);
+        Ok(())
+    }
+
+    #[test]
+    fn chid_bounded_by_num_chids() -> Result {
+        let pool = KBox::pin_init(ChannelIdPool::new(4), GFP_KERNEL)?;
+        let unaligned = Alignment::new::<1>();
+
+        assert_eq!(0, pool.alloc_area(nz(4), unaligned)?.start);
+
+        // An area larger than the pool is rejected by the maple tree itself.
+        assert_eq!(Some(EINVAL), pool.alloc_area(nz(5), unaligned).err());
+
+        let a = pool.alloc_area(nz(3), unaligned)?;
+        assert_eq!(0, a.start);
+        assert_eq!(Some(EBUSY), pool.alloc_area(nz(2), unaligned).err());
+
+        let b = pool.alloc_area(nz(1), unaligned)?;
+        assert_eq!(3, b.start);
+        assert_eq!(Some(EBUSY), pool.alloc_area(nz(1), unaligned).err());
+        Ok(())
+    }
+
+    #[test]
+    fn chid_area_aligned() -> Result {
+        let pool = KBox::pin_init(ChannelIdPool::new(16), GFP_KERNEL)?;
+        let unaligned = Alignment::new::<1>();
+        let align4 = Alignment::new::<4>();
+
+        // Alloc 0 so the first fit for the next area is unaligned.
+        let pad = pool.alloc_area(nz(1), unaligned)?;
+        assert_eq!(0, pad.start);
+
+        let a = pool.alloc_area(nz(4), align4)?;
+        assert_eq!(4, a.start);
+
+        // The area skipped over by the aligned allocation should still be available.
+        let b = pool.alloc_area(nz(1), unaligned)?;
+        assert_eq!(1, b.start);
+
+        let c = pool.alloc_area(nz(8), Alignment::new::<8>())?;
+        assert_eq!(8, c.start);
+
+        // Only 2 IDs left.
+        assert_eq!(Some(EBUSY), pool.alloc_area(nz(4), align4).err());
+        assert_eq!(
+            Some(EBUSY),
+            pool.alloc_area(nz(1), Alignment::new::<32>()).err()
+        );
+
+        assert_eq!(2, pool.alloc_area(nz(2), unaligned)?.start);
+        Ok(())
+    }
+}

---
base-commit: 78900738d981c97aad929e33620340d0d9bcfc82
change-id: 20260709-chid-maple-3d696554e0a3

Best regards,
--  
Eliot Courtney <ecourtney@nvidia.com>