1 files changed, 321 insertions, 140 deletions

diff --git a/drivers/accel/amdxdna/amdxdna_gem.c b/drivers/accel/amdxdna/amdxdna_gem.c
index 4b0d58d0329b..f14484bc31ec 100644
--- a/drivers/accel/amdxdna/amdxdna_gem.c
+++ b/drivers/accel/amdxdna/amdxdna_gem.c
@@ -14,6 +14,7 @@
 #include <linux/dma-direct.h>
 #include <linux/iosys-map.h>
 #include <linux/pagemap.h>
+#include <linux/swap.h>
 #include <linux/vmalloc.h>
 
 #include "amdxdna_cbuf.h"
@@ -37,6 +38,7 @@ amdxdna_init_dev_bo(struct amdxdna_gem_obj *dev_bo)
 	struct amdxdna_gem_obj *heap;
 	u64 heap_addr, exp_heap_uva;
 	u32 heap_id;
+	int ret;
 
 	if (xa_empty(&client->dev_heap_xa)) {
 		XDNA_DBG(xdna, "Empty heap xa");
@@ -58,24 +60,32 @@ amdxdna_init_dev_bo(struct amdxdna_gem_obj *dev_bo)
 	heap = xa_load(&client->dev_heap_xa, heap_id);
 	exp_heap_uva = amdxdna_gem_uva(heap);
 	heap_addr = amdxdna_gem_dev_addr(heap);
-	dev_bo->heap_start_id = heap_id;
 	dev_bo->mem.uva = dev_bo->mm_node.start - heap_addr + exp_heap_uva;
 
 	for (; heap_id < client->dev_heap_nid; heap_id++) {
 		heap = xa_load(&client->dev_heap_xa, heap_id);
 		if (!heap) {
 			XDNA_ERR(xdna, "Failed to load heap %d", heap_id);
-			return -EINVAL;
+			ret = -EINVAL;
+			goto cleanup_heap_xa;
 		}
 		heap_addr = amdxdna_gem_uva(heap);
 		if (heap_addr == AMDXDNA_INVALID_ADDR) {
 			XDNA_ERR(xdna, "Heap %d is not mapped", heap_id);
-			return -EAGAIN;
+			ret = -EAGAIN;
+			goto cleanup_heap_xa;
 		}
 
 		if (heap_addr != exp_heap_uva) {
 			XDNA_ERR(xdna, "Heap %d uva is not contiguous", heap_id);
-			return -EINVAL;
+			ret = -EINVAL;
+			goto cleanup_heap_xa;
+		}
+
+		ret = xa_insert(&dev_bo->heap_xa, heap_id, heap, GFP_KERNEL);
+		if (ret) {
+			ret = -ENOMEM;
+			goto cleanup_heap_xa;
 		}
 
 		if (heap->dev_addr + heap->mem.size >=
@@ -87,12 +97,15 @@ amdxdna_init_dev_bo(struct amdxdna_gem_obj *dev_bo)
 
 	if (heap_id == client->dev_heap_nid) {
 		XDNA_DBG(xdna, "Can not find heap end");
-		return -EAGAIN;
+		ret = -EAGAIN;
+		goto cleanup_heap_xa;
 	}
 
-	dev_bo->heap_end_id = heap_id;
-
 	return 0;
+
+cleanup_heap_xa:
+	xa_destroy(&dev_bo->heap_xa);
+	return ret;
 }
 
 static int
@@ -132,8 +145,7 @@ amdxdna_gem_heap_alloc(struct amdxdna_gem_obj *abo)
 	}
 
 	client->heap_usage += mem->size;
-	xa_for_each_range(&client->dev_heap_xa, heap_id, heap,
-			  abo->heap_start_id, abo->heap_end_id)
+	xa_for_each(&abo->heap_xa, heap_id, heap)
 		drm_gem_object_get(to_gobj(heap));
 
 unlock_out:
@@ -142,48 +154,16 @@ unlock_out:
 	return ret;
 }
 
-static void
-amdxdna_gem_heap_free(struct amdxdna_gem_obj *abo)
+void amdxdna_gem_heap_free(struct amdxdna_client *client, struct amdxdna_gem_obj *abo)
 {
-	struct amdxdna_client *client = abo->client;
-	struct amdxdna_gem_obj *heap;
-	unsigned long heap_id;
-
 	mutex_lock(&client->mm_lock);
 
 	drm_mm_remove_node(&abo->mm_node);
 	client->heap_usage -= abo->mem.size;
 
-	xa_for_each_range(&client->dev_heap_xa, heap_id, heap,
-			  abo->heap_start_id, abo->heap_end_id)
-		drm_gem_object_put(to_gobj(heap));
-
 	mutex_unlock(&client->mm_lock);
 }
 
-static struct amdxdna_gem_obj *
-amdxdna_gem_create_obj(struct drm_device *dev, size_t size)
-{
-	struct amdxdna_gem_obj *abo;
-
-	abo = kzalloc_obj(*abo);
-	if (!abo)
-		return ERR_PTR(-ENOMEM);
-
-	abo->pinned = false;
-	abo->assigned_hwctx = AMDXDNA_INVALID_CTX_HANDLE;
-	mutex_init(&abo->lock);
-
-	abo->mem.dma_addr = AMDXDNA_INVALID_ADDR;
-	abo->mem.uva = AMDXDNA_INVALID_ADDR;
-	abo->mem.size = size;
-	abo->open_ref = 0;
-	abo->internal = false;
-	INIT_LIST_HEAD(&abo->mem.umap_list);
-
-	return abo;
-}
-
 static void
 amdxdna_gem_destroy_obj(struct amdxdna_gem_obj *abo)
 {
@@ -253,13 +233,16 @@ static bool amdxdna_hmm_invalidate(struct mmu_interval_notifier *mni,
 	struct amdxdna_gem_obj *abo = mapp->abo;
 	struct amdxdna_dev *xdna;
 
+	if (!mmu_notifier_range_blockable(range))
+		return false;
+
+	if (mapp->unmapped)
+		return true;
+
 	xdna = to_xdna_dev(to_gobj(abo)->dev);
 	XDNA_DBG(xdna, "Invalidating range 0x%lx, 0x%lx, type %d",
 		 mapp->range.start, mapp->range.end, abo->type);
 
-	if (!mmu_notifier_range_blockable(range))
-		return false;
-
 	down_write(&xdna->notifier_lock);
 	abo->mem.map_invalid = true;
 	mapp->invalid = true;
@@ -271,10 +254,8 @@ static bool amdxdna_hmm_invalidate(struct mmu_interval_notifier *mni,
 
 	if (range->event == MMU_NOTIFY_UNMAP) {
 		down_write(&xdna->notifier_lock);
-		if (!mapp->unmapped) {
-			queue_work(xdna->notifier_wq, &mapp->hmm_unreg_work);
-			mapp->unmapped = true;
-		}
+		mapp->unmapped = true;
+		queue_work(xdna->notifier_wq, &abo->hmm_unreg_work);
 		up_write(&xdna->notifier_lock);
 	}
 
@@ -301,33 +282,42 @@ static void amdxdna_hmm_unregister(struct amdxdna_gem_obj *abo,
 
 	down_write(&xdna->notifier_lock);
 	list_for_each_entry(mapp, &abo->mem.umap_list, node) {
-		if (!vma || compare_range(mapp, vma->vm_mm, vma->vm_start, vma->vm_end)) {
-			if (!mapp->unmapped) {
-				queue_work(xdna->notifier_wq, &mapp->hmm_unreg_work);
-				mapp->unmapped = true;
-			}
-			if (vma)
-				break;
-		}
+		if (!compare_range(mapp, vma->vm_mm, vma->vm_start, vma->vm_end))
+			continue;
+
+		mapp->unmapped = true;
+		queue_work(xdna->notifier_wq, &abo->hmm_unreg_work);
 	}
 	up_write(&xdna->notifier_lock);
 }
 
-static void amdxdna_umap_release(struct kref *ref)
+static void amdxdna_hmm_unreg_umaps(struct amdxdna_gem_obj *abo, bool force)
 {
-	struct amdxdna_umap *mapp = container_of(ref, struct amdxdna_umap, refcnt);
-	struct amdxdna_gem_obj *abo = mapp->abo;
-	struct amdxdna_dev *xdna;
-
-	mmu_interval_notifier_remove(&mapp->notifier);
+	struct amdxdna_dev *xdna = to_xdna_dev(to_gobj(abo)->dev);
+	struct amdxdna_umap *mapp, *tmp;
+	LIST_HEAD(dead);
 
-	xdna = to_xdna_dev(to_gobj(mapp->abo)->dev);
 	down_write(&xdna->notifier_lock);
-	list_del(&mapp->node);
+	list_for_each_entry_safe(mapp, tmp, &abo->mem.umap_list, node) {
+		if (!force && !mapp->unmapped)
+			continue;
+
+		mapp->unmapped = true;
+		list_move(&mapp->node, &dead);
+	}
 	if (list_empty(&abo->mem.umap_list))
 		abo->mem.uva = AMDXDNA_INVALID_ADDR;
 	up_write(&xdna->notifier_lock);
 
+	list_for_each_entry_safe(mapp, tmp, &dead, node)
+		amdxdna_umap_put(mapp);
+}
+
+static void amdxdna_umap_release(struct kref *ref)
+{
+	struct amdxdna_umap *mapp = container_of(ref, struct amdxdna_umap, refcnt);
+
+	mmu_interval_notifier_remove(&mapp->notifier);
 	kvfree(mapp->range.hmm_pfns);
 	kfree(mapp);
 }
@@ -339,10 +329,10 @@ void amdxdna_umap_put(struct amdxdna_umap *mapp)
 
 static void amdxdna_hmm_unreg_work(struct work_struct *work)
 {
-	struct amdxdna_umap *mapp = container_of(work, struct amdxdna_umap,
-						 hmm_unreg_work);
+	struct amdxdna_gem_obj *abo = container_of(work, struct amdxdna_gem_obj,
+						   hmm_unreg_work);
 
-	amdxdna_umap_put(mapp);
+	amdxdna_hmm_unreg_umaps(abo, false);
 }
 
 static int amdxdna_hmm_register(struct amdxdna_gem_obj *abo,
@@ -387,6 +377,21 @@ static int amdxdna_hmm_register(struct amdxdna_gem_obj *abo,
 		goto free_map;
 	}
 
+	mapp->range.notifier = &mapp->notifier;
+	mapp->range.start = vma->vm_start;
+	mapp->range.end = vma->vm_end;
+	/*
+	 * Access permissions are fixed at mmap() time. Changing them later
+	 * with mprotect() is not supported: the range keeps requesting the
+	 * original permissions, so the application may see a fault failure
+	 * or an IOMMU fault.
+	 */
+	mapp->range.default_flags = HMM_PFN_REQ_FAULT;
+	if (vma->vm_flags & VM_WRITE)
+		mapp->range.default_flags |= HMM_PFN_REQ_WRITE;
+	mapp->abo = abo;
+	kref_init(&mapp->refcnt);
+
 	ret = mmu_interval_notifier_insert_locked(&mapp->notifier,
 						  current->mm,
 						  addr,
@@ -397,15 +402,6 @@ static int amdxdna_hmm_register(struct amdxdna_gem_obj *abo,
 		goto free_pfns;
 	}
 
-	mapp->range.notifier = &mapp->notifier;
-	mapp->range.start = vma->vm_start;
-	mapp->range.end = vma->vm_end;
-	mapp->range.default_flags = HMM_PFN_REQ_FAULT;
-	mapp->abo = abo;
-	kref_init(&mapp->refcnt);
-
-	INIT_WORK(&mapp->hmm_unreg_work, amdxdna_hmm_unreg_work);
-
 	down_write(&xdna->notifier_lock);
 	if (list_empty(&abo->mem.umap_list))
 		abo->mem.uva = addr;
@@ -421,17 +417,47 @@ free_map:
 	return ret;
 }
 
+static struct amdxdna_gem_obj *
+amdxdna_gem_create_obj(struct drm_device *dev, size_t size)
+{
+	struct amdxdna_gem_obj *abo;
+
+	abo = kzalloc_obj(*abo);
+	if (!abo)
+		return ERR_PTR(-ENOMEM);
+
+	abo->pinned = false;
+	abo->assigned_hwctx = AMDXDNA_INVALID_CTX_HANDLE;
+	mutex_init(&abo->lock);
+
+	abo->mem.dma_addr = AMDXDNA_INVALID_ADDR;
+	abo->mem.uva = AMDXDNA_INVALID_ADDR;
+	abo->mem.size = size;
+	abo->open_ref = 0;
+	abo->internal = false;
+	INIT_LIST_HEAD(&abo->mem.umap_list);
+	xa_init_flags(&abo->heap_xa, XA_FLAGS_ALLOC);
+	INIT_WORK(&abo->hmm_unreg_work, amdxdna_hmm_unreg_work);
+
+	return abo;
+}
+
 static void amdxdna_gem_dev_obj_free(struct drm_gem_object *gobj)
 {
 	struct amdxdna_dev *xdna = to_xdna_dev(gobj->dev);
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
+	struct amdxdna_gem_obj *heap;
+	unsigned long heap_id;
 
 	XDNA_DBG(xdna, "BO type %d xdna_addr 0x%llx", abo->type, amdxdna_gem_dev_addr(abo));
 	if (abo->pinned)
 		amdxdna_gem_unpin(abo);
 
 	amdxdna_gem_vunmap(abo);
-	amdxdna_gem_heap_free(abo);
+	xa_for_each(&abo->heap_xa, heap_id, heap)
+		drm_gem_object_put(to_gobj(heap));
+	xa_destroy(&abo->heap_xa);
+
 	drm_gem_object_release(gobj);
 	amdxdna_gem_destroy_obj(abo);
 }
@@ -458,16 +484,16 @@ static int amdxdna_insert_pages(struct amdxdna_gem_obj *abo,
 {
 	struct amdxdna_dev *xdna = to_xdna_dev(to_gobj(abo)->dev);
 	unsigned long num_pages = vma_pages(vma);
-	unsigned long offset = 0;
 	int ret;
 
-	if (!is_import_bo(abo)) {
-		ret = drm_gem_shmem_mmap(&abo->base, vma);
-		if (ret) {
-			XDNA_ERR(xdna, "Failed shmem mmap %d", ret);
-			return ret;
-		}
-	} else {
+	/*
+	 * Until today there is not any use case to mmap with non-zero
+	 * offset. Put an explicit check here.
+	 */
+	if (vma->vm_pgoff - drm_vma_node_start(&to_gobj(abo)->vma_node))
+		return -EINVAL;
+
+	if (is_import_bo(abo)) {
 		vma->vm_private_data = NULL;
 		vma->vm_ops = NULL;
 		ret = dma_buf_mmap(abo->dma_buf, vma, 0);
@@ -476,23 +502,28 @@ static int amdxdna_insert_pages(struct amdxdna_gem_obj *abo,
 			return ret;
 		}
 
+		amdxdna_mark_mapp_invalid(abo, vma);
+
 		/* Drop the reference drm_gem_mmap_obj() acquired.*/
 		drm_gem_object_put(to_gobj(abo));
+		return 0;
 	}
 
-	do {
-		vm_fault_t fault_ret;
-
-		fault_ret = handle_mm_fault(vma, vma->vm_start + offset,
-					    FAULT_FLAG_WRITE, NULL);
-		if (fault_ret & VM_FAULT_ERROR) {
-			XDNA_ERR(xdna, "Fault in page failed");
-			amdxdna_mark_mapp_invalid(abo, vma);
-			break;
-		}
+	ret = drm_gem_shmem_mmap(&abo->base, vma);
+	if (ret) {
+		XDNA_ERR(xdna, "Failed shmem mmap %d", ret);
+		return ret;
+	}
 
-		offset += PAGE_SIZE;
-	} while (--num_pages);
+	vm_flags_mod(vma, VM_MIXEDMAP, VM_PFNMAP);
+	ret = vm_insert_pages(vma, vma->vm_start, abo->base.pages, &num_pages);
+	if (ret) {
+		XDNA_ERR(xdna, "Failed to insert pages %d", ret);
+		dma_resv_lock(to_gobj(abo)->resv, NULL);
+		drm_gem_shmem_put_pages_locked(&abo->base);
+		dma_resv_unlock(to_gobj(abo)->resv);
+		return ret;
+	}
 
 	return 0;
 }
@@ -504,6 +535,10 @@ static int amdxdna_gem_obj_mmap(struct drm_gem_object *gobj,
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
 	int ret;
 
+	XDNA_DBG(xdna, "BO map_offset 0x%llx type %d userptr 0x%lx size 0x%lx",
+		 drm_vma_node_offset_addr(&gobj->vma_node), abo->type,
+		 vma->vm_start, gobj->size);
+
 	ret = amdxdna_hmm_register(abo, vma);
 	if (ret)
 		return ret;
@@ -514,9 +549,6 @@ static int amdxdna_gem_obj_mmap(struct drm_gem_object *gobj,
 		goto hmm_unreg;
 	}
 
-	XDNA_DBG(xdna, "BO map_offset 0x%llx type %d userptr 0x%lx size 0x%lx",
-		 drm_vma_node_offset_addr(&gobj->vma_node), abo->type,
-		 vma->vm_start, gobj->size);
 	return 0;
 
 hmm_unreg:
@@ -524,14 +556,99 @@ hmm_unreg:
 	return ret;
 }
 
+/*
+ * VM operations for amdxdna shmem VMAs that use VM_MIXEDMAP.
+ *
+ * drm_gem_shmem_vm_ops cannot be used on VM_MIXEDMAP VMAs because its fault
+ * handler calls vmf_insert_pfn() → vmf_insert_pfn_prot() which contains:
+ *   BUG_ON((vma->vm_flags & VM_MIXEDMAP) && pfn_valid(pfn))
+ * All amdxdna shmem pages are ordinary struct pages so pfn_valid() is always
+ * true, making the combination fatal.
+ *
+ * These ops use vmf_insert_page() (struct-page based) instead, which is the
+ * correct API for VM_MIXEDMAP VMAs backed by real struct pages.  The open and
+ * close handlers replicate drm_gem_shmem_vm_open/close using only exported
+ * symbols.
+ */
+static vm_fault_t amdxdna_gem_mixedmap_fault(struct vm_fault *vmf)
+{
+	struct vm_area_struct *vma = vmf->vma;
+	struct drm_gem_object *gobj = vma->vm_private_data;
+	struct drm_gem_shmem_object *shmem = to_drm_gem_shmem_obj(gobj);
+	loff_t num_pages = gobj->size >> PAGE_SHIFT;
+	vm_fault_t ret = VM_FAULT_SIGBUS;
+	pgoff_t page_offset;
+	struct page *page;
+
+	/*
+	 * Partial free of vma is unexpected. Otherwise, the wrong page
+	 * will be faulted in and the user application may crash itself.
+	 */
+	page_offset = linear_page_delta(vma, vmf->address);
+
+	dma_resv_lock(gobj->resv, NULL);
+
+	if (!shmem->pages || shmem->madv < 0 || page_offset >= num_pages)
+		goto out;
+
+	page = shmem->pages[page_offset];
+	if (WARN_ON_ONCE(!page))
+		goto out;
+
+	/*
+	 * Use vmf_insert_page() (struct-page path) not vmf_insert_pfn()
+	 * (PFN path) because this VMA carries VM_MIXEDMAP.
+	 */
+	ret = vmf_insert_page(vma, vmf->address, page);
+	if (ret == VM_FAULT_NOPAGE)
+		folio_mark_accessed(page_folio(page));
+
+out:
+	dma_resv_unlock(gobj->resv);
+	return ret;
+}
+
+static void amdxdna_gem_mixedmap_vm_open(struct vm_area_struct *vma)
+{
+	struct drm_gem_object *gobj = vma->vm_private_data;
+	struct drm_gem_shmem_object *shmem = to_drm_gem_shmem_obj(gobj);
+
+	/*
+	 * Bump pages_use_count so the page array stays alive for the new
+	 * mapping copy created by fork().  Mirrors drm_gem_shmem_vm_open().
+	 */
+	dma_resv_lock(gobj->resv, NULL);
+	drm_WARN_ON_ONCE(gobj->dev, !refcount_inc_not_zero(&shmem->pages_use_count));
+	dma_resv_unlock(gobj->resv);
+
+	drm_gem_vm_open(vma);
+}
+
+static void amdxdna_gem_mixedmap_vm_close(struct vm_area_struct *vma)
+{
+	struct drm_gem_object *gobj = vma->vm_private_data;
+	struct drm_gem_shmem_object *shmem = to_drm_gem_shmem_obj(gobj);
+
+	dma_resv_lock(gobj->resv, NULL);
+	drm_gem_shmem_put_pages_locked(shmem);
+	dma_resv_unlock(gobj->resv);
+
+	drm_gem_vm_close(vma);
+}
+
+static const struct vm_operations_struct amdxdna_gem_mixedmap_vm_ops = {
+	.fault  = amdxdna_gem_mixedmap_fault,
+	.open   = amdxdna_gem_mixedmap_vm_open,
+	.close  = amdxdna_gem_mixedmap_vm_close,
+};
+
 static int amdxdna_gem_dmabuf_mmap(struct dma_buf *dma_buf, struct vm_area_struct *vma)
 {
 	struct drm_gem_object *gobj = dma_buf->priv;
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
-	unsigned long num_pages = vma_pages(vma);
 	int ret;
 
-	vma->vm_ops = &drm_gem_shmem_vm_ops;
+	vma->vm_ops = &amdxdna_gem_mixedmap_vm_ops;
 	vma->vm_private_data = gobj;
 
 	drm_gem_object_get(gobj);
@@ -541,16 +658,9 @@ static int amdxdna_gem_dmabuf_mmap(struct dma_buf *dma_buf, struct vm_area_struc
 
 	/* The buffer is based on memory pages. Fix the flag. */
 	vm_flags_mod(vma, VM_MIXEDMAP, VM_PFNMAP);
-	ret = vm_insert_pages(vma, vma->vm_start, abo->base.pages,
-			      &num_pages);
-	if (ret)
-		goto close_vma;
 
 	return 0;
 
-close_vma:
-	vma->vm_ops->close(vma);
-	return ret;
 put_obj:
 	drm_gem_object_put(gobj);
 	return ret;
@@ -609,10 +719,8 @@ amdxdna_gem_skip_bo_usage(struct amdxdna_gem_obj *abo)
 }
 
 static void
-amdxdna_gem_add_bo_usage(struct amdxdna_gem_obj *abo)
+amdxdna_gem_add_bo_usage(struct amdxdna_client *client, struct amdxdna_gem_obj *abo)
 {
-	struct amdxdna_client *client = abo->client;
-
 	if (amdxdna_gem_skip_bo_usage(abo))
 		return;
 
@@ -624,10 +732,8 @@ amdxdna_gem_add_bo_usage(struct amdxdna_gem_obj *abo)
 }
 
 static void
-amdxdna_gem_del_bo_usage(struct amdxdna_gem_obj *abo)
+amdxdna_gem_del_bo_usage(struct amdxdna_client *client, struct amdxdna_gem_obj *abo)
 {
-	struct amdxdna_client *client = abo->client;
-
 	if (amdxdna_gem_skip_bo_usage(abo))
 		return;
 
@@ -643,8 +749,9 @@ static void amdxdna_gem_obj_free(struct drm_gem_object *gobj)
 	struct amdxdna_dev *xdna = to_xdna_dev(gobj->dev);
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
 
-	amdxdna_hmm_unregister(abo, NULL);
-	flush_workqueue(xdna->notifier_wq);
+	/* No notifier survives this, so no new work can be queued. */
+	amdxdna_hmm_unreg_umaps(abo, true);
+	cancel_work_sync(&abo->hmm_unreg_work);
 
 	if (abo->pinned)
 		amdxdna_gem_unpin(abo);
@@ -663,12 +770,20 @@ static int amdxdna_gem_obj_open(struct drm_gem_object *gobj, struct drm_file *fi
 {
 	struct amdxdna_dev *xdna = to_xdna_dev(gobj->dev);
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
+	struct amdxdna_client *client;
 	int ret;
 
-	guard(mutex)(&abo->lock);
+	mutex_lock(&abo->lock);
+	if (abo->open_ref > 0 && filp->driver_priv != abo->client) {
+		mutex_unlock(&abo->lock);
+		return -EPERM;
+	}
+
 	abo->open_ref++;
-	if (abo->open_ref > 1)
+	if (abo->open_ref > 1) {
+		mutex_unlock(&abo->lock);
 		return 0;
+	}
 
 	/* Attached to the client when first opened by it. */
 	abo->client = filp->driver_priv;
@@ -679,26 +794,34 @@ static int amdxdna_gem_obj_open(struct drm_gem_object *gobj, struct drm_file *fi
 		if (ret) {
 			abo->open_ref--;
 			abo->client = NULL;
+			mutex_unlock(&abo->lock);
 			return ret;
 		}
 	}
+	client = abo->client;
+	mutex_unlock(&abo->lock);
 
-	amdxdna_gem_add_bo_usage(abo);
+	amdxdna_gem_add_bo_usage(client, abo);
 	return 0;
 }
 
 static void amdxdna_gem_obj_close(struct drm_gem_object *gobj, struct drm_file *filp)
 {
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
+	struct amdxdna_client *client = NULL;
 
-	guard(mutex)(&abo->lock);
+	mutex_lock(&abo->lock);
 	abo->open_ref--;
 
 	if (abo->open_ref == 0) {
-		amdxdna_gem_del_bo_usage(abo);
 		/* Detach from the client when last closed by it. */
+		client = abo->client;
 		abo->client = NULL;
 	}
+	mutex_unlock(&abo->lock);
+
+	if (client)
+		amdxdna_gem_del_bo_usage(client, abo);
 }
 
 static int amdxdna_gem_obj_vmap(struct drm_gem_object *obj, struct iosys_map *map)
@@ -734,19 +857,72 @@ static void amdxdna_gem_obj_vunmap(struct drm_gem_object *obj, struct iosys_map
 		drm_gem_shmem_object_vunmap(obj, map);
 }
 
+static int amdxdna_gem_dev_obj_open(struct drm_gem_object *gobj, struct drm_file *filp)
+{
+	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
+
+	guard(mutex)(&abo->lock);
+	if (filp->driver_priv != abo->client)
+		return -EPERM;
+	abo->open_ref++;
+
+	return 0;
+}
+
+static void amdxdna_gem_dev_obj_close(struct drm_gem_object *gobj, struct drm_file *filp)
+{
+	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
+	struct amdxdna_client *client = NULL;
+
+	mutex_lock(&abo->lock);
+	abo->open_ref--;
+
+	/*
+	 * Freeing the heap allocation here, when the handle is closed, is
+	 * intentional.  DEV BOs are carved out of a per-client drm_mm heap;
+	 * any subsequent allocation that lands on the same device address will
+	 * also belong to the same client.  If the user closes the handle while
+	 * a job is still in flight the only consequence is self-inflicted
+	 * corruption within their own context -- it cannot affect other
+	 * processes.  The GEM reference held by the in-flight job keeps the
+	 * amdxdna_gem_obj struct alive until the job completes; it does not
+	 * prevent the device address from being reclaimed by the allocator.
+	 *
+	 * Cross-process handle creation for DEV BOs is rejected in
+	 * amdxdna_gem_dev_obj_open(), which prevents the following UAF:
+	 * if a second process shared the handle via GEM flink and the
+	 * original creator exited (freeing client), the importer would later
+	 * reach open_ref == 0 here and call amdxdna_gem_heap_free() with a
+	 * dangling abo->client pointer.  Because cross-process opens are
+	 * rejected, the process arriving here is always the owning client,
+	 * which is still alive.  abo->client is nulled out afterwards so that
+	 * any code path running on a lingering GEM reference (e.g. an
+	 * in-flight job) cannot silently dereference a stale pointer.
+	 */
+	if (abo->open_ref == 0) {
+		client = abo->client;
+		abo->client = NULL;
+	}
+	mutex_unlock(&abo->lock);
+
+	if (client)
+		amdxdna_gem_heap_free(client, abo);
+}
+
 static int amdxdna_gem_dev_obj_vmap(struct drm_gem_object *obj, struct iosys_map *map)
 {
 	struct amdxdna_gem_obj *abo = to_xdna_obj(obj);
 	struct amdxdna_gem_obj *heap;
+	unsigned long index = 0;
 	void *base;
 	u64 offset;
 
-	/* vmap dev bo which is across more than 1 heap is not allowed */
-	if (abo->heap_start_id != abo->heap_end_id)
+	heap = xa_find(&abo->heap_xa, &index, ULONG_MAX, XA_PRESENT);
+	if (!heap)
 		return -ENOMEM;
 
-	heap = xa_load(&abo->client->dev_heap_xa, abo->heap_start_id);
-	if (!heap)
+	/* vmap dev bo which is across more than 1 heap is not allowed */
+	if (xa_find_after(&abo->heap_xa, &index, ULONG_MAX, XA_PRESENT))
 		return -ENOMEM;
 
 	base = amdxdna_gem_vmap(heap);
@@ -765,6 +941,8 @@ static struct dma_buf *amdxdna_gem_dev_obj_export(struct drm_gem_object *gobj, i
 
 static const struct drm_gem_object_funcs amdxdna_gem_dev_obj_funcs = {
 	.free = amdxdna_gem_dev_obj_free,
+	.open = amdxdna_gem_dev_obj_open,
+	.close = amdxdna_gem_dev_obj_close,
 	.vmap = amdxdna_gem_dev_obj_vmap,
 	.export = amdxdna_gem_dev_obj_export,
 };
@@ -780,7 +958,7 @@ static const struct drm_gem_object_funcs amdxdna_gem_shmem_funcs = {
 	.vmap = amdxdna_gem_obj_vmap,
 	.vunmap = amdxdna_gem_obj_vunmap,
 	.mmap = amdxdna_gem_obj_mmap,
-	.vm_ops = &drm_gem_shmem_vm_ops,
+	.vm_ops = &amdxdna_gem_mixedmap_vm_ops,
 	.export = amdxdna_gem_prime_export,
 };
 
@@ -962,6 +1140,11 @@ amdxdna_drm_create_dev_heap_bo(struct drm_device *dev,
 	struct amdxdna_gem_obj *abo;
 	int ret;
 
+	if (!xdna->dev_info->dev_mem_size) {
+		XDNA_DBG(xdna, "Device heap is not supported");
+		return ERR_PTR(-EOPNOTSUPP);
+	}
+
 	WARN_ON(!is_power_of_2(xdna->dev_info->dev_mem_size));
 	XDNA_DBG(xdna, "Requested dev heap size 0x%llx", args->size);
 	if (!args->size || !IS_ALIGNED(args->size, xdna->dev_info->dev_mem_size)) {
@@ -1103,6 +1286,8 @@ int amdxdna_drm_create_bo_ioctl(struct drm_device *dev, void *data, struct drm_f
 		 args->handle, args->type, amdxdna_gem_uva(abo),
 		 amdxdna_gem_dev_addr(abo), abo->mem.size);
 put_obj:
+	if (ret && abo->type == AMDXDNA_BO_DEV)
+		amdxdna_gem_heap_free(client, abo);
 	/* Dereference object reference. Handle holds it now. */
 	drm_gem_object_put(to_gobj(abo));
 	return ret;
@@ -1136,7 +1321,6 @@ static void amdxdna_bo_unpin(struct amdxdna_gem_obj *abo)
 
 int amdxdna_gem_pin_nolock(struct amdxdna_gem_obj *abo)
 {
-	struct amdxdna_client *client = abo->client;
 	struct amdxdna_gem_obj *heap;
 	unsigned long heap_id, last = ULONG_MAX;
 	int ret = 0;
@@ -1144,17 +1328,15 @@ int amdxdna_gem_pin_nolock(struct amdxdna_gem_obj *abo)
 	if (abo->type != AMDXDNA_BO_DEV)
 		return amdxdna_bo_pin(abo);
 
-	xa_for_each_range(&client->dev_heap_xa, heap_id, heap,
-			  abo->heap_start_id, abo->heap_end_id) {
+	xa_for_each(&abo->heap_xa, heap_id, heap) {
 		ret = amdxdna_bo_pin(heap);
 		if (ret)
 			break;
 		last = heap_id;
 	}
 
-	if (ret && last <= abo->heap_end_id) {
-		xa_for_each_range(&client->dev_heap_xa, heap_id, heap,
-				  abo->heap_start_id, last)
+	if (ret && last != ULONG_MAX) {
+		xa_for_each_range(&abo->heap_xa, heap_id, heap, 0, last)
 			amdxdna_bo_unpin(heap);
 	}
 
@@ -1179,8 +1361,7 @@ void amdxdna_gem_unpin(struct amdxdna_gem_obj *abo)
 		struct amdxdna_gem_obj *heap;
 		unsigned long heap_id;
 
-		xa_for_each_range(&abo->client->dev_heap_xa, heap_id, heap,
-				  abo->heap_start_id, abo->heap_end_id)
+		xa_for_each(&abo->heap_xa, heap_id, heap)
 			amdxdna_bo_unpin(heap);
 	} else {
 		amdxdna_bo_unpin(abo);
@@ -1300,8 +1481,7 @@ int amdxdna_drm_sync_bo_ioctl(struct drm_device *dev,
 		u64 flush_start = bo_start + args->offset;
 		u64 flush_end = flush_start + args->size;
 
-		xa_for_each_range(&client->dev_heap_xa, heap_id, heap,
-				  abo->heap_start_id, abo->heap_end_id) {
+		xa_for_each(&abo->heap_xa, heap_id, heap) {
 			u64 heap_start = amdxdna_gem_dev_addr(heap);
 			u64 heap_end = heap_start + heap->mem.size;
 			u64 start = max(flush_start, heap_start);
@@ -1328,7 +1508,7 @@ int amdxdna_drm_sync_bo_ioctl(struct drm_device *dev,
 		amdxdna_gem_unpin(abo);
 
 		if (ret) {
-			drm_WARN(&xdna->ddev, 1, "Can not get flush memory");
+			XDNA_DBG(xdna, "Flush BO %d failed, ret %d", args->handle, ret);
 			goto put_obj;
 		}
 	}
@@ -1336,7 +1516,8 @@ int amdxdna_drm_sync_bo_ioctl(struct drm_device *dev,
 	XDNA_DBG(xdna, "Sync bo %d offset 0x%llx, size 0x%llx\n",
 		 args->handle, args->offset, args->size);
 
-	if (args->direction == SYNC_DIRECT_FROM_DEVICE)
+	if (abo->assigned_hwctx != AMDXDNA_INVALID_CTX_HANDLE &&
+	    args->direction == SYNC_DIRECT_FROM_DEVICE)
 		ret = amdxdna_hwctx_sync_debug_bo(client, args->handle);
 
 put_obj: