amdxdna/vxdna: support multiple virtio-gpu guest clients on the NPU

drivers/accel/amdxdna/aie2_message.c

@@ -947,7 +947,7 @@ static struct aie2_exec_msg_ops npu_exec_message_ops = {
 static int aie2_init_exec_req(void *req, struct amdxdna_gem_obj *cmd_abo,
 			      size_t *size, u32 *msg_op)
 {
-	struct amdxdna_dev *xdna = cmd_abo->client->xdna;
+	struct amdxdna_dev *xdna = to_xdna_dev(to_gobj(cmd_abo)->dev);
 	int ret;
 	u32 op;
 
@@ -981,7 +981,7 @@ static int
 aie2_cmdlist_fill_slot(void *slot, struct amdxdna_gem_obj *cmd_abo,
 		       size_t *size, u32 *cmd_op)
 {
-	struct amdxdna_dev *xdna = cmd_abo->client->xdna;
+	struct amdxdna_dev *xdna = to_xdna_dev(to_gobj(cmd_abo)->dev);
 	int ret;
 	u32 op;
 

drivers/accel/amdxdna/amdxdna_ctx.c

@@ -214,21 +214,63 @@ int amdxdna_cmd_set_error(struct amdxdna_gem_obj *abo,
 	return 0;
 }
 
+#define AMDXDNA_MIN_HWCTX_ID	(AMDXDNA_INVALID_CTX_HANDLE + 1)
+
+/*
+ * Allocate a device-wide-unique hwctx ID cyclically.
+ *
+ * The ID must be unique across the whole device (not just per client) so a
+ * context is identified unambiguously even when one process opens the device
+ * more than once. Allocation is therefore done from a device-wide IDA, while
+ * the id->hwctx map stays per-client (amdxdna_client.hwctx_xa) so lookups only
+ * ever see the caller's own contexts.
+ *
+ * The allocation must be cyclic. IDA has no cyclic variant (only xa_alloc_cyclic
+ * does), so keep an explicit @next_hwctxid cursor and search [cursor, MAX] then
+ * wrap to [MIN, MAX]. Cyclic (rather than IDA's default lowest-free) is required
+ * because the virtio-gpu guest shim derives its ring index from this ID by
+ * modulo and virtio-gpu exposes only 64 rings: not reusing a just-freed ID keeps
+ * a new context off a ring slot whose previous owner may not be fully torn down
+ * yet. The IDA has its own internal lock; @next_hwctxid is serialized by
+ * dev_lock, which the only caller (create) holds.
+ */
+static int amdxdna_hwctx_id_alloc(struct amdxdna_dev *xdna)
+{
+	int id;
+
+	if (xdna->next_hwctxid < AMDXDNA_MIN_HWCTX_ID)
+		xdna->next_hwctxid = AMDXDNA_MIN_HWCTX_ID;
+
+	id = ida_alloc_range(&xdna->hwctx_ida, xdna->next_hwctxid,
+			     MAX_HWCTX_ID, GFP_KERNEL);
+	if (id == -ENOSPC)
+		id = ida_alloc_range(&xdna->hwctx_ida, AMDXDNA_MIN_HWCTX_ID,
+				     MAX_HWCTX_ID, GFP_KERNEL);
+	if (id >= 0)
+		xdna->next_hwctxid = (id >= MAX_HWCTX_ID) ?
+			AMDXDNA_MIN_HWCTX_ID : id + 1;
+
+	return id;
+}
+
 /*
  * This should be called in close() and remove(). DO NOT call in other syscalls.
  * This guarantee that when hwctx and resources will be released, if user
  * doesn't call amdxdna_drm_destroy_hwctx_ioctl.
  */
 void amdxdna_hwctx_remove_all(struct amdxdna_client *client)
 {
+	struct amdxdna_dev *xdna = client->xdna;
 	struct amdxdna_hwctx *hwctx;
 	unsigned long hwctx_id;
 
 	amdxdna_for_each_hwctx(client, hwctx_id, hwctx) {
-		XDNA_DBG(client->xdna, "PID %d close HW context %d",
-			 client->pid, hwctx->id);
-		xa_erase(&client->hwctx_xa, hwctx->id);
+		u32 id = hwctx->id;
+
+		XDNA_DBG(xdna, "PID %d close HW context %d", client->pid, id);
+		xa_erase(&client->hwctx_xa, id);
 		amdxdna_hwctx_destroy_rcu(hwctx, &client->hwctx_srcu);
+		ida_free(&xdna->hwctx_ida, id);
 	}
 }
 
@@ -283,13 +325,19 @@ int amdxdna_drm_create_hwctx_ioctl(struct drm_device *dev, void *data, struct dr
 		goto fini_hwctx;
 	}
 
-	ret = xa_alloc_cyclic(&client->hwctx_xa, &hwctx->id, hwctx,
-			      XA_LIMIT(AMDXDNA_INVALID_CTX_HANDLE + 1, MAX_HWCTX_ID),
-			      &client->next_hwctxid, GFP_KERNEL);
+	ret = amdxdna_hwctx_id_alloc(xdna);
 	if (ret < 0) {
 		XDNA_ERR(xdna, "Allocate hwctx ID failed, ret %d", ret);
 		goto free_name;
 	}
+	hwctx->id = ret;
+
+	/* Publish into the per-client map; submit/wait/etc. look up here. */
+	ret = xa_err(xa_store(&client->hwctx_xa, hwctx->id, hwctx, GFP_KERNEL));
+	if (ret) {
+		XDNA_ERR(xdna, "Store hwctx %d failed, ret %d", hwctx->id, ret);
+		goto free_id;
+	}
 
 	args->handle = hwctx->id;
 	args->syncobj_handle = hwctx->syncobj_hdl;
@@ -301,6 +349,8 @@ int amdxdna_drm_create_hwctx_ioctl(struct drm_device *dev, void *data, struct dr
 	drm_dev_exit(idx);
 	return 0;
 
+free_id:
+	ida_free(&xdna->hwctx_ida, hwctx->id);
 free_name:
 	kfree(hwctx->name);
 fini_hwctx:
@@ -342,6 +392,7 @@ int amdxdna_drm_destroy_hwctx_ioctl(struct drm_device *dev, void *data, struct d
 	 * SRCU to synchronize with exec command ioctls.
 	 */
 	amdxdna_hwctx_destroy_rcu(hwctx, &client->hwctx_srcu);
+	ida_free(&xdna->hwctx_ida, args->handle);
 
 	XDNA_DBG(xdna, "PID %d destroyed HW context %d", client->pid, args->handle);
 out:
@@ -641,6 +692,11 @@ int amdxdna_cmd_submit(struct amdxdna_client *client,
 		goto put_bos;
 	}
 
+	/*
+	 * The lookup is in the caller's own hwctx_xa, so it can only ever return
+	 * this client's context; the per-client hwctx_srcu keeps it alive across
+	 * the blocking submit (drained by amdxdna_hwctx_destroy_rcu()).
+	 */
 	idx = srcu_read_lock(&client->hwctx_srcu);
 	hwctx = xa_load(&client->hwctx_xa, hwctx_hdl);
 	if (!hwctx) {
@@ -749,12 +805,15 @@ int amdxdna_cmd_wait(struct amdxdna_client *client, u32 ctx_hdl,
 	if (!xdna->dev_info->ops->cmd_wait)
 		return -EOPNOTSUPP;
 
-	/* For locking concerns, see amdxdna_drm_exec_cmd_ioctl. */
+	/*
+	 * Per-client lookup + per-client srcu: the wait (possibly unbounded) holds
+	 * only this client's srcu, so a parked waiter stalls only its own client's
+	 * teardown. See amdxdna_drm_exec_cmd_ioctl for the broader locking rationale.
+	 */
 	idx = srcu_read_lock(&client->hwctx_srcu);
 	hwctx = xa_load(&client->hwctx_xa, ctx_hdl);
 	if (!hwctx) {
-		XDNA_DBG(xdna, "PID %d failed to get ctx %d",
-			 client->pid, ctx_hdl);
+		XDNA_DBG(xdna, "PID %d failed to get ctx %d", client->pid, ctx_hdl);
 		ret = -EINVAL;
 		goto unlock_ctx_srcu;
 	}

drivers/accel/amdxdna/amdxdna_gem.c

@@ -218,6 +218,7 @@ amdxdna_gem_destroy_obj(struct amdxdna_gem_obj *abo)
  */
 void *amdxdna_gem_vmap(struct amdxdna_gem_obj *abo)
 {
+	struct amdxdna_dev *xdna = to_xdna_dev(to_gobj(abo)->dev);
 	struct iosys_map map = IOSYS_MAP_INIT_VADDR(NULL);
 	int ret;
 
@@ -230,7 +231,7 @@ void *amdxdna_gem_vmap(struct amdxdna_gem_obj *abo)
 	if (!abo->mem.kva) {
 		ret = drm_gem_vmap(to_gobj(abo), &map);
 		if (ret)
-			XDNA_ERR(abo->client->xdna, "Vmap bo failed, ret %d", ret);
+			XDNA_ERR(xdna, "Vmap bo failed, ret %d", ret);
 		else
 			abo->mem.kva = map.vaddr;
 	}
@@ -374,7 +375,13 @@ static int amdxdna_hmm_register(struct amdxdna_gem_obj *abo,
 	u32 nr_pages;
 	int ret;
 
-	if (!amdxdna_pasid_on(abo->client))
+	/*
+	 * When PASID is off, amdxdna_gem_obj_open() called amdxdna_dma_map_bo()
+	 * and mem.dma_addr is valid; use the DMA address directly and skip HMM.
+	 * Avoid dereferencing abo->client which may be NULL (cleared in close())
+	 * while internal kernel references are still held.
+	 */
+	if (abo->mem.dma_addr != AMDXDNA_INVALID_ADDR)
 		return 0;
 
 	down_read(&xdna->notifier_lock);

drivers/accel/amdxdna/amdxdna_gem.h

@@ -88,12 +88,19 @@ u64 amdxdna_gem_dev_addr(struct amdxdna_gem_obj *abo);
 
 static inline u64 amdxdna_dev_bo_offset(struct amdxdna_gem_obj *abo)
 {
-	return amdxdna_gem_dev_addr(abo) - abo->client->xdna->dev_info->dev_mem_base;
+	return amdxdna_gem_dev_addr(abo) - to_xdna_dev(to_gobj(abo)->dev)->dev_info->dev_mem_base;
 }
 
 static inline u64 amdxdna_obj_dma_addr(struct amdxdna_gem_obj *abo)
 {
-	return amdxdna_pasid_on(abo->client) ? amdxdna_gem_uva(abo) : abo->mem.dma_addr;
+	/*
+	 * amdxdna_gem_obj_open() calls amdxdna_dma_map_bo() only when PASID is
+	 * off, leaving mem.dma_addr at AMDXDNA_INVALID_ADDR when PASID is on.
+	 * Avoid dereferencing abo->client, which is cleared to NULL by
+	 * amdxdna_gem_obj_close() while internal kernel references remain.
+	 */
+	return (abo->mem.dma_addr != AMDXDNA_INVALID_ADDR) ?
+		abo->mem.dma_addr : amdxdna_gem_uva(abo);
 }
 
 void amdxdna_umap_put(struct amdxdna_umap *mapp);

drivers/accel/amdxdna/amdxdna_pci_drv.c

@@ -149,16 +149,18 @@ static void amdxdna_sva_fini(struct amdxdna_client *client)
 static int amdxdna_drm_open(struct drm_device *ddev, struct drm_file *filp)
 {
 	struct amdxdna_dev *xdna = to_xdna_dev(ddev);
-	struct amdxdna_client *tmp, *client;
+	struct amdxdna_client *client;
 	int ret;
 
 	client = kzalloc_obj(*client);
 	if (!client)
 		return -ENOMEM;
 
 	ret = init_srcu_struct(&client->hwctx_srcu);
-	if (ret)
-		goto free_client;
+	if (ret) {
+		kfree(client);
+		return ret;
+	}
 
 	client->pid = pid_nr(rcu_access_pointer(filp->pid));
 	client->xdna = xdna;
@@ -172,14 +174,15 @@ static int amdxdna_drm_open(struct drm_device *ddev, struct drm_file *filp)
 			XDNA_WARN(xdna, "PASID not available for pid %d", client->pid);
 			if (!amdxdna_use_carveout(xdna)) {
 				XDNA_ERR(xdna, "PASID unavailable and carveout not configured");
-				ret = -EINVAL;
-				goto cleanup_srcu;
+				cleanup_srcu_struct(&client->hwctx_srcu);
+				kfree(client);
+				return -EINVAL;
 			}
 		}
 	}
 #endif
 	mmgrab(client->mm);
-	xa_init_flags(&client->hwctx_xa, XA_FLAGS_ALLOC);
+	xa_init(&client->hwctx_xa);
 	xa_init_flags(&client->dev_heap_xa, XA_FLAGS_ALLOC);
 	/* Devices without a managed dev-heap aperture (e.g. PA-mode aie4) leave
 	 * dev_heap_max_size at 0; drm_mm_init() BUGs on a zero-sized range.
@@ -191,15 +194,6 @@ static int amdxdna_drm_open(struct drm_device *ddev, struct drm_file *filp)
 
 	mutex_lock(&xdna->client_lock);
 	mutex_lock(&xdna->dev_lock);
-	amdxdna_for_each_client(xdna, tmp) {
-		if (tmp->pid == client->pid) {
-			mutex_unlock(&xdna->dev_lock);
-			mutex_unlock(&xdna->client_lock);
-			XDNA_WARN(xdna, "pid %d already opened the device", client->pid);
-			ret = -EBUSY;
-			goto fail;
-		}
-	}
 	list_add_tail(&client->node, &xdna->client_list);
 	mutex_unlock(&xdna->dev_lock);
 	mutex_unlock(&xdna->client_lock);
@@ -211,22 +205,6 @@ static int amdxdna_drm_open(struct drm_device *ddev, struct drm_file *filp)
 
 	XDNA_DBG(xdna, "pid %d opened", client->pid);
 	return 0;
-
-fail:
-	if (xdna->dev_info->dev_heap_max_size)
-		drm_mm_takedown(&client->dev_heap_mm);
-	xa_destroy(&client->dev_heap_xa);
-	xa_destroy(&client->hwctx_xa);
-	mutex_destroy(&client->mm_lock);
-	mmdrop(client->mm);
-	amdxdna_sva_fini(client);
-#ifndef AMDXDNA_NPU3A
-cleanup_srcu:
-#endif
-	cleanup_srcu_struct(&client->hwctx_srcu);
-free_client:
-	kfree(client);
-	return ret;
 }
 
 static void amdxdna_client_cleanup(struct amdxdna_client *client)
@@ -448,6 +426,7 @@ static void amdxdna_xdna_drm_release(struct drm_device *drm, void *res)
 
 	amdxdna_carveout_fini(xdna);
 	cleanup_srcu_struct(&xdna->dpt_srcu);
+	ida_destroy(&xdna->hwctx_ida);
 }
 
 static int amdxdna_probe(struct pci_dev *pdev, const struct pci_device_id *id)
@@ -473,6 +452,7 @@ static int amdxdna_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 	drmm_mutex_init(ddev, &xdna->dev_lock);
 	init_rwsem(&xdna->notifier_lock);
 	INIT_LIST_HEAD(&xdna->client_list);
+	ida_init(&xdna->hwctx_ida);
 	pci_set_drvdata(pdev, xdna);
 
 	ret = init_srcu_struct(&xdna->dpt_srcu);

drivers/accel/amdxdna/amdxdna_pci_drv.h

@@ -13,6 +13,7 @@
 #include <drm/drm_print.h>
 #include <linux/capability.h>
 #include <linux/cred.h>
+#include <linux/idr.h>
 #include <linux/iommu.h>
 #include <linux/iova.h>
 #include <linux/srcu.h>
@@ -141,6 +142,14 @@ struct amdxdna_dev {
 	struct mutex			dev_lock; /* per device lock */
 	struct list_head		client_list;
 	struct mutex			client_lock; /* client_list */
+	/*
+	 * Device-wide cyclic allocator for hwctx IDs. Only the ID space is shared
+	 * so a context ID is unique across the whole device; the id->hwctx map and
+	 * its lifetime srcu live per-client (see struct amdxdna_client). next_hwctxid
+	 * is the cyclic cursor. Both are serialized by dev_lock.
+	 */
+	struct ida			hwctx_ida;
+	u32				next_hwctxid;
 	struct amdxdna_drm_query_firmware_version fw_ver;
 	struct rw_semaphore		notifier_lock; /* for mmu notifier*/
 	struct workqueue_struct		*notifier_wq;
@@ -186,9 +195,15 @@ struct amdxdna_io_stats {
 struct amdxdna_client {
 	struct list_head		node;
 	pid_t				pid;
+	/*
+	 * Guards hwctx lifetime against this client's own blocking submit/wait.
+	 * A lookup only ever returns this client's own hwctx (per-client xa), so a
+	 * single per-client srcu is sufficient and a parked waiter can only stall
+	 * its own client's teardown.
+	 */
 	struct srcu_struct		hwctx_srcu;
+	/* id->hwctx map for this client; IDs come from xdna->hwctx_ida. */
 	struct xarray			hwctx_xa;
-	u32				next_hwctxid;
 	struct amdxdna_dev		*xdna;
 	struct drm_file			*filp;
 

drivers/accel/amdxdna/amdxdna_ubuf.c

@@ -149,6 +149,7 @@ struct dma_buf *amdxdna_get_ubuf(struct drm_device *dev,
 	u32 npages, start = 0;
 	struct dma_buf *dbuf;
 	bool readonly = true;
+	bool need_contig;
 	int i, ret = 0;
 	DEFINE_DMA_BUF_EXPORT_INFO(exp_info);
 
@@ -174,6 +175,15 @@ struct dma_buf *amdxdna_get_ubuf(struct drm_device *dev,
 		goto free_ent;
 	}
 
+	/*
+	 * With an IOMMU domain the scattered pages are mapped to a single
+	 * contiguous device IOVA (iommu_map_sgtable), so the entries need not be
+	 * contiguous in user VA. Without one (PASID/SVA or PA mode) the device
+	 * addresses the BO at its user VA and the BO records only the first
+	 * entry's VA, so require the entries to describe one contiguous VA range.
+	 */
+	need_contig = !amdxdna_iova_on(xdna);
+
 	for (i = 0, exp_info.size = 0; i < num_entries; i++) {
 		if (!IS_ALIGNED(va_ent[i].vaddr, PAGE_SIZE) ||
 		    !IS_ALIGNED(va_ent[i].len, PAGE_SIZE)) {
@@ -183,6 +193,15 @@ struct dma_buf *amdxdna_get_ubuf(struct drm_device *dev,
 			goto free_ent;
 		}
 
+		if (need_contig && i &&
+		    va_ent[i].vaddr != va_ent[i - 1].vaddr + va_ent[i - 1].len) {
+			XDNA_ERR(xdna, "Non-contiguous va entry %d, %llx after %llx+%llx",
+				 i, va_ent[i].vaddr, va_ent[i - 1].vaddr,
+				 va_ent[i - 1].len);
+			ret = -EINVAL;
+			goto free_ent;
+		}
+
 		if (check_add_overflow(exp_info.size, va_ent[i].len, &exp_info.size)) {
 			ret = -EINVAL;
 			goto free_ent;