drm/asahi: G15 pin module before retaining firmware ownership

drivers/gpu/drm/asahi/g15/boot.rs

@@ -1,3 +1,3 @@
 // SPDX-License-Identifier: GPL-2.0-only OR MIT
 
 //! Admission and firmware start to READY.
@@ -61,7 +61,8 @@
         fn asahi_g15_machine(compatible: *const core::ffi::c_char) -> bool;
         fn asahi_g15_hold_power(dev: *mut core::ffi::c_void);
         fn asahi_g15_release_power(dev: *mut core::ffi::c_void);
-        fn asahi_g15_module_get();
+        fn asahi_g15_module_try_get() -> bool;
+        fn asahi_g15_module_put();
         pub(crate) fn asahi_g15_reprobe();
         pub(crate) fn asahi_g15_report_loader_gpu();
         fn asahi_g15_adt_phram(start: *mut u64, size: *mut u64) -> core::ffi::c_int;
@@ -121,12 +122,17 @@
         RETAINED.load(Ordering::Acquire)
     }
 
-    /// Pins this module until reboot, once (`__module_get(THIS_MODULE)`).
-    pub(crate) fn retain_module() {
-        if !RETAINED.swap(true, Ordering::AcqRel) {
-            // SAFETY: The module is live while its code runs; the reference is never put.
-            unsafe { asahi_g15_module_get() };
-        }
+    /// Pins this module until reboot, once (`try_module_get(THIS_MODULE)`, see `retain.rs`).
+    /// ENODEV, with nothing pinned or marked, once the module's removal has begun.
+    pub(crate) fn retain_module() -> Result {
+        crate::g15::retain::retain(
+            &RETAINED,
+            // SAFETY: Takes a reference to this module unless its removal has begun.
+            || unsafe { asahi_g15_module_try_get() },
+            // SAFETY: Drops only the reference this call took; the winner's is never put.
+            || unsafe { asahi_g15_module_put() },
+        )
+        .map_err(|_| ENODEV)
     }
 
     /// The phram "adt" export: start and size.
@@ -672,8 +678,16 @@
         return Err(e);
     }
 
+    // Pin the module before anything firmware-visible. Once its removal has begun this fails
+    // and the firmware is left untouched, cleaned up as for a failed self-check: the vote is
+    // released and the session stays unretained.
+    if let Err(e) = glue::retain_module() {
+        dev_err!(d, "G15: the module is being removed; the firmware is not started\n");
+        // Mapped with cached firmware classes: never unmapped.
+        core::mem::forget(objects);
+        return Err(e);
+    }
     // From here on the firmware is touched: nothing it can reach is freed, the vote included.
-    glue::retain_module();
     vote.keep();
     let rtkit = match Rtkit::new(&d, shared) {
         Ok(rtkit) => rtkit,

drivers/gpu/drm/asahi/g15/glue.c

@@ -1,3 +1,3 @@
 // SPDX-License-Identifier: GPL-2.0-only OR MIT
 /*
  * G15 (T6030/T6031, M3 Pro/Max) helpers that have no Rust bindings: static inlines,
@@ -26,7 +26,8 @@
 bool asahi_g15_machine(const char *compatible);
 void asahi_g15_hold_power(struct device *dev);
 void asahi_g15_release_power(struct device *dev);
-void asahi_g15_module_get(void);
+bool asahi_g15_module_try_get(void);
+void asahi_g15_module_put(void);
 void asahi_g15_reprobe(void);
 void asahi_g15_report_loader_gpu(void);
 int asahi_g15_adt_phram(u64 *start, u64 *size);
@@ -209,10 +210,20 @@
 
 }
 
-/* Pin this module until reboot. */
-void asahi_g15_module_get(void)
+/*
+ * Pin this module until reboot, unless its removal has begun: a deferred or
+ * asynchronous probe can run after an ordinary rmmod has passed its reference
+ * check, which __module_get() would not stop.
+ */
+bool asahi_g15_module_try_get(void)
+{
+	return try_module_get(THIS_MODULE);
+}
+
+/* Drop a reference asahi_g15_module_try_get() took. */
+void asahi_g15_module_put(void)
 {
-	__module_get(THIS_MODULE);
+	module_put(THIS_MODULE);
 }
 
 /*

drivers/gpu/drm/asahi/g15/mod.rs

@@ -1,3 +1,3 @@
 // SPDX-License-Identifier: GPL-2.0-only OR MIT
 
 //! G15 GPU support (T6030, M3 Pro: J514S, J516S; firmware RTKit 2419, the macOS 14.7-14.8 stub).
@@ -45,6 +45,7 @@
 mod adt;
 mod board;
 mod boot;
+mod retain;
 
 // One admission rule for the generation mappings and shared adopted UAT.
 pub(crate) use boot::glue::memory_mode;
@@ -452,17 +453,18 @@
             },
             GFP_KERNEL,
         )?;
-        gpu.retain();
+        gpu.retain()?;
         Ok(gpu)
     }
 
     /// Keeps the firmware session and this module alive until reboot.
     ///
-    /// Takes the module reference once. Nothing firmware-reachable is freed on any path: the
+    /// Takes the module reference once; `boot::start` already took it before touching the
+    /// firmware, so this only confirms it. Nothing firmware-reachable is freed on any path: the
     /// RTKit callbacks hold [`Shared`], which holds the [`Firmware`] (or the [`boot::Orphan`] of a
     /// failed start) until reboot.
-    pub(crate) fn retain(&self) {
-        boot::glue::retain_module();
+    pub(crate) fn retain(&self) -> Result {
+        boot::glue::retain_module()
     }
 }
 

drivers/gpu/drm/asahi/g15/retain.rs

@@ -0,0 +1,136 @@
+// SPDX-License-Identifier: GPL-2.0-only OR MIT
+
+//! The once-per-boot module pin taken before the firmware is first touched.
+//!
+//! A started G15 firmware session cannot be stopped, so this module must stay loaded until
+//! reboot from the first firmware-visible write on. `__module_get` alone does not ensure that:
+//! a deferred or asynchronous probe can run while an ordinary rmmod is already past its
+//! reference check, and the module is then unloaded under a running firmware.
+//! `try_module_get` fails once removal has begun, so the pin is taken with it, before the
+//! session is marked retained and before anything firmware-visible happens.
+
+use core::sync::atomic::{AtomicBool, Ordering};
+
+/// The module is being removed: nothing was pinned or marked, and nothing may be started.
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) struct Removing;
+
+/// Pins the module once for the session `retained` tracks. `get` is
+/// `try_module_get(THIS_MODULE)`, `put` is `module_put(THIS_MODULE)`.
+///
+/// Already retained: nothing more is taken. Otherwise the reference comes first; `retained` is
+/// set only once it is held. A caller that loses the race to set it drops only its own
+/// reference, so exactly one reference stays and it is never put.
+pub(crate) fn retain(
+    retained: &AtomicBool,
+    get: impl FnOnce() -> bool,
+    put: impl FnOnce(),
+) -> Result<(), Removing> {
+    if retained.load(Ordering::Acquire) {
+        return Ok(());
+    }
+    if !get() {
+        return Err(Removing);
+    }
+    if retained.swap(true, Ordering::AcqRel) {
+        put();
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use core::sync::atomic::AtomicI32;
+    use std::sync::Arc;
+
+    #[test]
+    fn removal_fails_closed_with_nothing_pinned_or_marked() {
+        let retained = AtomicBool::new(false);
+        let refs = AtomicI32::new(1); // base reference only: rmmod has released it
+        let put = || panic!("nothing to put");
+        assert_eq!(retain(&retained, || false, put), Err(Removing));
+        assert!(!retained.load(Ordering::Acquire), "a refused pin must not mark the session");
+        assert_eq!(refs.load(Ordering::Acquire), 1);
+        // A later live start can still pin: the refusal left no state behind.
+        let get = || {
+            refs.fetch_add(1, Ordering::AcqRel);
+            true
+        };
+        assert_eq!(retain(&retained, get, || panic!("winner keeps its reference")), Ok(()));
+        assert!(retained.load(Ordering::Acquire));
+        assert_eq!(refs.load(Ordering::Acquire), 2);
+    }
+
+    #[test]
+    fn first_pin_is_kept_and_later_calls_take_nothing() {
+        let retained = AtomicBool::new(false);
+        let refs = AtomicI32::new(1);
+        let get = || {
+            refs.fetch_add(1, Ordering::AcqRel);
+            true
+        };
+        assert_eq!(retain(&retained, get, || panic!("winner keeps its reference")), Ok(()));
+        // The post-start call in Gpu::retain and any later one: no new reference, no put.
+        for _ in 0..3 {
+            assert_eq!(
+                retain(&retained, || panic!("already retained"), || panic!("nothing to put")),
+                Ok(())
+            );
+        }
+        assert_eq!(refs.load(Ordering::Acquire), 2);
+    }
+
+    #[test]
+    fn a_lost_race_drops_only_its_own_reference() {
+        let retained = AtomicBool::new(false);
+        let refs = AtomicI32::new(2); // base + the reference the winner already holds
+        let puts = AtomicI32::new(0);
+        // The winner marks the session between this caller's check and its swap.
+        let get = || {
+            refs.fetch_add(1, Ordering::AcqRel);
+            retained.store(true, Ordering::Release);
+            true
+        };
+        let put = || {
+            refs.fetch_sub(1, Ordering::AcqRel);
+            puts.fetch_add(1, Ordering::AcqRel);
+        };
+        assert_eq!(retain(&retained, get, put), Ok(()));
+        assert_eq!(puts.load(Ordering::Acquire), 1);
+        assert_eq!(refs.load(Ordering::Acquire), 2);
+    }
+
+    #[test]
+    fn concurrent_pins_leave_exactly_one_reference() {
+        for _ in 0..500 {
+            let retained = Arc::new(AtomicBool::new(false));
+            let refs = Arc::new(AtomicI32::new(1));
+            let start = Arc::new(std::sync::Barrier::new(8));
+            let threads: Vec<_> = (0..8)
+                .map(|_| {
+                    let (retained, refs, start) = (retained.clone(), refs.clone(), start.clone());
+                    std::thread::spawn(move || {
+                        let (a, b) = (refs.clone(), refs.clone());
+                        start.wait();
+                        retain(
+                            &retained,
+                            move || {
+                                a.fetch_add(1, Ordering::AcqRel);
+                                true
+                            },
+                            move || {
+                                b.fetch_sub(1, Ordering::AcqRel);
+                            },
+                        )
+                    })
+                })
+                .collect();
+            for t in threads {
+                assert_eq!(t.join().unwrap(), Ok(()));
+            }
+            assert!(retained.load(Ordering::Acquire));
+            assert_eq!(refs.load(Ordering::Acquire), 2, "exactly one kept reference");
+        }
+    }
+}

tools/testing/selftests/drivers/gpu/asahi-g15/run.py

@@ -306,6 +306,7 @@ def main():
         harness += f'#[path = "{G15 / "recovery/logic.rs"}"] mod recovery;\n'
         harness += f'#[path = "{G15 / "power/idle.rs"}"] mod idle;\n'
         harness += f'#[path = "{G15 / "power/pstate.rs"}"] mod pstate;\n'
+        harness += f'#[path = "{G15 / "retain.rs"}"] mod retain;\n'
         harness += ("mod firmware_identity {\n" +
                     f'#[path = "{G15 / "identity/matching.rs"}"] pub mod matching;\n' +
                     f'#[path = "{G15 / "identity/lookup.rs"}"] pub mod lookup;\n' + "}\n")