rusticl/memory: Fix memory unmaps after rework

An application could map and unmap a host ptr allocation multiple times,
but because how the refcounting works, we might never ended up syncing the
written data to the mapped region.

This moves the refcounting out of the event processing.

Fixes: 7b22bc617b ("rusticl/memory: complete rework on how mapping is implemented")
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30528>
(cherry picked from commit 1fa288b224)

.pick_status.json

@@ -3644,7 +3644,7 @@
         "description": "rusticl/memory: Fix memory unmaps after rework",
         "nominated": true,
         "nomination_type": 1,
-        "resolution": 0,
+        "resolution": 1,
         "main_sha": null,
         "because_sha": "7b22bc617bf2db4120a438c1ad5e45992f638d82",
         "notes": null

src/gallium/frontends/rusticl/api/memory.rs

@@ -2192,13 +2192,20 @@ fn enqueue_unmap_mem_object(
 
     // SAFETY: it's required that applications do not cause data races
     let mapped_ptr = unsafe { MutMemoryPtr::from_ptr(mapped_ptr) };
+    let needs_sync = m.unmap(mapped_ptr)?;
     create_and_queue(
         q,
         CL_COMMAND_UNMAP_MEM_OBJECT,
         evs,
         event,
         false,
-        Box::new(move |q, ctx| m.unmap(q, ctx, mapped_ptr)),
+        Box::new(move |q, ctx| {
+            if needs_sync {
+                m.sync_unmap(q, ctx, mapped_ptr)
+            } else {
+                Ok(())
+            }
+        }),
     )
 }
 

src/gallium/frontends/rusticl/core/memory.rs

@@ -39,6 +39,8 @@ struct Mapping<T> {
     layout: Layout,
     writes: bool,
     ptr: Option<MutMemoryPtr>,
+    /// reference count from the API perspective. Once it reaches 0, we need to write back the
+    /// mappings content to the GPU resource.
     count: u32,
     inner: T,
 }
@@ -152,10 +154,17 @@ impl Mem {
         }
     }
 
-    pub fn unmap(&self, q: &Queue, ctx: &PipeContext, ptr: MutMemoryPtr) -> CLResult<()> {
+    pub fn sync_unmap(&self, q: &Queue, ctx: &PipeContext, ptr: MutMemoryPtr) -> CLResult<()> {
         match self {
-            Self::Buffer(b) => b.unmap(q, ctx, ptr),
-            Self::Image(i) => i.unmap(q, ctx, ptr),
+            Self::Buffer(b) => b.sync_unmap(q, ctx, ptr),
+            Self::Image(i) => i.sync_unmap(q, ctx, ptr),
+        }
+    }
+
+    pub fn unmap(&self, ptr: MutMemoryPtr) -> CLResult<bool> {
+        match self {
+            Self::Buffer(b) => b.unmap(ptr),
+            Self::Image(i) => i.unmap(ptr),
         }
     }
 }
@@ -897,7 +906,9 @@ impl Buffer {
     }
 
     fn is_mapped_ptr(&self, ptr: *mut c_void) -> bool {
-        self.maps.lock().unwrap().contains_key(ptr as usize)
+        let mut maps = self.maps.lock().unwrap();
+        let entry = maps.entry(ptr as usize);
+        matches!(entry, Entry::Occupied(entry) if entry.get().count > 0)
     }
 
     pub fn map(&self, size: usize, offset: usize, writes: bool) -> CLResult<MutMemoryPtr> {
@@ -978,6 +989,31 @@ impl Buffer {
         self.read(q, ctx, mapping.offset, ptr, mapping.size())
     }
 
+    pub fn sync_unmap(&self, q: &Queue, ctx: &PipeContext, ptr: MutMemoryPtr) -> CLResult<()> {
+        // no need to update
+        if self.is_pure_user_memory(q.device)? {
+            return Ok(());
+        }
+
+        match self.maps.lock().unwrap().entry(ptr.as_ptr() as usize) {
+            Entry::Vacant(_) => Err(CL_INVALID_VALUE),
+            Entry::Occupied(entry) => {
+                let mapping = entry.get();
+
+                if mapping.writes {
+                    self.write(q, ctx, mapping.offset, ptr.into(), mapping.size())?;
+                }
+
+                // only remove if the mapping wasn't reused in the meantime
+                if mapping.count == 0 {
+                    entry.remove();
+                }
+
+                Ok(())
+            }
+        }
+    }
+
     fn tx<'a>(
         &self,
         q: &Queue,
@@ -1002,22 +1038,16 @@ impl Buffer {
     }
 
     // TODO: only sync on unmap when the memory is not mapped for writing
-    pub fn unmap(&self, q: &Queue, ctx: &PipeContext, ptr: MutMemoryPtr) -> CLResult<()> {
-        let mapping = match self.maps.lock().unwrap().entry(ptr.as_ptr() as usize) {
-            Entry::Vacant(_) => return Err(CL_INVALID_VALUE),
+    pub fn unmap(&self, ptr: MutMemoryPtr) -> CLResult<bool> {
+        match self.maps.lock().unwrap().entry(ptr.as_ptr() as usize) {
+            Entry::Vacant(_) => Err(CL_INVALID_VALUE),
             Entry::Occupied(mut entry) => {
-                entry.get_mut().count -= 1;
-                (entry.get().count == 0).then(|| entry.remove())
+                let entry = entry.get_mut();
+                debug_assert!(entry.count > 0);
+                entry.count -= 1;
+                Ok(entry.count == 0)
             }
-        };
-
-        if let Some(mapping) = mapping {
-            if mapping.writes && !self.is_pure_user_memory(q.device)? {
-                self.write(q, ctx, mapping.offset, ptr.into(), mapping.size())?;
-            }
-        };
-
-        Ok(())
+        }
     }
 
     pub fn write(
@@ -1277,7 +1307,9 @@ impl Image {
     }
 
     fn is_mapped_ptr(&self, ptr: *mut c_void) -> bool {
-        self.maps.lock().unwrap().contains_key(ptr as usize)
+        let mut maps = self.maps.lock().unwrap();
+        let entry = maps.entry(ptr as usize);
+        matches!(entry, Entry::Occupied(entry) if entry.get().count > 0)
     }
 
     pub fn is_parent_buffer(&self) -> bool {
@@ -1406,6 +1438,41 @@ impl Image {
         )
     }
 
+    pub fn sync_unmap(&self, q: &Queue, ctx: &PipeContext, ptr: MutMemoryPtr) -> CLResult<()> {
+        // no need to update
+        if self.is_pure_user_memory(q.device)? {
+            return Ok(());
+        }
+
+        match self.maps.lock().unwrap().entry(ptr.as_ptr() as usize) {
+            Entry::Vacant(_) => Err(CL_INVALID_VALUE),
+            Entry::Occupied(entry) => {
+                let mapping = entry.get();
+                let row_pitch = self.image_desc.row_pitch()? as usize;
+                let slice_pitch = self.image_desc.slice_pitch();
+
+                if mapping.writes {
+                    self.write(
+                        ptr.into(),
+                        q,
+                        ctx,
+                        &mapping.region,
+                        row_pitch,
+                        slice_pitch,
+                        &mapping.origin,
+                    )?;
+                }
+
+                // only remove if the mapping wasn't reused in the meantime
+                if mapping.count == 0 {
+                    entry.remove();
+                }
+
+                Ok(())
+            }
+        }
+    }
+
     fn tx_image<'a>(
         &self,
         q: &Queue,
@@ -1421,33 +1488,16 @@ impl Image {
     }
 
     // TODO: only sync on unmap when the memory is not mapped for writing
-    pub fn unmap(&self, q: &Queue, ctx: &PipeContext, ptr: MutMemoryPtr) -> CLResult<()> {
-        let mapping = match self.maps.lock().unwrap().entry(ptr.as_ptr() as usize) {
-            Entry::Vacant(_) => return Err(CL_INVALID_VALUE),
+    pub fn unmap(&self, ptr: MutMemoryPtr) -> CLResult<bool> {
+        match self.maps.lock().unwrap().entry(ptr.as_ptr() as usize) {
+            Entry::Vacant(_) => Err(CL_INVALID_VALUE),
             Entry::Occupied(mut entry) => {
-                entry.get_mut().count -= 1;
-                (entry.get().count == 0).then(|| entry.remove())
-            }
-        };
-
-        let row_pitch = self.image_desc.row_pitch()? as usize;
-        let slice_pitch = self.image_desc.slice_pitch();
-
-        if let Some(mapping) = mapping {
-            if mapping.writes && !self.is_pure_user_memory(q.device)? {
-                self.write(
-                    ptr.into(),
-                    q,
-                    ctx,
-                    &mapping.region,
-                    row_pitch,
-                    slice_pitch,
-                    &mapping.origin,
-                )?;
+                let entry = entry.get_mut();
+                debug_assert!(entry.count > 0);
+                entry.count -= 1;
+                Ok(entry.count == 0)
             }
         }
-
-        Ok(())
     }
 
     pub fn write(