drm/asahi: mmu: Fix deadlock on remap ops

Signed-off-by: Asahi Lina <[email protected]>

Author: hoshinolina

drivers/gpu/drm/asahi/mmu.rs

@@ -228,6 +228,8 @@ struct StepContext {
     prev_va: Option<Pin<KBox<gpuvm::GpuVa<VmInner>>>>,
     next_va: Option<Pin<KBox<gpuvm::GpuVa<VmInner>>>>,
     vm_bo: Option<ARef<gpuvm::GpuVmBo<VmInner>>>,
+    prev_vm_bo_1: Option<ARef<gpuvm::GpuVmBo<VmInner>>>,
+    prev_vm_bo_2: Option<ARef<gpuvm::GpuVmBo<VmInner>>>,
     prot: u32,
 }
 
@@ -372,6 +374,18 @@ impl gpuvm::DriverGpuVm for VmInner {
 
         self.unmap_pages(unmap_start, UAT_PGSZ, (unmap_range >> UAT_PGBIT) as usize)?;
 
+        if let Some(asid) = self.slot() {
+            mem::tlbi_range(asid as u8, unmap_start as usize, unmap_range as usize);
+            mod_dev_dbg!(
+                self.dev,
+                "MMU: flush range: asid={:#x} start={:#x} len={:#x}\n",
+                asid,
+                unmap_start,
+                unmap_range,
+            );
+            mem::sync();
+        }
+
         if op.unmap().unmap_and_unlink_va().is_none() {
             dev_err!(self.dev.as_ref(), "step_unmap: could not unlink gpuva");
         }
@@ -390,6 +404,17 @@ impl gpuvm::DriverGpuVm for VmInner {
             }
         }
 
+        if ctx.prev_vm_bo_1.is_none() {
+            assert!(ctx.prev_vm_bo_1.replace(vm_bo).is_none());
+        } else if ctx.prev_vm_bo_2.is_none() {
+            assert!(ctx.prev_vm_bo_2.replace(vm_bo).is_none());
+        } else {
+            dev_crit!(
+                self.dev.as_ref(),
+                "step_remap: too many remap ops, deadlock expected"
+            );
+        }
+
         Ok(())
     }
 }
@@ -1194,8 +1219,8 @@ impl Vm {
             new_va: Some(gpuvm::GpuVa::<VmInner>::new(init::default())?),
             prev_va: Some(gpuvm::GpuVa::<VmInner>::new(init::default())?),
             next_va: Some(gpuvm::GpuVa::<VmInner>::new(init::default())?),
-            vm_bo: None,
             prot,
+            ..Default::default()
         };
 
         let sgt = gem.sg_table()?;