Now that we have a DeviceContext that we can use to represent whether a
Device is known to have been registered, we can make it so that drivers can
create their Devices but wait until the registration phase to assign their
private data to the Device. This is desirable as some drivers need to make
use of the DRM device early on before finalizing their private driver data.

As such, this change makes it so that the driver's private data can
currently only be accessed through Device<T, Registered> types and not
Device<T, Uninit>.

Signed-off-by: Lyude Paul <lyude@redhat.com>
---
 drivers/gpu/drm/nova/driver.rs |  4 +--
 drivers/gpu/drm/tyr/driver.rs  |  4 +--
 rust/kernel/drm/device.rs      | 63 ++++++++++++++++++++--------------
 rust/kernel/drm/driver.rs      | 19 ++++++++--
 4 files changed, 58 insertions(+), 32 deletions(-)

diff --git a/drivers/gpu/drm/nova/driver.rs b/drivers/gpu/drm/nova/driver.rs
index 99d6841b69cbc..8cea5f68c3b04 100644
--- a/drivers/gpu/drm/nova/driver.rs
+++ b/drivers/gpu/drm/nova/driver.rs
@@ -56,8 +56,8 @@ impl auxiliary::Driver for NovaDriver {
     fn probe(adev: &auxiliary::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
         let data = try_pin_init!(NovaData { adev: adev.into() });
 
-        let drm = drm::UnregisteredDevice::<Self>::new(adev.as_ref(), data)?;
-        let drm = drm::Registration::new_foreign_owned(drm, adev.as_ref(), 0)?;
+        let drm = drm::UnregisteredDevice::<Self>::new(adev.as_ref())?;
+        let drm = drm::Registration::new_foreign_owned(drm, adev.as_ref(), data, 0)?;
 
         Ok(Self { drm: drm.into() })
     }
diff --git a/drivers/gpu/drm/tyr/driver.rs b/drivers/gpu/drm/tyr/driver.rs
index ac265a60f5667..e73c56659ea75 100644
--- a/drivers/gpu/drm/tyr/driver.rs
+++ b/drivers/gpu/drm/tyr/driver.rs
@@ -133,8 +133,8 @@ fn probe(
                 gpu_info,
         });
 
-        let tdev = drm::UnregisteredDevice::<Self>::new(pdev.as_ref(), data)?;
-        let tdev = drm::driver::Registration::new_foreign_owned(tdev, pdev.as_ref(), 0)?;
+        let tdev = drm::UnregisteredDevice::<Self>::new(pdev.as_ref())?;
+        let tdev = drm::driver::Registration::new_foreign_owned(tdev, pdev.as_ref(), data, 0)?;
 
         let driver = TyrDriver {
             device: tdev.into(),
diff --git a/rust/kernel/drm/device.rs b/rust/kernel/drm/device.rs
index f529bc7fc4032..0e81957cf8c28 100644
--- a/rust/kernel/drm/device.rs
+++ b/rust/kernel/drm/device.rs
@@ -26,13 +26,15 @@
 };
 use core::{
     alloc::Layout,
+    cell::UnsafeCell,
     marker::PhantomData,
-    mem,
-    ops::Deref,
-    ptr::{
+    mem::{
         self,
-        NonNull, //
+        MaybeUninit, //
     },
+    ops::Deref,
+    ptr::NonNull,
+    sync::atomic::*,
 };
 
 #[cfg(CONFIG_DRM_LEGACY)]
@@ -141,7 +143,7 @@ impl DeviceContext for Uninit {}
 ///
 /// The device in `self.0` is guaranteed to be a newly created [`Device`] that has not yet been
 /// registered with userspace until this type is dropped.
-pub struct UnregisteredDevice<T: drm::Driver>(ARef<Device<T, Uninit>>, NotThreadSafe);
+pub struct UnregisteredDevice<T: drm::Driver>(pub(crate) ARef<Device<T, Uninit>>, NotThreadSafe);
 
 impl<T: drm::Driver> Deref for UnregisteredDevice<T> {
     type Target = Device<T, Uninit>;
@@ -188,7 +190,7 @@ impl<T: drm::Driver> UnregisteredDevice<T> {
     /// Create a new `UnregisteredDevice` for a `drm::Driver`.
     ///
     /// This can be used to create a [`Registration`](kernel::drm::Registration).
-    pub fn new(dev: &device::Device, data: impl PinInit<T::Data, Error>) -> Result<Self> {
+    pub fn new(dev: &device::Device) -> Result<Self> {
         // `__drm_dev_alloc` uses `kmalloc()` to allocate memory, hence ensure a `kmalloc()`
         // compatible `Layout`.
         let layout = Kmalloc::aligned_layout(Layout::new::<Self>());
@@ -207,22 +209,6 @@ pub fn new(dev: &device::Device, data: impl PinInit<T::Data, Error>) -> Result<S
         .cast();
         let raw_drm = NonNull::new(from_err_ptr(raw_drm)?).ok_or(ENOMEM)?;
 
-        // SAFETY: `raw_drm` is a valid pointer to `Self`.
-        let raw_data = unsafe { ptr::addr_of_mut!((*raw_drm.as_ptr()).data) };
-
-        // SAFETY:
-        // - `raw_data` is a valid pointer to uninitialized memory.
-        // - `raw_data` will not move until it is dropped.
-        unsafe { data.__pinned_init(raw_data) }.inspect_err(|_| {
-            // SAFETY: `raw_drm` is a valid pointer to `Self`, given that `__drm_dev_alloc` was
-            // successful.
-            let drm_dev = unsafe { Device::into_drm_device(raw_drm) };
-
-            // SAFETY: `__drm_dev_alloc()` was successful, hence `drm_dev` must be valid and the
-            // refcount must be non-zero.
-            unsafe { bindings::drm_dev_put(drm_dev) };
-        })?;
-
         // SAFETY: The reference count is one, and now we take ownership of that reference as a
         // `drm::Device`.
         // INVARIANT: We just created the device above, but have yet to call `drm_dev_register`.
@@ -254,7 +240,15 @@ pub fn new(dev: &device::Device, data: impl PinInit<T::Data, Error>) -> Result<S
 #[repr(C)]
 pub struct Device<T: drm::Driver, C: DeviceContext = Registered> {
     dev: Opaque<bindings::drm_device>,
-    data: T::Data,
+
+    /// Keeps track of whether we've initialized the device data yet.
+    pub(crate) data_is_init: AtomicBool,
+
+    /// The Driver's private data.
+    ///
+    /// This must only be written to from [`drm::Registration::new`].
+    pub(crate) data: UnsafeCell<MaybeUninit<T::Data>>,
+
     _ctx: PhantomData<C>,
 }
 
@@ -305,6 +299,21 @@ extern "C" fn release(ptr: *mut bindings::drm_device) {
         // SAFETY: `ptr` is a valid pointer to a `struct drm_device` and embedded in `Self`.
         let this = unsafe { Self::from_drm_device(ptr) };
 
+        {
+            // SAFETY:
+            // - Since we are in release(), we are guaranteed that no one else has access to `this`.
+            // - We confirmed above that `this` is a valid pointer to an initialized `Self`.
+            let this = unsafe { &mut *this };
+            if this.data_is_init.load(Ordering::Relaxed) {
+                // SAFETY:
+                // - Since we are in release(), we are guaranteed that no one else has access to
+                //   `this`.
+                // - We checked that the data is initialized above.
+                // - We do not use `data` any point after calling this function.
+                unsafe { (&mut *this.data.get()).assume_init_drop() };
+            }
+        }
+
         // SAFETY:
         // - When `release` runs it is guaranteed that there is no further access to `this`.
         // - `this` is valid for dropping.
@@ -323,11 +332,15 @@ pub(crate) unsafe fn assume_ctx<NewCtx: DeviceContext>(&self) -> &Device<T, NewC
     }
 }
 
-impl<T: drm::Driver, C: DeviceContext> Deref for Device<T, C> {
+impl<T: drm::Driver> Deref for Device<T, Registered> {
     type Target = T::Data;
 
     fn deref(&self) -> &Self::Target {
-        &self.data
+        // SAFETY:
+        // - `data` is initialized before any `Device`s with the `Registered` context are available
+        //   to the user.
+        // - `data` is only written to once in `Registration::new()`, so this read will never race.
+        unsafe { (&*self.data.get()).assume_init_ref() }
     }
 }
 
diff --git a/rust/kernel/drm/driver.rs b/rust/kernel/drm/driver.rs
index da9f1bfef1f14..a16605b407159 100644
--- a/rust/kernel/drm/driver.rs
+++ b/rust/kernel/drm/driver.rs
@@ -15,7 +15,8 @@
 };
 use core::{
     mem,
-    ptr::NonNull, //
+    ptr::NonNull,
+    sync::atomic::*, //
 };
 
 /// Driver use the GEM memory manager. This should be set for all modern drivers.
@@ -127,7 +128,18 @@ pub trait Driver {
 pub struct Registration<T: Driver>(ARef<drm::Device<T>>);
 
 impl<T: Driver> Registration<T> {
-    fn new(drm: drm::UnregisteredDevice<T>, flags: usize) -> Result<Self> {
+    fn new(
+        drm: drm::UnregisteredDevice<T>,
+        data: impl PinInit<T::Data, Error>,
+        flags: usize,
+    ) -> Result<Self> {
+        // SAFETY:
+        // - `raw_data` is a valid pointer to uninitialized memory.
+        // - `raw_data` will not move until it is dropped.
+        unsafe { data.__pinned_init(drm.0.data.get().cast()) }?;
+
+        drm.data_is_init.store(true, Ordering::Relaxed);
+
         // SAFETY: `drm.as_raw()` is valid by the invariants of `drm::Device`.
         to_result(unsafe { bindings::drm_dev_register(drm.as_raw(), flags) })?;
 
@@ -150,6 +162,7 @@ fn new(drm: drm::UnregisteredDevice<T>, flags: usize) -> Result<Self> {
     pub fn new_foreign_owned<'a>(
         drm: drm::UnregisteredDevice<T>,
         dev: &'a device::Device<device::Bound>,
+        data: impl PinInit<T::Data, Error>,
         flags: usize,
     ) -> Result<&'a drm::Device<T>>
     where
@@ -160,7 +173,7 @@ pub fn new_foreign_owned<'a>(
             return Err(EINVAL);
         }
 
-        let reg = Registration::<T>::new(drm, flags)?;
+        let reg = Registration::<T>::new(drm, data, flags)?;
         let drm = NonNull::from(reg.device());
 
         devres::register(dev, reg, GFP_KERNEL)?;
-- 
2.52.0

Hi Lyude,

> On 22 Jan 2026, at 19:46, Lyude Paul <lyude@redhat.com> wrote:
> 
> Now that we have a DeviceContext that we can use to represent whether a
> Device is known to have been registered, we can make it so that drivers can
> create their Devices but wait until the registration phase to assign their
> private data to the Device. This is desirable as some drivers need to make
> use of the DRM device early on before finalizing their private driver data.
> 
> As such, this change makes it so that the driver's private data can
> currently only be accessed through Device<T, Registered> types and not
> Device<T, Uninit>.
> 
> Signed-off-by: Lyude Paul <lyude@redhat.com>
> ---
> drivers/gpu/drm/nova/driver.rs |  4 +--
> drivers/gpu/drm/tyr/driver.rs  |  4 +--
> rust/kernel/drm/device.rs      | 63 ++++++++++++++++++++--------------
> rust/kernel/drm/driver.rs      | 19 ++++++++--
> 4 files changed, 58 insertions(+), 32 deletions(-)
> 
> diff --git a/drivers/gpu/drm/nova/driver.rs b/drivers/gpu/drm/nova/driver.rs
> index 99d6841b69cbc..8cea5f68c3b04 100644
> --- a/drivers/gpu/drm/nova/driver.rs
> +++ b/drivers/gpu/drm/nova/driver.rs
> @@ -56,8 +56,8 @@ impl auxiliary::Driver for NovaDriver {
>     fn probe(adev: &auxiliary::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
>         let data = try_pin_init!(NovaData { adev: adev.into() });
> 
> -        let drm = drm::UnregisteredDevice::<Self>::new(adev.as_ref(), data)?;
> -        let drm = drm::Registration::new_foreign_owned(drm, adev.as_ref(), 0)?;
> +        let drm = drm::UnregisteredDevice::<Self>::new(adev.as_ref())?;
> +        let drm = drm::Registration::new_foreign_owned(drm, adev.as_ref(), data, 0)?;
> 
>         Ok(Self { drm: drm.into() })
>     }
> diff --git a/drivers/gpu/drm/tyr/driver.rs b/drivers/gpu/drm/tyr/driver.rs
> index ac265a60f5667..e73c56659ea75 100644
> --- a/drivers/gpu/drm/tyr/driver.rs
> +++ b/drivers/gpu/drm/tyr/driver.rs
> @@ -133,8 +133,8 @@ fn probe(
>                 gpu_info,
>         });
> 
> -        let tdev = drm::UnregisteredDevice::<Self>::new(pdev.as_ref(), data)?;
> -        let tdev = drm::driver::Registration::new_foreign_owned(tdev, pdev.as_ref(), 0)?;
> +        let tdev = drm::UnregisteredDevice::<Self>::new(pdev.as_ref())?;
> +        let tdev = drm::driver::Registration::new_foreign_owned(tdev, pdev.as_ref(), data, 0)?;
> 
>         let driver = TyrDriver {
>             device: tdev.into(),
> diff --git a/rust/kernel/drm/device.rs b/rust/kernel/drm/device.rs
> index f529bc7fc4032..0e81957cf8c28 100644
> --- a/rust/kernel/drm/device.rs
> +++ b/rust/kernel/drm/device.rs
> @@ -26,13 +26,15 @@
> };
> use core::{
>     alloc::Layout,
> +    cell::UnsafeCell,
>     marker::PhantomData,
> -    mem,
> -    ops::Deref,
> -    ptr::{
> +    mem::{
>         self,
> -        NonNull, //
> +        MaybeUninit, //
>     },
> +    ops::Deref,
> +    ptr::NonNull,
> +    sync::atomic::*,
> };
> 
> #[cfg(CONFIG_DRM_LEGACY)]
> @@ -141,7 +143,7 @@ impl DeviceContext for Uninit {}
> ///
> /// The device in `self.0` is guaranteed to be a newly created [`Device`] that has not yet been
> /// registered with userspace until this type is dropped.
> -pub struct UnregisteredDevice<T: drm::Driver>(ARef<Device<T, Uninit>>, NotThreadSafe);
> +pub struct UnregisteredDevice<T: drm::Driver>(pub(crate) ARef<Device<T, Uninit>>, NotThreadSafe);
> 
> impl<T: drm::Driver> Deref for UnregisteredDevice<T> {
>     type Target = Device<T, Uninit>;
> @@ -188,7 +190,7 @@ impl<T: drm::Driver> UnregisteredDevice<T> {
>     /// Create a new `UnregisteredDevice` for a `drm::Driver`.
>     ///
>     /// This can be used to create a [`Registration`](kernel::drm::Registration).
> -    pub fn new(dev: &device::Device, data: impl PinInit<T::Data, Error>) -> Result<Self> {
> +    pub fn new(dev: &device::Device) -> Result<Self> {
>         // `__drm_dev_alloc` uses `kmalloc()` to allocate memory, hence ensure a `kmalloc()`
>         // compatible `Layout`.
>         let layout = Kmalloc::aligned_layout(Layout::new::<Self>());
> @@ -207,22 +209,6 @@ pub fn new(dev: &device::Device, data: impl PinInit<T::Data, Error>) -> Result<S
>         .cast();
>         let raw_drm = NonNull::new(from_err_ptr(raw_drm)?).ok_or(ENOMEM)?;
> 
> -        // SAFETY: `raw_drm` is a valid pointer to `Self`.
> -        let raw_data = unsafe { ptr::addr_of_mut!((*raw_drm.as_ptr()).data) };
> -
> -        // SAFETY:
> -        // - `raw_data` is a valid pointer to uninitialized memory.
> -        // - `raw_data` will not move until it is dropped.
> -        unsafe { data.__pinned_init(raw_data) }.inspect_err(|_| {
> -            // SAFETY: `raw_drm` is a valid pointer to `Self`, given that `__drm_dev_alloc` was
> -            // successful.
> -            let drm_dev = unsafe { Device::into_drm_device(raw_drm) };
> -
> -            // SAFETY: `__drm_dev_alloc()` was successful, hence `drm_dev` must be valid and the
> -            // refcount must be non-zero.
> -            unsafe { bindings::drm_dev_put(drm_dev) };
> -        })?;
> -
>         // SAFETY: The reference count is one, and now we take ownership of that reference as a
>         // `drm::Device`.
>         // INVARIANT: We just created the device above, but have yet to call `drm_dev_register`.
> @@ -254,7 +240,15 @@ pub fn new(dev: &device::Device, data: impl PinInit<T::Data, Error>) -> Result<S
> #[repr(C)]
> pub struct Device<T: drm::Driver, C: DeviceContext = Registered> {
>     dev: Opaque<bindings::drm_device>,
> -    data: T::Data,
> +
> +    /// Keeps track of whether we've initialized the device data yet.
> +    pub(crate) data_is_init: AtomicBool,

Why don’t we make the data a member of the Registered context?

> +
> +    /// The Driver's private data.
> +    ///
> +    /// This must only be written to from [`drm::Registration::new`].
> +    pub(crate) data: UnsafeCell<MaybeUninit<T::Data>>,
> +
>     _ctx: PhantomData<C>,
> }
> 
> @@ -305,6 +299,21 @@ extern "C" fn release(ptr: *mut bindings::drm_device) {
>         // SAFETY: `ptr` is a valid pointer to a `struct drm_device` and embedded in `Self`.
>         let this = unsafe { Self::from_drm_device(ptr) };
> 
> +        {
> +            // SAFETY:
> +            // - Since we are in release(), we are guaranteed that no one else has access to `this`.
> +            // - We confirmed above that `this` is a valid pointer to an initialized `Self`.
> +            let this = unsafe { &mut *this };
> +            if this.data_is_init.load(Ordering::Relaxed) {
> +                // SAFETY:
> +                // - Since we are in release(), we are guaranteed that no one else has access to
> +                //   `this`.
> +                // - We checked that the data is initialized above.
> +                // - We do not use `data` any point after calling this function.
> +                unsafe { (&mut *this.data.get()).assume_init_drop() };
> +            }
> +        }
> +
>         // SAFETY:
>         // - When `release` runs it is guaranteed that there is no further access to `this`.
>         // - `this` is valid for dropping.
> @@ -323,11 +332,15 @@ pub(crate) unsafe fn assume_ctx<NewCtx: DeviceContext>(&self) -> &Device<T, NewC
>     }
> }
> 
> -impl<T: drm::Driver, C: DeviceContext> Deref for Device<T, C> {
> +impl<T: drm::Driver> Deref for Device<T, Registered> {
>     type Target = T::Data;
> 
>     fn deref(&self) -> &Self::Target {
> -        &self.data
> +        // SAFETY:
> +        // - `data` is initialized before any `Device`s with the `Registered` context are available
> +        //   to the user.
> +        // - `data` is only written to once in `Registration::new()`, so this read will never race.
> +        unsafe { (&*self.data.get()).assume_init_ref() }
>     }
> }
> 
> diff --git a/rust/kernel/drm/driver.rs b/rust/kernel/drm/driver.rs
> index da9f1bfef1f14..a16605b407159 100644
> --- a/rust/kernel/drm/driver.rs
> +++ b/rust/kernel/drm/driver.rs
> @@ -15,7 +15,8 @@
> };
> use core::{
>     mem,
> -    ptr::NonNull, //
> +    ptr::NonNull,
> +    sync::atomic::*, //
> };
> 
> /// Driver use the GEM memory manager. This should be set for all modern drivers.
> @@ -127,7 +128,18 @@ pub trait Driver {
> pub struct Registration<T: Driver>(ARef<drm::Device<T>>);
> 
> impl<T: Driver> Registration<T> {
> -    fn new(drm: drm::UnregisteredDevice<T>, flags: usize) -> Result<Self> {
> +    fn new(
> +        drm: drm::UnregisteredDevice<T>,
> +        data: impl PinInit<T::Data, Error>,
> +        flags: usize,
> +    ) -> Result<Self> {

IIUC this consumes UnregisteredDevice and returns 

Registration<T: Driver>(ARef<drm::Device<T>>);

i.e.:

Registration<T: Driver>(ARef<drm::Device<T, C = Registered>>);

Again, the Registered typestate seems like the perfect place to store T::Data.


> +        // SAFETY:
> +        // - `raw_data` is a valid pointer to uninitialized memory.
> +        // - `raw_data` will not move until it is dropped.
> +        unsafe { data.__pinned_init(drm.0.data.get().cast()) }?;
> +
> +        drm.data_is_init.store(true, Ordering::Relaxed);
> +
>         // SAFETY: `drm.as_raw()` is valid by the invariants of `drm::Device`.
>         to_result(unsafe { bindings::drm_dev_register(drm.as_raw(), flags) })?;
> 
> @@ -150,6 +162,7 @@ fn new(drm: drm::UnregisteredDevice<T>, flags: usize) -> Result<Self> {
>     pub fn new_foreign_owned<'a>(
>         drm: drm::UnregisteredDevice<T>,
>         dev: &'a device::Device<device::Bound>,
> +        data: impl PinInit<T::Data, Error>,
>         flags: usize,
>     ) -> Result<&'a drm::Device<T>>
>     where
> @@ -160,7 +173,7 @@ pub fn new_foreign_owned<'a>(
>             return Err(EINVAL);
>         }
> 
> -        let reg = Registration::<T>::new(drm, flags)?;
> +        let reg = Registration::<T>::new(drm, data, flags)?;
>         let drm = NonNull::from(reg.device());
> 
>         devres::register(dev, reg, GFP_KERNEL)?;
> -- 
> 2.52.0
> 

On Thu, 2026-01-22 at 22:52 -0300, Daniel Almeida wrote:
> > @@ -254,7 +240,15 @@ pub fn new(dev: &device::Device, data: impl
> > PinInit<T::Data, Error>) -> Result<S
> > #[repr(C)]
> > pub struct Device<T: drm::Driver, C: DeviceContext = Registered> {
> >      dev: Opaque<bindings::drm_device>,
> > -    data: T::Data,
> > +
> > +    /// Keeps track of whether we've initialized the device data
> > yet.
> > +    pub(crate) data_is_init: AtomicBool,
> 
> Why don’t we make the data a member of the Registered context?

A couple of different reasons. For one: not having it embedded as part
of Device would complicate trying to go from the device's private data
using container_of! to the actual Device struct, which isn't great for
workqueues.

The other much more important reason is that Registered isn't going to
be the only typestate that has access to the driver's private data in
the future. With KMS, nearly all of the modesetting callbacks for a
driver can be invoked before userspace registration. E.g. consider a
modesetting driver that needs to perform a modeset pre-registration so
that the hardware is in a known good state before being exposed to
userspace.

To clarify what this looks like: you'll recall that I made a diagram
showing a high-level overview of the DRM initialization process for the
documentation for DeviceContext. The second stage in that diagram,
which I'm currently calling Init, is the context that we're going to
need to eventually add a typestate for.

FWIW, this is more or less what the flow will look like with this new
context. Indenting indicates calling down to a function from within the
function above

 * UnregisteredDevice::new(d: device::Device) -> UnregisteredDevice<T>
    - // Creates Crtcs, Connectors, etc.
      KmsDriver::probe(d: &UninitializedKmsDevice)
 * // The driver sets stuff up
 * Registration::new_foreign_owned(
     dev: UnregisteredDevice,
     data: impl PinInit<T::Data>
   ):
    - // Initialize `data`, so driver data is ready at this point
    - KmsDriver::pre_registration_init(d: &Device<T, Init>)
    - // Perform actual userspace registration

If it wasn't clear too: this means tyr won't really need to do anything
when we add the new DeviceContext typestate :)

> 
> > +
> > +    /// The Driver's private data.
> > +    ///
> > +    /// This must only be written to from
> > [`drm::Registration::new`].
> > +    pub(crate) data: UnsafeCell<MaybeUninit<T::Data>>,
> > +
> >      _ctx: PhantomData<C>,