[PATCH 0/9] rust: dma: add CoherentArray for compile-time sized allocations

[PATCH 0/9] rust: dma: add CoherentArray for compile-time sized allocations

[Eliot Courtney]

This series extends the DMA coherent allocation API to support compile-time
known sizes. This lets bounds checking to be moved from runtime to build
time, which is useful to avoid runtime panics from index typos. It also
removes the need for a Result return type in some places.

The compile time size is specified via a marker type: StaticSize<N>.
Statically sized allocations can decay to runtime sized ones via deref
coercion for code that doesn't need to know the size at compile time, or to
avoid having to carry around extra type parameters. The implementation
follows a similar pattern to Device/DeviceContext.

The series defines three type aliases: CoherentSlice<T> (for runtime size),
CoherentArray<T, N> (for compile-time size N), and CoherentObject<T> (for
single object allocations). It also adds infallible dma_read!/dma_write!
macros and methods to CoherentArray, while prefixing the existing fallible
methods and macros with `try_`.

The macros keep the same syntax (i.e.
coherent_allocation[index].optional_fields = expression) even for
CoherentObject, because the [] syntax is needed to know where to split the
actual CoherentAllocation object from the fields. This means that
CoherentObject is indexed with [0] in dma_write!/dma_read! macros. The
alternative is defining a separate macro for single object access, but it
still would need a way to delineate between the allocation and the fields,
perhaps by using commas (dma_read_obj!(object, fields),
dma_write_obj!(object, fields, value)). This would be inconsistent with the
array/slice syntax.

The last patch in the series may be useful as an example of what this
looks like to use. Also, there is probably a better name than
CoherentSlice. I found that specifying a default of RuntimeSize on
CoherentAllocation stopped the compiler from being able to resolve
which alloc_attrs to call in usages like e.g.
CoherentAllocation<u8>::alloc_attrs. Also, we probably want to encourage
people to use the statically sized one if possible, so it may be nice to
avoid defaulting CoherentAllocation to RuntimeSize.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
Eliot Courtney (9):
      rust: dma: rename CoherentAllocation fallible methods
      rust: dma: parameterize CoherentAllocation with AllocationSize
      rust: dma: add CoherentArray for compile-time sized allocations
      rust: dma: simplify try_dma_read! and try_dma_write!
      rust: dma: rename try_item_from_index to try_ptr_at
      rust: dma: add dma_read! and dma_write! macros
      rust: dma: implement decay from CoherentArray to CoherentSlice
      rust: dma: add CoherentObject for single element allocations
      gpu: nova-core: migrate to CoherentArray and CoherentObject

 drivers/gpu/nova-core/dma.rs            |  10 +-
 drivers/gpu/nova-core/falcon.rs         |   2 +-
 drivers/gpu/nova-core/firmware/fwsec.rs |   4 +-
 drivers/gpu/nova-core/gsp.rs            |  44 +--
 drivers/gpu/nova-core/gsp/boot.rs       |   6 +-
 drivers/gpu/nova-core/gsp/cmdq.rs       |  20 +-
 drivers/gpu/nova-core/gsp/fw.rs         |  12 +-
 rust/kernel/dma.rs                      | 555 +++++++++++++++++++++++++-------
 samples/rust/rust_dma.rs                |  14 +-
 9 files changed, 489 insertions(+), 178 deletions(-)
---
base-commit: c71257394bc9c59ea727803f6e55e83fe63db74e
change-id: 20260128-coherent-array-0321eb723d4c

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

[PATCH 1/9] rust: dma: rename CoherentAllocation fallible methods

[Eliot Courtney]

Prefix fallible methods in CoherentAllocation with try_. Prefix
dma_write! and dma_read! macros with try_ to better indicate they
can fail.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 drivers/gpu/nova-core/dma.rs            |  2 +-
 drivers/gpu/nova-core/falcon.rs         |  2 +-
 drivers/gpu/nova-core/firmware/fwsec.rs |  4 +-
 drivers/gpu/nova-core/gsp.rs            | 16 +++----
 drivers/gpu/nova-core/gsp/boot.rs       |  6 +--
 drivers/gpu/nova-core/gsp/cmdq.rs       | 14 +++---
 rust/kernel/dma.rs                      | 85 +++++++++++++++++----------------
 samples/rust/rust_dma.rs                |  6 +--
 8 files changed, 69 insertions(+), 66 deletions(-)

diff --git a/drivers/gpu/nova-core/dma.rs b/drivers/gpu/nova-core/dma.rs
index 7215398969da..f77754f12f02 100644
--- a/drivers/gpu/nova-core/dma.rs
+++ b/drivers/gpu/nova-core/dma.rs
@@ -33,7 +33,7 @@ pub(crate) fn from_data(dev: &device::Device<device::Bound>, data: &[u8]) -> Res
         Self::new(dev, data.len()).and_then(|mut dma_obj| {
             // SAFETY: We have just allocated the DMA memory, we are the only users and
             // we haven't made the device aware of the handle yet.
-            unsafe { dma_obj.write(data, 0)? }
+            unsafe { dma_obj.try_write(data, 0)? }
             Ok(dma_obj)
         })
     }
diff --git a/drivers/gpu/nova-core/falcon.rs b/drivers/gpu/nova-core/falcon.rs
index 82c661aef594..9cd271de0554 100644
--- a/drivers/gpu/nova-core/falcon.rs
+++ b/drivers/gpu/nova-core/falcon.rs
@@ -460,7 +460,7 @@ fn dma_wr<F: FalconFirmware<Target = E>>(
             FalconMem::Imem => (load_offsets.src_start, fw.dma_handle()),
             FalconMem::Dmem => (
                 0,
-                fw.dma_handle_with_offset(load_offsets.src_start.into_safe_cast())?,
+                fw.try_dma_handle_with_offset(load_offsets.src_start.into_safe_cast())?,
             ),
         };
         if dma_start % DmaAddress::from(DMA_LEN) > 0 {
diff --git a/drivers/gpu/nova-core/firmware/fwsec.rs b/drivers/gpu/nova-core/firmware/fwsec.rs
index b28e34d279f4..515b19926b49 100644
--- a/drivers/gpu/nova-core/firmware/fwsec.rs
+++ b/drivers/gpu/nova-core/firmware/fwsec.rs
@@ -191,7 +191,7 @@ unsafe fn transmute<T: Sized + FromBytes>(fw: &DmaObject, offset: usize) -> Resu
     // SAFETY: The safety requirements of the function guarantee the device won't read
     // or write to memory while the reference is alive and that this call won't race
     // with writes to the same memory region.
-    T::from_bytes(unsafe { fw.as_slice(offset, size_of::<T>())? }).ok_or(EINVAL)
+    T::from_bytes(unsafe { fw.try_as_slice(offset, size_of::<T>())? }).ok_or(EINVAL)
 }
 
 /// Reinterpret the area starting from `offset` in `fw` as a mutable instance of `T` (which must
@@ -210,7 +210,7 @@ unsafe fn transmute_mut<T: Sized + FromBytes + AsBytes>(
     // SAFETY: The safety requirements of the function guarantee the device won't read
     // or write to memory while the reference is alive and that this call won't race
     // with writes or reads to the same memory region.
-    T::from_bytes_mut(unsafe { fw.as_slice_mut(offset, size_of::<T>())? }).ok_or(EINVAL)
+    T::from_bytes_mut(unsafe { fw.try_as_slice_mut(offset, size_of::<T>())? }).ok_or(EINVAL)
 }
 
 /// The FWSEC microcode, extracted from the BIOS and to be run on the GSP falcon.
diff --git a/drivers/gpu/nova-core/gsp.rs b/drivers/gpu/nova-core/gsp.rs
index fb6f74797178..43bc35fd3b55 100644
--- a/drivers/gpu/nova-core/gsp.rs
+++ b/drivers/gpu/nova-core/gsp.rs
@@ -8,10 +8,10 @@
         CoherentAllocation,
         DmaAddress, //
     },
-    dma_write,
     pci,
     prelude::*,
-    transmute::AsBytes, //
+    transmute::AsBytes,
+    try_dma_write, //
 };
 
 pub(crate) mod cmdq;
@@ -92,7 +92,7 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         unsafe {
             // Copy the self-mapping PTE at the expected location.
             obj.0
-                .as_slice_mut(size_of::<u64>(), size_of_val(&ptes))?
+                .try_as_slice_mut(size_of::<u64>(), size_of_val(&ptes))?
                 .copy_from_slice(ptes.as_bytes())
         };
 
@@ -131,13 +131,13 @@ pub(crate) fn new(pdev: &pci::Device<device::Bound>) -> Result<impl PinInit<Self
         // _kgspInitLibosLoggingStructures (allocates memory for buffers)
         // kgspSetupLibosInitArgs_IMPL (creates pLibosInitArgs[] array)
         let loginit = LogBuffer::new(dev)?;
-        dma_write!(libos[0] = LibosMemoryRegionInitArgument::new("LOGINIT", &loginit.0))?;
+        try_dma_write!(libos[0] = LibosMemoryRegionInitArgument::new("LOGINIT", &loginit.0))?;
 
         let logintr = LogBuffer::new(dev)?;
-        dma_write!(libos[1] = LibosMemoryRegionInitArgument::new("LOGINTR", &logintr.0))?;
+        try_dma_write!(libos[1] = LibosMemoryRegionInitArgument::new("LOGINTR", &logintr.0))?;
 
         let logrm = LogBuffer::new(dev)?;
-        dma_write!(libos[2] = LibosMemoryRegionInitArgument::new("LOGRM", &logrm.0))?;
+        try_dma_write!(libos[2] = LibosMemoryRegionInitArgument::new("LOGRM", &logrm.0))?;
 
         let cmdq = Cmdq::new(dev)?;
 
@@ -146,8 +146,8 @@ pub(crate) fn new(pdev: &pci::Device<device::Bound>) -> Result<impl PinInit<Self
             1,
             GFP_KERNEL | __GFP_ZERO,
         )?;
-        dma_write!(rmargs[0] = fw::GspArgumentsCached::new(&cmdq))?;
-        dma_write!(libos[3] = LibosMemoryRegionInitArgument::new("RMARGS", &rmargs))?;
+        try_dma_write!(rmargs[0] = fw::GspArgumentsCached::new(&cmdq))?;
+        try_dma_write!(libos[3] = LibosMemoryRegionInitArgument::new("RMARGS", &rmargs))?;
 
         Ok(try_pin_init!(Self {
             libos,
diff --git a/drivers/gpu/nova-core/gsp/boot.rs b/drivers/gpu/nova-core/gsp/boot.rs
index 54937606b5b0..69e2fb064220 100644
--- a/drivers/gpu/nova-core/gsp/boot.rs
+++ b/drivers/gpu/nova-core/gsp/boot.rs
@@ -3,11 +3,11 @@
 use kernel::{
     device,
     dma::CoherentAllocation,
-    dma_write,
     io::poll::read_poll_timeout,
     pci,
     prelude::*,
-    time::Delta, //
+    time::Delta,
+    try_dma_write, //
 };
 
 use crate::{
@@ -160,7 +160,7 @@ pub(crate) fn boot(
 
         let wpr_meta =
             CoherentAllocation::<GspFwWprMeta>::alloc_coherent(dev, 1, GFP_KERNEL | __GFP_ZERO)?;
-        dma_write!(wpr_meta[0] = GspFwWprMeta::new(&gsp_fw, &fb_layout))?;
+        try_dma_write!(wpr_meta[0] = GspFwWprMeta::new(&gsp_fw, &fb_layout))?;
 
         self.cmdq
             .send_command(bar, commands::SetSystemInfo::new(pdev))?;
diff --git a/drivers/gpu/nova-core/gsp/cmdq.rs b/drivers/gpu/nova-core/gsp/cmdq.rs
index 3991ccc0c10f..9c94f4c6ff6d 100644
--- a/drivers/gpu/nova-core/gsp/cmdq.rs
+++ b/drivers/gpu/nova-core/gsp/cmdq.rs
@@ -15,7 +15,6 @@
         CoherentAllocation,
         DmaAddress, //
     },
-    dma_write,
     io::poll::read_poll_timeout,
     prelude::*,
     sync::aref::ARef,
@@ -24,6 +23,7 @@
         AsBytes,
         FromBytes, //
     },
+    try_dma_write, //
 };
 
 use crate::{
@@ -201,9 +201,11 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
 
         let gsp_mem =
             CoherentAllocation::<GspMem>::alloc_coherent(dev, 1, GFP_KERNEL | __GFP_ZERO)?;
-        dma_write!(gsp_mem[0].ptes = PteArray::new(gsp_mem.dma_handle())?)?;
-        dma_write!(gsp_mem[0].cpuq.tx = MsgqTxHeader::new(MSGQ_SIZE, RX_HDR_OFF, MSGQ_NUM_PAGES))?;
-        dma_write!(gsp_mem[0].cpuq.rx = MsgqRxHeader::new())?;
+        try_dma_write!(gsp_mem[0].ptes = PteArray::new(gsp_mem.dma_handle())?)?;
+        try_dma_write!(
+            gsp_mem[0].cpuq.tx = MsgqTxHeader::new(MSGQ_SIZE, RX_HDR_OFF, MSGQ_NUM_PAGES)
+        )?;
+        try_dma_write!(gsp_mem[0].cpuq.rx = MsgqRxHeader::new())?;
 
         Ok(Self(gsp_mem))
     }
@@ -221,7 +223,7 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         // - The `CoherentAllocation` contains exactly one object.
         // - We will only access the driver-owned part of the shared memory.
         // - Per the safety statement of the function, no concurrent access will be performed.
-        let gsp_mem = &mut unsafe { self.0.as_slice_mut(0, 1) }.unwrap()[0];
+        let gsp_mem = &mut unsafe { self.0.try_as_slice_mut(0, 1) }.unwrap()[0];
         // PANIC: per the invariant of `cpu_write_ptr`, `tx` is `<= MSGQ_NUM_PAGES`.
         let (before_tx, after_tx) = gsp_mem.cpuq.msgq.data.split_at_mut(tx);
 
@@ -256,7 +258,7 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         // - The `CoherentAllocation` contains exactly one object.
         // - We will only access the driver-owned part of the shared memory.
         // - Per the safety statement of the function, no concurrent access will be performed.
-        let gsp_mem = &unsafe { self.0.as_slice(0, 1) }.unwrap()[0];
+        let gsp_mem = &unsafe { self.0.try_as_slice(0, 1) }.unwrap()[0];
         // PANIC: per the invariant of `cpu_read_ptr`, `xx` is `<= MSGQ_NUM_PAGES`.
         let (before_rx, after_rx) = gsp_mem.gspq.msgq.data.split_at(rx);
 
diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 909d56fd5118..02321d5f3f06 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -482,7 +482,7 @@ pub fn dma_handle(&self) -> DmaAddress {
     /// device as the DMA address base of the region.
     ///
     /// Returns `EINVAL` if `offset` is not within the bounds of the allocation.
-    pub fn dma_handle_with_offset(&self, offset: usize) -> Result<DmaAddress> {
+    pub fn try_dma_handle_with_offset(&self, offset: usize) -> Result<DmaAddress> {
         if offset >= self.count {
             Err(EINVAL)
         } else {
@@ -494,7 +494,7 @@ pub fn dma_handle_with_offset(&self, offset: usize) -> Result<DmaAddress> {
 
     /// Common helper to validate a range applied from the allocated region in the CPU's virtual
     /// address space.
-    fn validate_range(&self, offset: usize, count: usize) -> Result {
+    fn try_validate_range(&self, offset: usize, count: usize) -> Result {
         if offset.checked_add(count).ok_or(EOVERFLOW)? > self.count {
             return Err(EINVAL);
         }
@@ -514,8 +514,8 @@ fn validate_range(&self, offset: usize, count: usize) -> Result {
     ///   slice is live.
     /// * Callers must ensure that this call does not race with a write to the same region while
     ///   the returned slice is live.
-    pub unsafe fn as_slice(&self, offset: usize, count: usize) -> Result<&[T]> {
-        self.validate_range(offset, count)?;
+    pub unsafe fn try_as_slice(&self, offset: usize, count: usize) -> Result<&[T]> {
+        self.try_validate_range(offset, count)?;
         // SAFETY:
         // - The pointer is valid due to type invariant on `CoherentAllocation`,
         //   we've just checked that the range and index is within bounds. The immutability of the
@@ -525,8 +525,8 @@ pub unsafe fn as_slice(&self, offset: usize, count: usize) -> Result<&[T]> {
         Ok(unsafe { core::slice::from_raw_parts(self.start_ptr().add(offset), count) })
     }
 
-    /// Performs the same functionality as [`CoherentAllocation::as_slice`], except that a mutable
-    /// slice is returned.
+    /// Performs the same functionality as [`CoherentAllocation::try_as_slice`], except that a
+    /// mutable slice is returned.
     ///
     /// # Safety
     ///
@@ -534,8 +534,8 @@ pub unsafe fn as_slice(&self, offset: usize, count: usize) -> Result<&[T]> {
     ///   slice is live.
     /// * Callers must ensure that this call does not race with a read or write to the same region
     ///   while the returned slice is live.
-    pub unsafe fn as_slice_mut(&mut self, offset: usize, count: usize) -> Result<&mut [T]> {
-        self.validate_range(offset, count)?;
+    pub unsafe fn try_as_slice_mut(&mut self, offset: usize, count: usize) -> Result<&mut [T]> {
+        self.try_validate_range(offset, count)?;
         // SAFETY:
         // - The pointer is valid due to type invariant on `CoherentAllocation`,
         //   we've just checked that the range and index is within bounds. The immutability of the
@@ -561,11 +561,11 @@ pub unsafe fn as_slice_mut(&mut self, offset: usize, count: usize) -> Result<&mu
     /// let buf: &[u8] = &somedata;
     /// // SAFETY: There is no concurrent HW operation on the device and no other R/W access to the
     /// // region.
-    /// unsafe { alloc.write(buf, 0)?; }
+    /// unsafe { alloc.try_write(buf, 0)?; }
     /// # Ok::<(), Error>(()) }
     /// ```
-    pub unsafe fn write(&mut self, src: &[T], offset: usize) -> Result {
-        self.validate_range(offset, src.len())?;
+    pub unsafe fn try_write(&mut self, src: &[T], offset: usize) -> Result {
+        self.try_validate_range(offset, src.len())?;
         // SAFETY:
         // - The pointer is valid due to type invariant on `CoherentAllocation`
         //   and we've just checked that the range and index is within bounds.
@@ -581,12 +581,13 @@ pub unsafe fn write(&mut self, src: &[T], offset: usize) -> Result {
         Ok(())
     }
 
-    /// Returns a pointer to an element from the region with bounds checking. `offset` is in
-    /// units of `T`, not the number of bytes.
+    /// Returns a pointer to an element from the region with bounds checking. `offset` is in units
+    /// of `T`, not the number of bytes.
     ///
-    /// Public but hidden since it should only be used from [`dma_read`] and [`dma_write`] macros.
+    /// Public but hidden since it should only be used from [`try_dma_read`] and [`try_dma_write`]
+    /// macros.
     #[doc(hidden)]
-    pub fn item_from_index(&self, offset: usize) -> Result<*mut T> {
+    pub fn try_item_from_index(&self, offset: usize) -> Result<*mut T> {
         if offset >= self.count {
             return Err(EINVAL);
         }
@@ -602,10 +603,10 @@ pub fn item_from_index(&self, offset: usize) -> Result<*mut T> {
     ///
     /// # Safety
     ///
-    /// This must be called from the [`dma_read`] macro which ensures that the `field` pointer is
-    /// validated beforehand.
+    /// This must be called from the [`try_dma_read`] macro which ensures that the `field` pointer
+    /// is validated beforehand.
     ///
-    /// Public but hidden since it should only be used from [`dma_read`] macro.
+    /// Public but hidden since it should only be used from [`try_dma_read`] macro.
     #[doc(hidden)]
     pub unsafe fn field_read<F: FromBytes>(&self, field: *const F) -> F {
         // SAFETY:
@@ -625,10 +626,10 @@ pub unsafe fn field_read<F: FromBytes>(&self, field: *const F) -> F {
     ///
     /// # Safety
     ///
-    /// This must be called from the [`dma_write`] macro which ensures that the `field` pointer is
-    /// validated beforehand.
+    /// This must be called from the [`try_dma_write`] macro which ensures that the `field` pointer
+    /// is validated beforehand.
     ///
-    /// Public but hidden since it should only be used from [`dma_write`] macro.
+    /// Public but hidden since it should only be used from [`try_dma_write`] macro.
     #[doc(hidden)]
     pub unsafe fn field_write<F: AsBytes>(&self, field: *mut F, val: F) {
         // SAFETY:
@@ -684,18 +685,18 @@ unsafe impl<T: AsBytes + FromBytes + Send> Send for CoherentAllocation<T> {}
 /// unsafe impl kernel::transmute::AsBytes for MyStruct{};
 ///
 /// # fn test(alloc: &kernel::dma::CoherentAllocation<MyStruct>) -> Result {
-/// let whole = kernel::dma_read!(alloc[2]);
-/// let field = kernel::dma_read!(alloc[1].field);
+/// let whole = kernel::try_dma_read!(alloc[2]);
+/// let field = kernel::try_dma_read!(alloc[1].field);
 /// # Ok::<(), Error>(()) }
 /// ```
 #[macro_export]
-macro_rules! dma_read {
-    ($dma:expr, $idx: expr, $($field:tt)*) => {{
+macro_rules! try_dma_read {
+    ($dma:expr, $idx:expr, $($field:tt)*) => {{
         (|| -> ::core::result::Result<_, $crate::error::Error> {
-            let item = $crate::dma::CoherentAllocation::item_from_index(&$dma, $idx)?;
-            // SAFETY: `item_from_index` ensures that `item` is always a valid pointer and can be
-            // dereferenced. The compiler also further validates the expression on whether `field`
-            // is a member of `item` when expanded by the macro.
+            let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
+            // SAFETY: `try_item_from_index` ensures that `item` is always a valid pointer
+            // and can be dereferenced. The compiler also further validates the expression
+            // on whether `field` is a member of `item` when expanded by the macro.
             unsafe {
                 let ptr_field = ::core::ptr::addr_of!((*item) $($field)*);
                 ::core::result::Result::Ok(
@@ -705,10 +706,10 @@ macro_rules! dma_read {
         })()
     }};
     ($dma:ident [ $idx:expr ] $($field:tt)* ) => {
-        $crate::dma_read!($dma, $idx, $($field)*)
+        $crate::try_dma_read!($dma, $idx, $($field)*)
     };
     ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {
-        $crate::dma_read!($($dma).*, $idx, $($field)*)
+        $crate::try_dma_read!($($dma).*, $idx, $($field)*)
     };
 }
 
@@ -728,32 +729,32 @@ macro_rules! dma_read {
 /// unsafe impl kernel::transmute::AsBytes for MyStruct{};
 ///
 /// # fn test(alloc: &kernel::dma::CoherentAllocation<MyStruct>) -> Result {
-/// kernel::dma_write!(alloc[2].member = 0xf);
-/// kernel::dma_write!(alloc[1] = MyStruct { member: 0xf });
+/// kernel::try_dma_write!(alloc[2].member = 0xf);
+/// kernel::try_dma_write!(alloc[1] = MyStruct { member: 0xf });
 /// # Ok::<(), Error>(()) }
 /// ```
 #[macro_export]
-macro_rules! dma_write {
+macro_rules! try_dma_write {
     ($dma:ident [ $idx:expr ] $($field:tt)*) => {{
-        $crate::dma_write!($dma, $idx, $($field)*)
+        $crate::try_dma_write!($dma, $idx, $($field)*)
     }};
     ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {{
-        $crate::dma_write!($($dma).*, $idx, $($field)*)
+        $crate::try_dma_write!($($dma).*, $idx, $($field)*)
     }};
     ($dma:expr, $idx: expr, = $val:expr) => {
         (|| -> ::core::result::Result<_, $crate::error::Error> {
-            let item = $crate::dma::CoherentAllocation::item_from_index(&$dma, $idx)?;
-            // SAFETY: `item_from_index` ensures that `item` is always a valid item.
+            let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
+            // SAFETY: `try_item_from_index` ensures that `item` is always a valid item.
             unsafe { $crate::dma::CoherentAllocation::field_write(&$dma, item, $val) }
             ::core::result::Result::Ok(())
         })()
     };
     ($dma:expr, $idx: expr, $(.$field:ident)* = $val:expr) => {
         (|| -> ::core::result::Result<_, $crate::error::Error> {
-            let item = $crate::dma::CoherentAllocation::item_from_index(&$dma, $idx)?;
-            // SAFETY: `item_from_index` ensures that `item` is always a valid pointer and can be
-            // dereferenced. The compiler also further validates the expression on whether `field`
-            // is a member of `item` when expanded by the macro.
+            let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
+            // SAFETY: `try_item_from_index` ensures that `item` is always a valid pointer
+            // and can be dereferenced. The compiler also further validates the expression
+            // on whether `field` is a member of `item` when expanded by the macro.
             unsafe {
                 let ptr_field = ::core::ptr::addr_of_mut!((*item) $(.$field)*);
                 $crate::dma::CoherentAllocation::field_write(&$dma, ptr_field, $val)
diff --git a/samples/rust/rust_dma.rs b/samples/rust/rust_dma.rs
index 9c45851c876e..7a87048575df 100644
--- a/samples/rust/rust_dma.rs
+++ b/samples/rust/rust_dma.rs
@@ -68,7 +68,7 @@ fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, E
                 CoherentAllocation::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
 
             for (i, value) in TEST_VALUES.into_iter().enumerate() {
-                kernel::dma_write!(ca[i] = MyStruct::new(value.0, value.1))?;
+                kernel::try_dma_write!(ca[i] = MyStruct::new(value.0, value.1))?;
             }
 
             let size = 4 * page::PAGE_SIZE;
@@ -91,8 +91,8 @@ fn drop(self: Pin<&mut Self>) {
         dev_info!(self.pdev, "Unload DMA test driver.\n");
 
         for (i, value) in TEST_VALUES.into_iter().enumerate() {
-            let val0 = kernel::dma_read!(self.ca[i].h);
-            let val1 = kernel::dma_read!(self.ca[i].b);
+            let val0 = kernel::try_dma_read!(self.ca[i].h);
+            let val1 = kernel::try_dma_read!(self.ca[i].b);
             assert!(val0.is_ok());
             assert!(val1.is_ok());
 

-- 
2.52.0

[PATCH 2/9] rust: dma: parameterize CoherentAllocation with AllocationSize

[Eliot Courtney]

Parameterize CoherentAllocation with AllocationSize. This lets it
carry information about whether it knows its size at compile time.
This follows a similar design to Device and DeviceContext.

This is useful to be able to read/write without having to handle
a Result, and to move indexing errors from runtime to build time.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/dma.rs | 185 +++++++++++++++++++++++++++++++++--------------------
 1 file changed, 117 insertions(+), 68 deletions(-)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 02321d5f3f06..6e6d91a9cd62 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -12,7 +12,7 @@
     sync::aref::ARef,
     transmute::{AsBytes, FromBytes},
 };
-use core::ptr::NonNull;
+use core::{marker::PhantomData, ptr::NonNull};
 
 /// DMA address type.
 ///
@@ -344,6 +344,29 @@ fn from(direction: DataDirection) -> Self {
     }
 }
 
+/// Marker trait for the size parameter of a [`CoherentAllocation`].
+///
+/// [`AllocationSize`] is a marker trait for the size parameter of a [`CoherentAllocation`].
+///
+/// The specific types of size are `RuntimeSize` and `StaticSize<N>`.
+pub trait AllocationSize: private::Sealed {}
+
+/// Marker type for a [`CoherentAllocation`] with a runtime-determined size.
+pub struct RuntimeSize;
+
+/// Marker type for a [`CoherentAllocation`] with a compile-time-known size of `N` elements.
+pub struct StaticSize<const N: usize>;
+
+mod private {
+    pub trait Sealed {}
+
+    impl Sealed for super::RuntimeSize {}
+    impl<const N: usize> Sealed for super::StaticSize<N> {}
+}
+
+impl AllocationSize for RuntimeSize {}
+impl<const N: usize> AllocationSize for StaticSize<N> {}
+
 /// An abstraction of the `dma_alloc_coherent` API.
 ///
 /// This is an abstraction around the `dma_alloc_coherent` API which is used to allocate and map
@@ -361,6 +384,12 @@ fn from(direction: DataDirection) -> Self {
 ///   region.
 /// - The size in bytes of the allocation is equal to `size_of::<T> * count`.
 /// - `size_of::<T> * count` fits into a `usize`.
+/// - If parameterized by `StaticSize<N>`, then `count == N`.
+///
+/// # Allocation size
+///
+/// [`CoherentAllocation`] is generic over an [`AllocationSize`], which lets it record a compile
+/// time known size (in number of elements of `T`).
 // TODO
 //
 // DMA allocations potentially carry device resources (e.g.IOMMU mappings), hence for soundness
@@ -373,79 +402,19 @@ fn from(direction: DataDirection) -> Self {
 //
 // Hence, find a way to revoke the device resources of a `CoherentAllocation`, but not the
 // entire `CoherentAllocation` including the allocated memory itself.
-pub struct CoherentAllocation<T: AsBytes + FromBytes> {
+pub struct CoherentAllocation<T: AsBytes + FromBytes, Size: AllocationSize = RuntimeSize> {
     dev: ARef<device::Device>,
     dma_handle: DmaAddress,
     count: usize,
     cpu_addr: NonNull<T>,
     dma_attrs: Attrs,
+    _size: PhantomData<Size>,
 }
 
-impl<T: AsBytes + FromBytes> CoherentAllocation<T> {
-    /// Allocates a region of `size_of::<T> * count` of coherent memory.
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// # use kernel::device::{Bound, Device};
-    /// use kernel::dma::{attrs::*, CoherentAllocation};
-    ///
-    /// # fn test(dev: &Device<Bound>) -> Result {
-    /// let c: CoherentAllocation<u64> =
-    ///     CoherentAllocation::alloc_attrs(dev, 4, GFP_KERNEL, DMA_ATTR_NO_WARN)?;
-    /// # Ok::<(), Error>(()) }
-    /// ```
-    pub fn alloc_attrs(
-        dev: &device::Device<Bound>,
-        count: usize,
-        gfp_flags: kernel::alloc::Flags,
-        dma_attrs: Attrs,
-    ) -> Result<CoherentAllocation<T>> {
-        build_assert!(
-            core::mem::size_of::<T>() > 0,
-            "It doesn't make sense for the allocated type to be a ZST"
-        );
-
-        let size = count
-            .checked_mul(core::mem::size_of::<T>())
-            .ok_or(EOVERFLOW)?;
-        let mut dma_handle = 0;
-        // SAFETY: Device pointer is guaranteed as valid by the type invariant on `Device`.
-        let addr = unsafe {
-            bindings::dma_alloc_attrs(
-                dev.as_raw(),
-                size,
-                &mut dma_handle,
-                gfp_flags.as_raw(),
-                dma_attrs.as_raw(),
-            )
-        };
-        let addr = NonNull::new(addr).ok_or(ENOMEM)?;
-        // INVARIANT:
-        // - We just successfully allocated a coherent region which is accessible for
-        //   `count` elements, hence the cpu address is valid. We also hold a refcounted reference
-        //   to the device.
-        // - The allocated `size` is equal to `size_of::<T> * count`.
-        // - The allocated `size` fits into a `usize`.
-        Ok(Self {
-            dev: dev.into(),
-            dma_handle,
-            count,
-            cpu_addr: addr.cast(),
-            dma_attrs,
-        })
-    }
-
-    /// Performs the same functionality as [`CoherentAllocation::alloc_attrs`], except the
-    /// `dma_attrs` is 0 by default.
-    pub fn alloc_coherent(
-        dev: &device::Device<Bound>,
-        count: usize,
-        gfp_flags: kernel::alloc::Flags,
-    ) -> Result<CoherentAllocation<T>> {
-        CoherentAllocation::alloc_attrs(dev, count, gfp_flags, Attrs(0))
-    }
+/// A coherent DMA allocation with a runtime-determined size.
+pub type CoherentSlice<T> = CoherentAllocation<T, RuntimeSize>;
 
+impl<T: AsBytes + FromBytes, Size: AllocationSize> CoherentAllocation<T, Size> {
     /// Returns the number of elements `T` in this allocation.
     ///
     /// Note that this is not the size of the allocation in bytes, which is provided by
@@ -644,10 +613,87 @@ pub unsafe fn field_write<F: AsBytes>(&self, field: *mut F, val: F) {
         // the UB caused by racing between two kernel functions nor do they provide atomicity.
         unsafe { field.write_volatile(val) }
     }
+
+    // Allocates a region of `size_of::<T> * count` of coherent memory.
+    fn alloc_impl(
+        dev: &device::Device<Bound>,
+        count: usize,
+        gfp_flags: kernel::alloc::Flags,
+        dma_attrs: Attrs,
+    ) -> Result<Self> {
+        build_assert!(
+            core::mem::size_of::<T>() > 0,
+            "It doesn't make sense for the allocated type to be a ZST"
+        );
+
+        let size = count
+            .checked_mul(core::mem::size_of::<T>())
+            .ok_or(EOVERFLOW)?;
+        let mut dma_handle = 0;
+        // SAFETY: Device pointer is guaranteed as valid by the type invariant on `Device`.
+        let addr = unsafe {
+            bindings::dma_alloc_attrs(
+                dev.as_raw(),
+                size,
+                &mut dma_handle,
+                gfp_flags.as_raw(),
+                dma_attrs.as_raw(),
+            )
+        };
+        let addr = NonNull::new(addr).ok_or(ENOMEM)?;
+        // INVARIANT:
+        // - We just successfully allocated a coherent region which is accessible for
+        //   `count` elements, hence the cpu address is valid. We also hold a refcounted reference
+        //   to the device.
+        // - The allocated `size` is equal to `size_of::<T> * count`.
+        // - The allocated `size` fits into a `usize`.
+        Ok(Self {
+            dev: dev.into(),
+            dma_handle,
+            count,
+            cpu_addr: addr.cast(),
+            dma_attrs,
+            _size: PhantomData,
+        })
+    }
+}
+
+impl<T: AsBytes + FromBytes> CoherentSlice<T> {
+    /// Allocates a region of `size_of::<T> * count` of coherent memory.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use kernel::device::{Bound, Device};
+    /// use kernel::dma::{attrs::*, CoherentSlice};
+    ///
+    /// # fn test(dev: &Device<Bound>) -> Result {
+    /// let c: CoherentSlice<u64> =
+    ///     CoherentSlice::alloc_attrs(dev, 4, GFP_KERNEL, DMA_ATTR_NO_WARN)?;
+    /// # Ok::<(), Error>(()) }
+    /// ```
+    pub fn alloc_attrs(
+        dev: &device::Device<Bound>,
+        count: usize,
+        gfp_flags: kernel::alloc::Flags,
+        dma_attrs: Attrs,
+    ) -> Result<Self> {
+        Self::alloc_impl(dev, count, gfp_flags, dma_attrs)
+    }
+
+    /// Performs the same functionality as [`CoherentSlice::alloc_attrs`], except the
+    /// `dma_attrs` is 0 by default.
+    pub fn alloc_coherent(
+        dev: &device::Device<Bound>,
+        count: usize,
+        gfp_flags: kernel::alloc::Flags,
+    ) -> Result<Self> {
+        Self::alloc_attrs(dev, count, gfp_flags, Attrs(0))
+    }
 }
 
 /// Note that the device configured to do DMA must be halted before this object is dropped.
-impl<T: AsBytes + FromBytes> Drop for CoherentAllocation<T> {
+impl<T: AsBytes + FromBytes, Size: AllocationSize> Drop for CoherentAllocation<T, Size> {
     fn drop(&mut self) {
         let size = self.count * core::mem::size_of::<T>();
         // SAFETY: Device pointer is guaranteed as valid by the type invariant on `Device`.
@@ -667,7 +713,10 @@ fn drop(&mut self) {
 
 // SAFETY: It is safe to send a `CoherentAllocation` to another thread if `T`
 // can be sent to another thread.
-unsafe impl<T: AsBytes + FromBytes + Send> Send for CoherentAllocation<T> {}
+unsafe impl<T: AsBytes + FromBytes + Send, Size: AllocationSize> Send
+    for CoherentAllocation<T, Size>
+{
+}
 
 /// Reads a field of an item from an allocated region of structs.
 ///

-- 
2.52.0

[PATCH 3/9] rust: dma: add CoherentArray for compile-time sized allocations

[Eliot Courtney]

Add a CoherentArray type alias which takes the size parameter directly,
without using the StaticSize<N> marker type. This makes it a bit nicer
to use.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 drivers/gpu/nova-core/dma.rs |   8 +--
 rust/kernel/dma.rs           | 127 ++++++++++++++++++++++++++++++++++++++++++-
 samples/rust/rust_dma.rs     |   8 +--
 3 files changed, 132 insertions(+), 11 deletions(-)

diff --git a/drivers/gpu/nova-core/dma.rs b/drivers/gpu/nova-core/dma.rs
index f77754f12f02..c217cdb14223 100644
--- a/drivers/gpu/nova-core/dma.rs
+++ b/drivers/gpu/nova-core/dma.rs
@@ -9,13 +9,13 @@
 
 use kernel::{
     device,
-    dma::CoherentAllocation,
+    dma::CoherentSlice,
     page::PAGE_SIZE,
     prelude::*, //
 };
 
 pub(crate) struct DmaObject {
-    dma: CoherentAllocation<u8>,
+    dma: CoherentSlice<u8>,
 }
 
 impl DmaObject {
@@ -24,7 +24,7 @@ pub(crate) fn new(dev: &device::Device<device::Bound>, len: usize) -> Result<Sel
             .map_err(|_| EINVAL)?
             .pad_to_align()
             .size();
-        let dma = CoherentAllocation::alloc_coherent(dev, len, GFP_KERNEL | __GFP_ZERO)?;
+        let dma = CoherentSlice::alloc_coherent(dev, len, GFP_KERNEL | __GFP_ZERO)?;
 
         Ok(Self { dma })
     }
@@ -40,7 +40,7 @@ pub(crate) fn from_data(dev: &device::Device<device::Bound>, data: &[u8]) -> Res
 }
 
 impl Deref for DmaObject {
-    type Target = CoherentAllocation<u8>;
+    type Target = CoherentSlice<u8>;
 
     fn deref(&self) -> &Self::Target {
         &self.dma
diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 6e6d91a9cd62..43ed0dfdbc08 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -194,12 +194,12 @@ pub const fn value(&self) -> u64 {
 ///
 /// ```
 /// # use kernel::device::{Bound, Device};
-/// use kernel::dma::{attrs::*, CoherentAllocation};
+/// use kernel::dma::{attrs::*, CoherentArray};
 ///
 /// # fn test(dev: &Device<Bound>) -> Result {
 /// let attribs = DMA_ATTR_FORCE_CONTIGUOUS | DMA_ATTR_NO_WARN;
-/// let c: CoherentAllocation<u64> =
-///     CoherentAllocation::alloc_attrs(dev, 4, GFP_KERNEL, attribs)?;
+/// let c: CoherentArray<u64, 4> =
+///     CoherentArray::alloc_attrs(dev, GFP_KERNEL, attribs)?;
 /// # Ok::<(), Error>(()) }
 /// ```
 #[derive(Clone, Copy, PartialEq)]
@@ -414,6 +414,9 @@ pub struct CoherentAllocation<T: AsBytes + FromBytes, Size: AllocationSize = Run
 /// A coherent DMA allocation with a runtime-determined size.
 pub type CoherentSlice<T> = CoherentAllocation<T, RuntimeSize>;
 
+/// A coherent DMA allocation for an array of `N` elements.
+pub type CoherentArray<T, const N: usize> = CoherentAllocation<T, StaticSize<N>>;
+
 impl<T: AsBytes + FromBytes, Size: AllocationSize> CoherentAllocation<T, Size> {
     /// Returns the number of elements `T` in this allocation.
     ///
@@ -692,6 +695,124 @@ pub fn alloc_coherent(
     }
 }
 
+impl<T: AsBytes + FromBytes, const N: usize> CoherentArray<T, N> {
+    /// Allocates a region of `size_of::<T> * N` of coherent memory.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use kernel::device::{Bound, Device};
+    /// use kernel::dma::{attrs::*, CoherentArray};
+    ///
+    /// # fn test(dev: &Device<Bound>) -> Result {
+    /// let c: CoherentArray<u64, 4> =
+    ///     CoherentArray::alloc_attrs(dev, GFP_KERNEL, DMA_ATTR_NO_WARN)?;
+    /// # Ok::<(), Error>(()) }
+    /// ```
+    pub fn alloc_attrs(
+        dev: &device::Device<Bound>,
+        gfp_flags: kernel::alloc::Flags,
+        dma_attrs: Attrs,
+    ) -> Result<Self> {
+        Self::alloc_impl(dev, N, gfp_flags, dma_attrs)
+    }
+
+    /// Performs the same functionality as [`CoherentArray::alloc_attrs`], except the
+    /// `dma_attrs` is 0 by default.
+    pub fn alloc_coherent(
+        dev: &device::Device<Bound>,
+        gfp_flags: kernel::alloc::Flags,
+    ) -> Result<Self> {
+        Self::alloc_attrs(dev, gfp_flags, Attrs(0))
+    }
+
+    /// Returns a DMA handle starting at `OFFSET` (in units of `T`) which may be given to the
+    /// device as the DMA address base of the region.
+    pub fn dma_handle_with_offset<const OFFSET: usize>(&self) -> DmaAddress {
+        build_assert!(OFFSET < N, "Offset is out of bounds for the allocation.");
+
+        // INVARIANT: The type invariant of `Self` guarantees that `size_of::<T> * N` fits
+        // into a `usize`, and `OFFSET` is inferior to `N`.
+        self.dma_handle + (OFFSET * core::mem::size_of::<T>()) as DmaAddress
+    }
+
+    /// Returns the data from the region starting from `OFFSET` as a slice.
+    /// `OFFSET` and `COUNT` are in units of `T`, not the number of bytes.
+    ///
+    /// For ringbuffer type of r/w access or use-cases where the pointer to the live data is needed,
+    /// [`CoherentAllocation::start_ptr`] or [`CoherentAllocation::start_ptr_mut`] could be used
+    /// instead.
+    ///
+    /// # Safety
+    ///
+    /// * Callers must ensure that the device does not read/write to/from memory while the returned
+    ///   slice is live.
+    /// * Callers must ensure that this call does not race with a write to the same region while
+    ///   the returned slice is live.
+    pub unsafe fn as_slice<const OFFSET: usize, const COUNT: usize>(&self) -> &[T] {
+        build_assert!(
+            OFFSET + COUNT <= N,
+            "Range is out of bounds for the allocation."
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`,
+        //   we've just checked that the range and index is within bounds. The immutability of the
+        //   data is also guaranteed by the safety requirements of the function.
+        // - `OFFSET + COUNT` can't overflow since it is smaller than `N` and we've checked
+        //   that `N` won't overflow early in the constructor.
+        unsafe { core::slice::from_raw_parts(self.start_ptr().add(OFFSET), COUNT) }
+    }
+
+    /// Performs the same functionality as [`CoherentArray::as_slice`], except that a mutable
+    /// slice is returned.
+    ///
+    /// # Safety
+    ///
+    /// * Callers must ensure that the device does not read/write to/from memory while the returned
+    ///   slice is live.
+    /// * Callers must ensure that this call does not race with a read or write to the same region
+    ///   while the returned slice is live.
+    pub unsafe fn as_slice_mut<const OFFSET: usize, const COUNT: usize>(&mut self) -> &mut [T] {
+        build_assert!(
+            OFFSET + COUNT <= N,
+            "Range is out of bounds for the allocation."
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`,
+        //   we've just checked that the range and index is within bounds. The immutability of the
+        //   data is also guaranteed by the safety requirements of the function.
+        // - `OFFSET + COUNT` can't overflow since it is smaller than `N` and we've checked
+        //   that `N` won't overflow early in the constructor.
+        unsafe { core::slice::from_raw_parts_mut(self.start_ptr_mut().add(OFFSET), COUNT) }
+    }
+
+    /// Writes data to the region starting from `OFFSET`. `OFFSET` is in units of `T`, not the
+    /// number of bytes.
+    ///
+    /// # Safety
+    ///
+    /// * Callers must ensure that this call does not race with a read or write to the same region
+    ///   that overlaps with this write.
+    pub unsafe fn write<const OFFSET: usize, const SIZE: usize>(&mut self, src: &[T; SIZE]) {
+        build_assert!(
+            OFFSET + SIZE <= N,
+            "Range is out of bounds for the allocation."
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`
+        //   and we've just checked that the range and index is within bounds.
+        // - `OFFSET + SIZE` can't overflow since it is smaller than `N` and we've checked
+        //   that `N` won't overflow early in the constructor.
+        unsafe {
+            core::ptr::copy_nonoverlapping(
+                src.as_ptr(),
+                self.start_ptr_mut().add(OFFSET),
+                src.len(),
+            )
+        };
+    }
+}
+
 /// Note that the device configured to do DMA must be halted before this object is dropped.
 impl<T: AsBytes + FromBytes, Size: AllocationSize> Drop for CoherentAllocation<T, Size> {
     fn drop(&mut self) {
diff --git a/samples/rust/rust_dma.rs b/samples/rust/rust_dma.rs
index 7a87048575df..97711a99ac8b 100644
--- a/samples/rust/rust_dma.rs
+++ b/samples/rust/rust_dma.rs
@@ -6,7 +6,7 @@
 
 use kernel::{
     device::Core,
-    dma::{CoherentAllocation, DataDirection, Device, DmaMask},
+    dma::{CoherentSlice, DataDirection, Device, DmaMask},
     page, pci,
     prelude::*,
     scatterlist::{Owned, SGTable},
@@ -16,7 +16,7 @@
 #[pin_data(PinnedDrop)]
 struct DmaSampleDriver {
     pdev: ARef<pci::Device>,
-    ca: CoherentAllocation<MyStruct>,
+    ca: CoherentSlice<MyStruct>,
     #[pin]
     sgt: SGTable<Owned<VVec<u8>>>,
 }
@@ -64,8 +64,8 @@ fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, E
             // SAFETY: There are no concurrent calls to DMA allocation and mapping primitives.
             unsafe { pdev.dma_set_mask_and_coherent(mask)? };
 
-            let ca: CoherentAllocation<MyStruct> =
-                CoherentAllocation::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
+            let ca: CoherentSlice<MyStruct> =
+                CoherentSlice::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
 
             for (i, value) in TEST_VALUES.into_iter().enumerate() {
                 kernel::try_dma_write!(ca[i] = MyStruct::new(value.0, value.1))?;

-- 
2.52.0

[PATCH 4/9] rust: dma: simplify try_dma_read! and try_dma_write!

[Eliot Courtney]

Remove unnecessary branch and unify branch order.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/dma.rs | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 43ed0dfdbc08..03fc001eb983 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -875,9 +875,6 @@ macro_rules! try_dma_read {
             }
         })()
     }};
-    ($dma:ident [ $idx:expr ] $($field:tt)* ) => {
-        $crate::try_dma_read!($dma, $idx, $($field)*)
-    };
     ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {
         $crate::try_dma_read!($($dma).*, $idx, $($field)*)
     };
@@ -905,13 +902,7 @@ macro_rules! try_dma_read {
 /// ```
 #[macro_export]
 macro_rules! try_dma_write {
-    ($dma:ident [ $idx:expr ] $($field:tt)*) => {{
-        $crate::try_dma_write!($dma, $idx, $($field)*)
-    }};
-    ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {{
-        $crate::try_dma_write!($($dma).*, $idx, $($field)*)
-    }};
-    ($dma:expr, $idx: expr, = $val:expr) => {
+    ($dma:expr, $idx:expr, = $val:expr) => {
         (|| -> ::core::result::Result<_, $crate::error::Error> {
             let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
             // SAFETY: `try_item_from_index` ensures that `item` is always a valid item.
@@ -919,7 +910,7 @@ macro_rules! try_dma_write {
             ::core::result::Result::Ok(())
         })()
     };
-    ($dma:expr, $idx: expr, $(.$field:ident)* = $val:expr) => {
+    ($dma:expr, $idx:expr, $(.$field:ident)* = $val:expr) => {
         (|| -> ::core::result::Result<_, $crate::error::Error> {
             let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
             // SAFETY: `try_item_from_index` ensures that `item` is always a valid pointer
@@ -932,4 +923,7 @@ macro_rules! try_dma_write {
             ::core::result::Result::Ok(())
         })()
     };
+    ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {{
+        $crate::try_dma_write!($($dma).*, $idx, $($field)*)
+    }};
 }

-- 
2.52.0

[PATCH 5/9] rust: dma: rename try_item_from_index to try_ptr_at

[Eliot Courtney]

This function returns a pointer, so rename it to be clearer about what
it is getting.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/dma.rs | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 03fc001eb983..e4bca7a18ac1 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -559,7 +559,7 @@ pub unsafe fn try_write(&mut self, src: &[T], offset: usize) -> Result {
     /// Public but hidden since it should only be used from [`try_dma_read`] and [`try_dma_write`]
     /// macros.
     #[doc(hidden)]
-    pub fn try_item_from_index(&self, offset: usize) -> Result<*mut T> {
+    pub fn try_ptr_at(&self, offset: usize) -> Result<*mut T> {
         if offset >= self.count {
             return Err(EINVAL);
         }
@@ -863,12 +863,12 @@ unsafe impl<T: AsBytes + FromBytes + Send, Size: AllocationSize> Send
 macro_rules! try_dma_read {
     ($dma:expr, $idx:expr, $($field:tt)*) => {{
         (|| -> ::core::result::Result<_, $crate::error::Error> {
-            let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
-            // SAFETY: `try_item_from_index` ensures that `item` is always a valid pointer
-            // and can be dereferenced. The compiler also further validates the expression
-            // on whether `field` is a member of `item` when expanded by the macro.
+            let ptr = $crate::dma::CoherentAllocation::try_ptr_at(&$dma, $idx)?;
+            // SAFETY: `try_ptr_at` ensures that `ptr` is always a valid pointer and can be
+            // dereferenced. The compiler also further validates the expression on whether `field`
+            // is a member of `ptr` when expanded by the macro.
             unsafe {
-                let ptr_field = ::core::ptr::addr_of!((*item) $($field)*);
+                let ptr_field = ::core::ptr::addr_of!((*ptr) $($field)*);
                 ::core::result::Result::Ok(
                     $crate::dma::CoherentAllocation::field_read(&$dma, ptr_field)
                 )
@@ -904,20 +904,20 @@ macro_rules! try_dma_read {
 macro_rules! try_dma_write {
     ($dma:expr, $idx:expr, = $val:expr) => {
         (|| -> ::core::result::Result<_, $crate::error::Error> {
-            let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
-            // SAFETY: `try_item_from_index` ensures that `item` is always a valid item.
-            unsafe { $crate::dma::CoherentAllocation::field_write(&$dma, item, $val) }
+            let ptr = $crate::dma::CoherentAllocation::try_ptr_at(&$dma, $idx)?;
+            // SAFETY: `try_ptr_at` ensures that `ptr` is always a valid ptr.
+            unsafe { $crate::dma::CoherentAllocation::field_write(&$dma, ptr, $val) }
             ::core::result::Result::Ok(())
         })()
     };
     ($dma:expr, $idx:expr, $(.$field:ident)* = $val:expr) => {
         (|| -> ::core::result::Result<_, $crate::error::Error> {
-            let item = $crate::dma::CoherentAllocation::try_item_from_index(&$dma, $idx)?;
-            // SAFETY: `try_item_from_index` ensures that `item` is always a valid pointer
-            // and can be dereferenced. The compiler also further validates the expression
-            // on whether `field` is a member of `item` when expanded by the macro.
+            let ptr = $crate::dma::CoherentAllocation::try_ptr_at(&$dma, $idx)?;
+            // SAFETY: `try_ptr_at` ensures that `ptr` is always a valid pointer and can be
+            // dereferenced. The compiler also further validates the expression on whether `field`
+            // is a member of `ptr` when expanded by the macro.
             unsafe {
-                let ptr_field = ::core::ptr::addr_of_mut!((*item) $(.$field)*);
+                let ptr_field = ::core::ptr::addr_of_mut!((*ptr) $(.$field)*);
                 $crate::dma::CoherentAllocation::field_write(&$dma, ptr_field, $val)
             }
             ::core::result::Result::Ok(())

-- 
2.52.0

[PATCH 6/9] rust: dma: add dma_read! and dma_write! macros

[Eliot Courtney]

Add dma_read! and dma_write! macros using the new infallible methods
on CoherentArray.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/dma.rs | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 103 insertions(+)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index e4bca7a18ac1..f3920f74583a 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -811,6 +811,24 @@ pub unsafe fn as_slice_mut<const OFFSET: usize, const COUNT: usize>(&mut self) -
             )
         };
     }
+
+    /// Returns a pointer to an element from the region with bounds checking. `OFFSET` is in
+    /// units of `T`, not the number of bytes.
+    ///
+    /// Public but hidden since it should only be used from [`dma_read`] and [`dma_write`] macros.
+    #[doc(hidden)]
+    pub fn ptr_at<const OFFSET: usize>(&self) -> *mut T {
+        build_assert!(
+            OFFSET < N,
+            "Index out of bounds when accessing CoherentArray"
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`
+        // and we've just checked that the range and index is within bounds.
+        // - `OFFSET` can't overflow since it is smaller than `N` and we've checked
+        // that `N` won't overflow early in the constructor.
+        unsafe { self.cpu_addr.as_ptr().add(OFFSET) }
+    }
 }
 
 /// Note that the device configured to do DMA must be halted before this object is dropped.
@@ -927,3 +945,88 @@ macro_rules! try_dma_write {
         $crate::try_dma_write!($($dma).*, $idx, $($field)*)
     }};
 }
+
+/// Reads a field of an item from a [`CoherentArray`] with compile-time bounds checking.
+///
+/// # Examples
+///
+/// ```
+/// use kernel::device::Device;
+/// use kernel::dma::{attrs::*, CoherentArray};
+///
+/// struct MyStruct { field: u32, }
+///
+/// // SAFETY: All bit patterns are acceptable values for `MyStruct`.
+/// unsafe impl kernel::transmute::FromBytes for MyStruct{};
+/// // SAFETY: Instances of `MyStruct` have no uninitialized portions.
+/// unsafe impl kernel::transmute::AsBytes for MyStruct{};
+///
+/// # fn test(alloc: &kernel::dma::CoherentArray<MyStruct, 3>) {
+/// let whole = kernel::dma_read!(alloc[2]);
+/// let field = kernel::dma_read!(alloc[1].field);
+/// # }
+/// ```
+#[macro_export]
+macro_rules! dma_read {
+    ($dma:expr, $idx:expr, $($field:tt)*) => {{
+        (|| {
+            let ptr = $crate::dma::CoherentArray::ptr_at::<$idx>(&$dma);
+            // SAFETY: `ptr_at` ensures that `ptr` is always a valid pointer and can be
+            // dereferenced. The compiler also further validates the expression on whether `field`
+            // is a member of `ptr` when expanded by the macro.
+            unsafe {
+                let ptr_field = ::core::ptr::addr_of!((*ptr) $($field)*);
+                $crate::dma::CoherentAllocation::field_read(&$dma, ptr_field)
+            }
+        })()
+    }};
+    ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {
+        $crate::dma_read!($($dma).*, $idx, $($field)*)
+    };
+}
+
+/// Writes to a field of an item in a [`CoherentArray`] with compile-time bounds checking.
+///
+/// # Examples
+///
+/// ```
+/// use kernel::device::Device;
+/// use kernel::dma::{attrs::*, CoherentArray};
+///
+/// struct MyStruct { member: u32, }
+///
+/// // SAFETY: All bit patterns are acceptable values for `MyStruct`.
+/// unsafe impl kernel::transmute::FromBytes for MyStruct{};
+/// // SAFETY: Instances of `MyStruct` have no uninitialized portions.
+/// unsafe impl kernel::transmute::AsBytes for MyStruct{};
+///
+/// # fn test(alloc: &kernel::dma::CoherentArray<MyStruct, 3>) {
+/// kernel::dma_write!(alloc[2].member = 0xf);
+/// kernel::dma_write!(alloc[1] = MyStruct { member: 0xf });
+/// # }
+/// ```
+#[macro_export]
+macro_rules! dma_write {
+    ($dma:expr, $idx:expr, = $val:expr) => {
+        (|| {
+            let ptr = $crate::dma::CoherentArray::ptr_at::<$idx>(&$dma);
+            // SAFETY: `ptr_at` ensures that `ptr` is always a valid ptr.
+            unsafe { $crate::dma::CoherentAllocation::field_write(&$dma, ptr, $val) }
+        })()
+    };
+    ($dma:expr, $idx:expr, $(.$field:ident)* = $val:expr) => {
+        (|| {
+            let ptr = $crate::dma::CoherentArray::ptr_at::<$idx>(&$dma);
+            // SAFETY: `ptr_at` ensures that `ptr` is always a valid pointer and can be
+            // dereferenced. The compiler also further validates the expression on whether `field`
+            // is a member of `ptr` when expanded by the macro.
+            unsafe {
+                let ptr_field = ::core::ptr::addr_of_mut!((*ptr) $(.$field)*);
+                $crate::dma::CoherentAllocation::field_write(&$dma, ptr_field, $val)
+            }
+        })()
+    };
+    ($($dma:ident).* [ $idx:expr ] $($field:tt)* ) => {{
+        $crate::dma_write!($($dma).*, $idx, $($field)*)
+    }};
+}

-- 
2.52.0

[PATCH 7/9] rust: dma: implement decay from CoherentArray to CoherentSlice

[Eliot Courtney]

Implement Deref, DerefMut, AsRef, AsMut, From for various methods
of decaying CoherentArray to CoherentSlice. This is so statically
sized CoherentArrays can be used as if they were CoherentSlices by
code that doesn't care about knowing the compile time size.

This also helps avoid having to annotate static sizes on types all
the time.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/dma.rs | 48 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index f3920f74583a..25da678c863b 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -12,7 +12,11 @@
     sync::aref::ARef,
     transmute::{AsBytes, FromBytes},
 };
-use core::{marker::PhantomData, ptr::NonNull};
+use core::{
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    ptr::NonNull, //
+};
 
 /// DMA address type.
 ///
@@ -389,7 +393,8 @@ impl<const N: usize> AllocationSize for StaticSize<N> {}
 /// # Allocation size
 ///
 /// [`CoherentAllocation`] is generic over an [`AllocationSize`], which lets it record a compile
-/// time known size (in number of elements of `T`).
+/// time known size (in number of elements of `T`). A statically sized [`CoherentAllocation`] can
+/// decay to a runtime sized one via deref coercion.
 // TODO
 //
 // DMA allocations potentially carry device resources (e.g.IOMMU mappings), hence for soundness
@@ -402,6 +407,7 @@ impl<const N: usize> AllocationSize for StaticSize<N> {}
 //
 // Hence, find a way to revoke the device resources of a `CoherentAllocation`, but not the
 // entire `CoherentAllocation` including the allocated memory itself.
+#[repr(C)]
 pub struct CoherentAllocation<T: AsBytes + FromBytes, Size: AllocationSize = RuntimeSize> {
     dev: ARef<device::Device>,
     dma_handle: DmaAddress,
@@ -857,6 +863,44 @@ unsafe impl<T: AsBytes + FromBytes + Send, Size: AllocationSize> Send
 {
 }
 
+impl<T: AsBytes + FromBytes, const N: usize> Deref for CoherentArray<T, N> {
+    type Target = CoherentSlice<T>;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: `CoherentArray<T, N>` and `CoherentSlice<T>` are both `CoherentAllocation<T, S>`
+        // with different `S: AllocationSize` marker types. Since `AllocationSize` is only stored as
+        // `PhantomData<S>` (a ZST) and CoherentAllocation<T, S> is `repr(C)`, both types have
+        // identical memory layouts.
+        unsafe { &*core::ptr::from_ref(self).cast::<CoherentSlice<T>>() }
+    }
+}
+
+impl<T: AsBytes + FromBytes, const N: usize> DerefMut for CoherentArray<T, N> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        // SAFETY: Same as `Deref::deref`.
+        unsafe { &mut *core::ptr::from_mut(self).cast::<CoherentSlice<T>>() }
+    }
+}
+
+impl<T: AsBytes + FromBytes, const N: usize> AsRef<CoherentSlice<T>> for CoherentArray<T, N> {
+    fn as_ref(&self) -> &CoherentSlice<T> {
+        self
+    }
+}
+
+impl<T: AsBytes + FromBytes, const N: usize> AsMut<CoherentSlice<T>> for CoherentArray<T, N> {
+    fn as_mut(&mut self) -> &mut CoherentSlice<T> {
+        self
+    }
+}
+
+impl<T: AsBytes + FromBytes, const N: usize> From<CoherentArray<T, N>> for CoherentSlice<T> {
+    fn from(array: CoherentArray<T, N>) -> Self {
+        // SAFETY: Same as `Deref::deref`.
+        unsafe { core::mem::transmute(array) }
+    }
+}
+
 /// Reads a field of an item from an allocated region of structs.
 ///
 /// # Examples

-- 
2.52.0

[PATCH 8/9] rust: dma: add CoherentObject for single element allocations

[Eliot Courtney]

Add a type alias for convenience and readability for single element
allocations.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/dma.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 25da678c863b..120fc01e57a3 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -423,6 +423,9 @@ pub struct CoherentAllocation<T: AsBytes + FromBytes, Size: AllocationSize = Run
 /// A coherent DMA allocation for an array of `N` elements.
 pub type CoherentArray<T, const N: usize> = CoherentAllocation<T, StaticSize<N>>;
 
+/// A coherent DMA allocation for a single object.
+pub type CoherentObject<T> = CoherentAllocation<T, StaticSize<1>>;
+
 impl<T: AsBytes + FromBytes, Size: AllocationSize> CoherentAllocation<T, Size> {
     /// Returns the number of elements `T` in this allocation.
     ///

-- 
2.52.0

[PATCH 9/9] gpu: nova-core: migrate to CoherentArray and CoherentObject

[Eliot Courtney]

Migrate to CoherentArray and CoherentObject. This enables removing
a few unwrap()s.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 drivers/gpu/nova-core/gsp.rs      | 46 ++++++++++++++++++++-------------------
 drivers/gpu/nova-core/gsp/boot.rs | 10 ++++-----
 drivers/gpu/nova-core/gsp/cmdq.rs | 24 ++++++++++----------
 drivers/gpu/nova-core/gsp/fw.rs   | 12 +++-------
 4 files changed, 43 insertions(+), 49 deletions(-)

diff --git a/drivers/gpu/nova-core/gsp.rs b/drivers/gpu/nova-core/gsp.rs
index 43bc35fd3b55..2513822d43fb 100644
--- a/drivers/gpu/nova-core/gsp.rs
+++ b/drivers/gpu/nova-core/gsp.rs
@@ -5,13 +5,15 @@
 use kernel::{
     device,
     dma::{
-        CoherentAllocation,
+        CoherentArray,
+        CoherentObject,
+        CoherentSlice,
         DmaAddress, //
     },
+    dma_write,
     pci,
     prelude::*,
-    transmute::AsBytes,
-    try_dma_write, //
+    transmute::AsBytes, //
 };
 
 pub(crate) mod cmdq;
@@ -74,14 +76,14 @@ fn new(start: DmaAddress) -> Result<Self> {
 /// then pp points to index into the buffer where the next logging entry will
 /// be written. Therefore, the logging data is valid if:
 ///   1 <= pp < sizeof(buffer)/sizeof(u64)
-struct LogBuffer(CoherentAllocation<u8>);
+struct LogBuffer(CoherentSlice<u8>);
 
 impl LogBuffer {
     /// Creates a new `LogBuffer` mapped on `dev`.
     fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         const NUM_PAGES: usize = RM_LOG_BUFFER_NUM_PAGES;
 
-        let mut obj = Self(CoherentAllocation::<u8>::alloc_coherent(
+        let mut obj = Self(CoherentSlice::<u8>::alloc_coherent(
             dev,
             NUM_PAGES * GSP_PAGE_SIZE,
             GFP_KERNEL | __GFP_ZERO,
@@ -100,11 +102,14 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
     }
 }
 
+/// Number of `LibosMemoryRegionInitArgument` entries that fit in a GSP page.
+const LIBOS_REGION_SIZE: usize = GSP_PAGE_SIZE / size_of::<LibosMemoryRegionInitArgument>();
+
 /// GSP runtime data.
 #[pin_data]
 pub(crate) struct Gsp {
     /// Libos arguments.
-    pub(crate) libos: CoherentAllocation<LibosMemoryRegionInitArgument>,
+    pub(crate) libos: CoherentArray<LibosMemoryRegionInitArgument, LIBOS_REGION_SIZE>,
     /// Init log buffer.
     loginit: LogBuffer,
     /// Interrupts log buffer.
@@ -114,40 +119,37 @@ pub(crate) struct Gsp {
     /// Command queue.
     pub(crate) cmdq: Cmdq,
     /// RM arguments.
-    rmargs: CoherentAllocation<GspArgumentsCached>,
+    rmargs: CoherentObject<GspArgumentsCached>,
 }
 
 impl Gsp {
     // Creates an in-place initializer for a `Gsp` manager for `pdev`.
     pub(crate) fn new(pdev: &pci::Device<device::Bound>) -> Result<impl PinInit<Self, Error>> {
         let dev = pdev.as_ref();
-        let libos = CoherentAllocation::<LibosMemoryRegionInitArgument>::alloc_coherent(
-            dev,
-            GSP_PAGE_SIZE / size_of::<LibosMemoryRegionInitArgument>(),
-            GFP_KERNEL | __GFP_ZERO,
-        )?;
+        let libos =
+            CoherentArray::<LibosMemoryRegionInitArgument, LIBOS_REGION_SIZE>::alloc_coherent(
+                dev,
+                GFP_KERNEL | __GFP_ZERO,
+            )?;
 
         // Initialise the logging structures. The OpenRM equivalents are in:
         // _kgspInitLibosLoggingStructures (allocates memory for buffers)
         // kgspSetupLibosInitArgs_IMPL (creates pLibosInitArgs[] array)
         let loginit = LogBuffer::new(dev)?;
-        try_dma_write!(libos[0] = LibosMemoryRegionInitArgument::new("LOGINIT", &loginit.0))?;
+        dma_write!(libos[0] = LibosMemoryRegionInitArgument::new("LOGINIT", &loginit.0));
 
         let logintr = LogBuffer::new(dev)?;
-        try_dma_write!(libos[1] = LibosMemoryRegionInitArgument::new("LOGINTR", &logintr.0))?;
+        dma_write!(libos[1] = LibosMemoryRegionInitArgument::new("LOGINTR", &logintr.0));
 
         let logrm = LogBuffer::new(dev)?;
-        try_dma_write!(libos[2] = LibosMemoryRegionInitArgument::new("LOGRM", &logrm.0))?;
+        dma_write!(libos[2] = LibosMemoryRegionInitArgument::new("LOGRM", &logrm.0));
 
         let cmdq = Cmdq::new(dev)?;
 
-        let rmargs = CoherentAllocation::<GspArgumentsCached>::alloc_coherent(
-            dev,
-            1,
-            GFP_KERNEL | __GFP_ZERO,
-        )?;
-        try_dma_write!(rmargs[0] = fw::GspArgumentsCached::new(&cmdq))?;
-        try_dma_write!(libos[3] = LibosMemoryRegionInitArgument::new("RMARGS", &rmargs))?;
+        let rmargs =
+            CoherentObject::<GspArgumentsCached>::alloc_coherent(dev, GFP_KERNEL | __GFP_ZERO)?;
+        dma_write!(rmargs[0] = fw::GspArgumentsCached::new(&cmdq));
+        dma_write!(libos[3] = LibosMemoryRegionInitArgument::new("RMARGS", &rmargs));
 
         Ok(try_pin_init!(Self {
             libos,
diff --git a/drivers/gpu/nova-core/gsp/boot.rs b/drivers/gpu/nova-core/gsp/boot.rs
index 69e2fb064220..7888a39356f3 100644
--- a/drivers/gpu/nova-core/gsp/boot.rs
+++ b/drivers/gpu/nova-core/gsp/boot.rs
@@ -2,12 +2,12 @@
 
 use kernel::{
     device,
-    dma::CoherentAllocation,
+    dma::CoherentObject,
+    dma_write,
     io::poll::read_poll_timeout,
     pci,
     prelude::*,
-    time::Delta,
-    try_dma_write, //
+    time::Delta, //
 };
 
 use crate::{
@@ -159,8 +159,8 @@ pub(crate) fn boot(
         )?;
 
         let wpr_meta =
-            CoherentAllocation::<GspFwWprMeta>::alloc_coherent(dev, 1, GFP_KERNEL | __GFP_ZERO)?;
-        try_dma_write!(wpr_meta[0] = GspFwWprMeta::new(&gsp_fw, &fb_layout))?;
+            CoherentObject::<GspFwWprMeta>::alloc_coherent(dev, GFP_KERNEL | __GFP_ZERO)?;
+        dma_write!(wpr_meta[0] = GspFwWprMeta::new(&gsp_fw, &fb_layout));
 
         self.cmdq
             .send_command(bar, commands::SetSystemInfo::new(pdev))?;
diff --git a/drivers/gpu/nova-core/gsp/cmdq.rs b/drivers/gpu/nova-core/gsp/cmdq.rs
index 9c94f4c6ff6d..845c9e176b93 100644
--- a/drivers/gpu/nova-core/gsp/cmdq.rs
+++ b/drivers/gpu/nova-core/gsp/cmdq.rs
@@ -12,9 +12,10 @@
 use kernel::{
     device,
     dma::{
-        CoherentAllocation,
+        CoherentObject,
         DmaAddress, //
     },
+    dma_write,
     io::poll::read_poll_timeout,
     prelude::*,
     sync::aref::ARef,
@@ -22,8 +23,7 @@
     transmute::{
         AsBytes,
         FromBytes, //
-    },
-    try_dma_write, //
+    }, //
 };
 
 use crate::{
@@ -191,7 +191,7 @@ unsafe impl FromBytes for GspMem {}
 ///   pointer and the GSP read pointer. This region is returned by [`Self::driver_write_area`].
 /// * The driver owns (i.e. can read from) the part of the GSP message queue between the CPU read
 ///   pointer and the GSP write pointer. This region is returned by [`Self::driver_read_area`].
-struct DmaGspMem(CoherentAllocation<GspMem>);
+struct DmaGspMem(CoherentObject<GspMem>);
 
 impl DmaGspMem {
     /// Allocate a new instance and map it for `dev`.
@@ -199,13 +199,11 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         const MSGQ_SIZE: u32 = num::usize_into_u32::<{ size_of::<Msgq>() }>();
         const RX_HDR_OFF: u32 = num::usize_into_u32::<{ mem::offset_of!(Msgq, rx) }>();
 
-        let gsp_mem =
-            CoherentAllocation::<GspMem>::alloc_coherent(dev, 1, GFP_KERNEL | __GFP_ZERO)?;
-        try_dma_write!(gsp_mem[0].ptes = PteArray::new(gsp_mem.dma_handle())?)?;
-        try_dma_write!(
-            gsp_mem[0].cpuq.tx = MsgqTxHeader::new(MSGQ_SIZE, RX_HDR_OFF, MSGQ_NUM_PAGES)
-        )?;
-        try_dma_write!(gsp_mem[0].cpuq.rx = MsgqRxHeader::new())?;
+        let gsp_mem = CoherentObject::<GspMem>::alloc_coherent(dev, GFP_KERNEL | __GFP_ZERO)?;
+        let ptes = PteArray::new(gsp_mem.dma_handle())?;
+        dma_write!(gsp_mem[0].ptes = ptes);
+        dma_write!(gsp_mem[0].cpuq.tx = MsgqTxHeader::new(MSGQ_SIZE, RX_HDR_OFF, MSGQ_NUM_PAGES));
+        dma_write!(gsp_mem[0].cpuq.rx = MsgqRxHeader::new());
 
         Ok(Self(gsp_mem))
     }
@@ -223,7 +221,7 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         // - The `CoherentAllocation` contains exactly one object.
         // - We will only access the driver-owned part of the shared memory.
         // - Per the safety statement of the function, no concurrent access will be performed.
-        let gsp_mem = &mut unsafe { self.0.try_as_slice_mut(0, 1) }.unwrap()[0];
+        let gsp_mem = &mut unsafe { self.0.as_slice_mut::<0, 1>() }[0];
         // PANIC: per the invariant of `cpu_write_ptr`, `tx` is `<= MSGQ_NUM_PAGES`.
         let (before_tx, after_tx) = gsp_mem.cpuq.msgq.data.split_at_mut(tx);
 
@@ -258,7 +256,7 @@ fn new(dev: &device::Device<device::Bound>) -> Result<Self> {
         // - The `CoherentAllocation` contains exactly one object.
         // - We will only access the driver-owned part of the shared memory.
         // - Per the safety statement of the function, no concurrent access will be performed.
-        let gsp_mem = &unsafe { self.0.try_as_slice(0, 1) }.unwrap()[0];
+        let gsp_mem = &unsafe { self.0.as_slice::<0, 1>() }[0];
         // PANIC: per the invariant of `cpu_read_ptr`, `xx` is `<= MSGQ_NUM_PAGES`.
         let (before_rx, after_rx) = gsp_mem.gspq.msgq.data.split_at(rx);
 
diff --git a/drivers/gpu/nova-core/gsp/fw.rs b/drivers/gpu/nova-core/gsp/fw.rs
index caeb0d251fe5..1877b727cc22 100644
--- a/drivers/gpu/nova-core/gsp/fw.rs
+++ b/drivers/gpu/nova-core/gsp/fw.rs
@@ -9,17 +9,14 @@
 use core::ops::Range;
 
 use kernel::{
-    dma::CoherentAllocation,
+    dma::CoherentSlice,
     fmt,
     prelude::*,
     ptr::{
         Alignable,
         Alignment, //
     },
-    sizes::{
-        SZ_128K,
-        SZ_1M, //
-    },
+    sizes::{SZ_128K, SZ_1M},
     transmute::{
         AsBytes,
         FromBytes, //
@@ -652,10 +649,7 @@ unsafe impl AsBytes for LibosMemoryRegionInitArgument {}
 unsafe impl FromBytes for LibosMemoryRegionInitArgument {}
 
 impl LibosMemoryRegionInitArgument {
-    pub(crate) fn new<A: AsBytes + FromBytes>(
-        name: &'static str,
-        obj: &CoherentAllocation<A>,
-    ) -> Self {
+    pub(crate) fn new<A: AsBytes + FromBytes>(name: &'static str, obj: &CoherentSlice<A>) -> Self {
         /// Generates the `ID8` identifier required for some GSP objects.
         fn id8(name: &str) -> u64 {
             let mut bytes = [0u8; core::mem::size_of::<u64>()];

-- 
2.52.0

Re: [PATCH 6/9] rust: dma: add dma_read! and dma_write! macros

[Alice Ryhl]

On Fri, Jan 30, 2026 at 9:35 AM Eliot Courtney <ecourtney@nvidia.com> wrote:
>
> Add dma_read! and dma_write! macros using the new infallible methods
> on CoherentArray.
>
> Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
> ---
>  rust/kernel/dma.rs | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 103 insertions(+)
>
> diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
> index e4bca7a18ac1..f3920f74583a 100644
> --- a/rust/kernel/dma.rs
> +++ b/rust/kernel/dma.rs
> @@ -811,6 +811,24 @@ pub unsafe fn as_slice_mut<const OFFSET: usize, const COUNT: usize>(&mut self) -
>              )
>          };
>      }
> +
> +    /// Returns a pointer to an element from the region with bounds checking. `OFFSET` is in
> +    /// units of `T`, not the number of bytes.
> +    ///
> +    /// Public but hidden since it should only be used from [`dma_read`] and [`dma_write`] macros.
> +    #[doc(hidden)]
> +    pub fn ptr_at<const OFFSET: usize>(&self) -> *mut T {
> +        build_assert!(
> +            OFFSET < N,
> +            "Index out of bounds when accessing CoherentArray"
> +        );

This build assert does not depend on runtime values, so it can be
written as a const block:
const { assert!(OFFSET < N, "Index out of bounds when accessing
CoherentArray"); }

Alice

Re: [PATCH 0/9] rust: dma: add CoherentArray for compile-time sized allocations

[Danilo Krummrich]

(Cc: Lyude)

On Fri Jan 30, 2026 at 9:34 AM CET, Eliot Courtney wrote:
> This series extends the DMA coherent allocation API to support compile-time
> known sizes. This lets bounds checking to be moved from runtime to build
> time, which is useful to avoid runtime panics from index typos. It also
> removes the need for a Result return type in some places.
>
> The compile time size is specified via a marker type: StaticSize<N>.
> Statically sized allocations can decay to runtime sized ones via deref
> coercion for code that doesn't need to know the size at compile time, or to
> avoid having to carry around extra type parameters. The implementation
> follows a similar pattern to Device/DeviceContext.
>
> The series defines three type aliases: CoherentSlice<T> (for runtime size),
> CoherentArray<T, N> (for compile-time size N), and CoherentObject<T> (for
> single object allocations). It also adds infallible dma_read!/dma_write!
> macros and methods to CoherentArray, while prefixing the existing fallible
> methods and macros with `try_`.
>
> The macros keep the same syntax (i.e.
> coherent_allocation[index].optional_fields = expression) even for
> CoherentObject, because the [] syntax is needed to know where to split the
> actual CoherentAllocation object from the fields. This means that
> CoherentObject is indexed with [0] in dma_write!/dma_read! macros. The
> alternative is defining a separate macro for single object access, but it
> still would need a way to delineate between the allocation and the fields,
> perhaps by using commas (dma_read_obj!(object, fields),
> dma_write_obj!(object, fields, value)). This would be inconsistent with the
> array/slice syntax.

We've just generalized I/O to support arbitrary I/O backends (busses, backing
storage, etc.).

With this we can wire up the I/O traits to DMA and generalize the dma_read() and
dma_write() macros accordingly. I.e. we can extend the I/O traits with
field_write() and field_read().

(Lyude is going to work on this as a more integrated alternative to iosys_map.
It would be good to align with her regarding this work.)

This has the advantage that we don't have to duplicate all this infrastructure
for I/O memory, DMA, etc.

I also think that CoherentSlice is too specific of a type. I'd rather have a
generic type, maybe UnsafeSlice or IoSlice, that just uses the I/O backend for
accesses.

Re: [PATCH 0/9] rust: dma: add CoherentArray for compile-time sized allocations

[Alexandre Courbot]

On Sat Jan 31, 2026 at 9:27 PM JST, Danilo Krummrich wrote:
> (Cc: Lyude)
>
> On Fri Jan 30, 2026 at 9:34 AM CET, Eliot Courtney wrote:
>> This series extends the DMA coherent allocation API to support compile-time
>> known sizes. This lets bounds checking to be moved from runtime to build
>> time, which is useful to avoid runtime panics from index typos. It also
>> removes the need for a Result return type in some places.
>>
>> The compile time size is specified via a marker type: StaticSize<N>.
>> Statically sized allocations can decay to runtime sized ones via deref
>> coercion for code that doesn't need to know the size at compile time, or to
>> avoid having to carry around extra type parameters. The implementation
>> follows a similar pattern to Device/DeviceContext.
>>
>> The series defines three type aliases: CoherentSlice<T> (for runtime size),
>> CoherentArray<T, N> (for compile-time size N), and CoherentObject<T> (for
>> single object allocations). It also adds infallible dma_read!/dma_write!
>> macros and methods to CoherentArray, while prefixing the existing fallible
>> methods and macros with `try_`.
>>
>> The macros keep the same syntax (i.e.
>> coherent_allocation[index].optional_fields = expression) even for
>> CoherentObject, because the [] syntax is needed to know where to split the
>> actual CoherentAllocation object from the fields. This means that
>> CoherentObject is indexed with [0] in dma_write!/dma_read! macros. The
>> alternative is defining a separate macro for single object access, but it
>> still would need a way to delineate between the allocation and the fields,
>> perhaps by using commas (dma_read_obj!(object, fields),
>> dma_write_obj!(object, fields, value)). This would be inconsistent with the
>> array/slice syntax.
>
> We've just generalized I/O to support arbitrary I/O backends (busses, backing
> storage, etc.).
>
> With this we can wire up the I/O traits to DMA and generalize the dma_read() and
> dma_write() macros accordingly. I.e. we can extend the I/O traits with
> field_write() and field_read().

With the caveat that the I/O traits for now only support accessing
primitive types; is the plan to add a function to read any type
implementing `FromBytes`?

>
> (Lyude is going to work on this as a more integrated alternative to iosys_map.
> It would be good to align with her regarding this work.)

Heads up, I am also doing some plumbing in `io.rs` related to the
register macro. Maybe we should have a thread on Zulip to discuss what
everyone is working on.

>
> This has the advantage that we don't have to duplicate all this infrastructure
> for I/O memory, DMA, etc.
>
> I also think that CoherentSlice is too specific of a type. I'd rather have a
> generic type, maybe UnsafeSlice or IoSlice, that just uses the I/O backend for
> accesses.

For me the main appeal of this patchset is that it provides a way to
work infallibly with a single object or a fixed-size array. I hope
that's something we can preserve.

Re: [PATCH 0/9] rust: dma: add CoherentArray for compile-time sized allocations

[Danilo Krummrich]

(Cc: Lyude)

I assume something odd is going on with your mail client for people that have
been added to a thread later on?

On Sat Jan 31, 2026 at 2:16 PM CET, Alexandre Courbot wrote:
> On Sat Jan 31, 2026 at 9:27 PM JST, Danilo Krummrich wrote:
>> We've just generalized I/O to support arbitrary I/O backends (busses, backing
>> storage, etc.).
>>
>> With this we can wire up the I/O traits to DMA and generalize the dma_read() and
>> dma_write() macros accordingly. I.e. we can extend the I/O traits with
>> field_write() and field_read().
>
> With the caveat that the I/O traits for now only support accessing
> primitive types; is the plan to add a function to read any type
> implementing `FromBytes`?

That's exactly what I say above: generalize the dma_read!() and dma_write!()
macros by adding field_write() and field_read() to the I/O traits. :)

For reference, this is where I brought this up originally [1].

[1] https://lore.kernel.org/all/DFOP5BY09539.AFY5L5FV1HNV@kernel.org/

>> (Lyude is going to work on this as a more integrated alternative to iosys_map.
>> It would be good to align with her regarding this work.)
>
> Heads up, I am also doing some plumbing in `io.rs` related to the
> register macro. Maybe we should have a thread on Zulip to discuss what
> everyone is working on.

Done!

Link: https://rust-for-linux.zulipchat.com/#narrow/channel/288089-General/topic/Generic.20I.2FO.20backends/with/571198078

>> This has the advantage that we don't have to duplicate all this infrastructure
>> for I/O memory, DMA, etc.
>>
>> I also think that CoherentSlice is too specific of a type. I'd rather have a
>> generic type, maybe UnsafeSlice or IoSlice, that just uses the I/O backend for
>> accesses.
>
> For me the main appeal of this patchset is that it provides a way to
> work infallibly with a single object or a fixed-size array. I hope
> that's something we can preserve.

Of course, the generic I/O backend infrastructure is based on the distinction
between compile-time and run-time.

Re: [PATCH 0/9] rust: dma: add CoherentArray for compile-time sized allocations

[Gary Guo]

On Fri Jan 30, 2026 at 8:34 AM GMT, Eliot Courtney wrote:
> This series extends the DMA coherent allocation API to support compile-time
> known sizes. This lets bounds checking to be moved from runtime to build
> time, which is useful to avoid runtime panics from index typos. It also
> removes the need for a Result return type in some places.
>
> The compile time size is specified via a marker type: StaticSize<N>.
> Statically sized allocations can decay to runtime sized ones via deref
> coercion for code that doesn't need to know the size at compile time, or to
> avoid having to carry around extra type parameters. The implementation
> follows a similar pattern to Device/DeviceContext.
>
> The series defines three type aliases: CoherentSlice<T> (for runtime size),
> CoherentArray<T, N> (for compile-time size N), and CoherentObject<T> (for
> single object allocations). It also adds infallible dma_read!/dma_write!
> macros and methods to CoherentArray, while prefixing the existing fallible
> methods and macros with `try_`.
>
> The macros keep the same syntax (i.e.
> coherent_allocation[index].optional_fields = expression) even for
> CoherentObject, because the [] syntax is needed to know where to split the
> actual CoherentAllocation object from the fields. This means that
> CoherentObject is indexed with [0] in dma_write!/dma_read! macros. The
> alternative is defining a separate macro for single object access, but it
> still would need a way to delineate between the allocation and the fields,
> perhaps by using commas (dma_read_obj!(object, fields),
> dma_write_obj!(object, fields, value)). This would be inconsistent with the
> array/slice syntax.
>
> The last patch in the series may be useful as an example of what this
> looks like to use. Also, there is probably a better name than
> CoherentSlice. I found that specifying a default of RuntimeSize on
> CoherentAllocation stopped the compiler from being able to resolve
> which alloc_attrs to call in usages like e.g.
> CoherentAllocation<u8>::alloc_attrs. Also, we probably want to encourage
> people to use the statically sized one if possible, so it may be nice to
> avoid defaulting CoherentAllocation to RuntimeSize.

I've already posted an example on Zulip but for visibility I'll post it here
too:

I think the design should be `CoherentObject<T: ?Sized>` so that if you need a
`CoherentArray`, it's `CoherentObject<[T; N]>`, and `CoherentSlice<T>` is
`CoherentObject<[T]>`. The existing `Io` that has a fixed minimum size but
variable actual size can be abstracted as a new type.

Link: https://rust-for-linux.zulipchat.com/#narrow/channel/288089-General/topic/Generic.20I.2FO.20backends/near/571228593

Best,
Gary

>
> Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
> ---
> Eliot Courtney (9):
>       rust: dma: rename CoherentAllocation fallible methods
>       rust: dma: parameterize CoherentAllocation with AllocationSize
>       rust: dma: add CoherentArray for compile-time sized allocations
>       rust: dma: simplify try_dma_read! and try_dma_write!
>       rust: dma: rename try_item_from_index to try_ptr_at
>       rust: dma: add dma_read! and dma_write! macros
>       rust: dma: implement decay from CoherentArray to CoherentSlice
>       rust: dma: add CoherentObject for single element allocations
>       gpu: nova-core: migrate to CoherentArray and CoherentObject
>
>  drivers/gpu/nova-core/dma.rs            |  10 +-
>  drivers/gpu/nova-core/falcon.rs         |   2 +-
>  drivers/gpu/nova-core/firmware/fwsec.rs |   4 +-
>  drivers/gpu/nova-core/gsp.rs            |  44 +--
>  drivers/gpu/nova-core/gsp/boot.rs       |   6 +-
>  drivers/gpu/nova-core/gsp/cmdq.rs       |  20 +-
>  drivers/gpu/nova-core/gsp/fw.rs         |  12 +-
>  rust/kernel/dma.rs                      | 555 +++++++++++++++++++++++++-------
>  samples/rust/rust_dma.rs                |  14 +-
>  9 files changed, 489 insertions(+), 178 deletions(-)
> ---
> base-commit: c71257394bc9c59ea727803f6e55e83fe63db74e
> change-id: 20260128-coherent-array-0321eb723d4c
>
> Best regards,