[PATCH v3 7/8] drm/tyr: track VMs in a registry

[Alvin Sun <alvin.sun@linux.dev>]

Add VmRegistry to track all VMs for debugfs enumeration. VMs are
explicitly registered and unregistered to keep the list in sync.

Signed-off-by: Alvin Sun <alvin.sun@linux.dev>
---
 drivers/gpu/drm/tyr/driver.rs |  6 +++++-
 drivers/gpu/drm/tyr/fw.rs     | 18 +++++++++++++++---
 drivers/gpu/drm/tyr/vm.rs     | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/drivers/gpu/drm/tyr/driver.rs b/drivers/gpu/drm/tyr/driver.rs
index a6694400be659..f26c3bb332e70 100644
--- a/drivers/gpu/drm/tyr/driver.rs
+++ b/drivers/gpu/drm/tyr/driver.rs
@@ -57,7 +57,8 @@
     gpu,
     gpu::GpuInfo,
     mmu::Mmu,
-    regs::gpu_control::*, //
+    regs::gpu_control::*,
+    vm::VmRegistry, //
 };
 
 pub(crate) type IoMem<'a> = kernel::io::mem::IoMem<'a, SZ_2M>;
@@ -167,11 +168,14 @@ fn probe<'bound>(
 
         let mmu = Mmu::new(pdev.as_ref(), iomem.as_arc_borrow(), &gpu_info)?;
 
+        let registry = Arc::pin_init(VmRegistry::new(), GFP_KERNEL)?;
+
         let firmware = Firmware::new(
             pdev.as_ref(),
             iomem.clone(),
             &unreg_dev,
             mmu.as_arc_borrow(),
+            registry.clone(),
             &gpu_info,
         )?;
 
diff --git a/drivers/gpu/drm/tyr/fw.rs b/drivers/gpu/drm/tyr/fw.rs
index 65ac18b92b4f2..8b09ea6c7c6d2 100644
--- a/drivers/gpu/drm/tyr/fw.rs
+++ b/drivers/gpu/drm/tyr/fw.rs
@@ -64,7 +64,6 @@
         KernelBoVaAlloc, //
     },
     gpu::GpuInfo,
-
     mmu::Mmu,
     regs::{
         gpu_control::{
@@ -76,7 +75,8 @@
         }, //
         job_control::JOB_IRQ_CLEAR,
     },
-    vm::Vm, //
+    vm::Vm,
+    vm::VmRegistry, //
 };
 
 mod interfaces;
@@ -174,6 +174,9 @@ pub(crate) struct Firmware<'drm> {
     /// MCU VM.
     vm: Arc<Vm<'drm>>,
 
+    /// VM registry, used to unregister `vm` on drop.
+    vm_registry: Arc<VmRegistry<'drm>>,
+
     /// List of firmware sections.
     sections: KVec<Section<'drm>>,
 
@@ -195,6 +198,7 @@ fn drop(self: Pin<&mut Self>) {
         let _ = self.stop();
 
         // AS slots retain a VM ref, we need to kill the circular ref manually.
+        self.vm_registry.unregister(&self.vm);
         self.vm.kill();
     }
 }
@@ -251,12 +255,18 @@ pub(crate) fn new(
         iomem: Arc<IoMem<'drm>>,
         ddev: &TyrDrmDevice,
         mmu: ArcBorrow<'_, Mmu<'drm>>,
+        vm_registry: Arc<VmRegistry<'drm>>,
         gpu_info: &GpuInfo,
     ) -> Result<Arc<Firmware<'drm>>> {
         let vm = Vm::new(dev, ddev, mmu, gpu_info)?;
-        vm.activate()?;
+        if let Err(e) = vm_registry.register(vm.clone()) {
+            dev_warn!(dev, "failed to register VM: {e:?}\n");
+        }
+        let registry = vm_registry.clone();
 
         let result = (|| {
+            vm.activate()?;
+
             let vm = &vm;
             let (fw, parsed_sections) = Self::load(dev, ddev, gpu_info)?;
             let mut sections = KVec::new();
@@ -291,6 +301,7 @@ pub(crate) fn new(
                 try_pin_init!(Firmware {
                     iomem,
                     vm: vm.clone(),
+                    vm_registry: registry,
                     sections,
                     global_iface <- new_mutex!(GlobalInterface::new()?),
                     job_irq_wait: Arc::pin_init(new_waitqueue!(), GFP_KERNEL)?,
@@ -301,6 +312,7 @@ pub(crate) fn new(
         })();
 
         if result.is_err() {
+            vm_registry.unregister(&vm);
             vm.kill();
         }
 
diff --git a/drivers/gpu/drm/tyr/vm.rs b/drivers/gpu/drm/tyr/vm.rs
index 74c3d6c8efc49..690d08b8aa901 100644
--- a/drivers/gpu/drm/tyr/vm.rs
+++ b/drivers/gpu/drm/tyr/vm.rs
@@ -311,6 +311,34 @@ pub(crate) struct GpuVmData<'drm> {
     _phantom: PhantomData<&'drm ()>,
 }
 
+/// Tracks all VMs for this device.
+#[pin_data]
+pub(crate) struct VmRegistry<'drm> {
+    #[pin]
+    vms: Mutex<KVec<Arc<Vm<'drm>>>>,
+}
+
+impl<'drm> VmRegistry<'drm> {
+    pub(crate) fn new() -> impl PinInit<Self> {
+        pin_init!(Self { vms <- new_mutex!(KVec::new()) })
+    }
+
+    pub(crate) fn register(&self, vm: Arc<Vm<'drm>>) -> Result {
+        Ok(self.vms.lock().push(vm, GFP_KERNEL)?)
+    }
+
+    pub(crate) fn unregister(&self, vm: &Vm<'drm>) {
+        self.vms.lock().retain(|v| !core::ptr::eq(&**v, vm));
+    }
+
+    pub(crate) fn for_each(&self, mut f: impl FnMut(&Vm<'drm>) -> Result) -> Result {
+        for vm in self.vms.lock().iter() {
+            f(vm)?;
+        }
+        Ok(())
+    }
+}
+
 /// GPU virtual address space.
 ///
 /// Each VM can be mapped into a hardware address space slot.
@@ -401,6 +429,11 @@ pub(crate) fn activate(&self) -> Result {
             })
     }
 
+    /// Dumps GPU VA space info into a seq_file.
+    pub(crate) fn dump_gpuva_info(&self, m: &kernel::seq_file::SeqFile) -> Result {
+        self.gpuvm_unique.lock().dump_gpuva_info(m)
+    }
+
     /// Deactivate the VM by evicting it from its address space slot.
     fn deactivate(&self) -> Result {
         self.mmu.deactivate_vm(&self.as_data).inspect_err(|e| {

-- 
2.43.0