rkopnu: honor vendor iommu_domain_id -> multi-domain NPU-resident memory

The vendor rknpu ABI (rknpu_mem_create.iommu_domain_id,
rknpu_submit.iommu_domain_id) has always let a client spread
allocations across multiple IOMMU paging domains, and our own
ggml-rknpu2 glue (IOMMUDomainManager) already assigns tensors across
up to 16 domains for exactly this reason: each domain is an
independent page-table root with its own ~4GiB IOVA aperture (capped
by the vendor rknn_matmul API's int32 B-matrix offsets, not by IOMMU
hardware). rkopnu was reading and silently dropping this field in both
ioctls, collapsing every allocation into the single domain created
once at open() -- so userspace's domain-spreading was a no-op and any
model with >~4GiB of NPU-resident weights hit ENOSPC ("MEM_CREATE VA
insert failed").

Replace the single per-fd domain with a lazily-created 16-slot pool
keyed by the vendor domain id, each owning its own drm_mm IOVA
allocator (previously one drm_mm shared per-fd; now one per domain,
matching each domain's own independent address space). Wire
iommu_domain_id through in rkopnu_ioctl_mem_create and
rkopnu_ioctl_submit. The domain attach/detach-on-change logic in
rocket_job_run() already existed and needed no changes -- it was
comparing core->attached_domain against job->domain per core with a
skip-if-unchanged optimization, just never fed more than one domain
value before now.

Also fixes two latent bugs surfaced while making domain creation
lazy/fallible (rocket_iommu_domain_get() could never fail before this;
now it can, e.g. -ENOMEM from iommu_paging_domain_alloc):
 - rocket_iommu_domain_put(NULL) would crash; made it NULL-tolerant
   like the dma_fence_put()/kfree() calls next to it, since
   rocket_job_cleanup() unconditionally puts job->domain.
 - MEM_CREATE's error paths never dropped the domain reference on
   failure (harmless before, since the fd's one domain outlived any
   failed allocation regardless; now leaks a whole IOMMU paging domain
   per failed MEM_CREATE under memory pressure).

Verified on boltzmann (RK3588, 7.0.0-rc3-npuclk+): Qwen2.5-3B-Instruct
f16 (5.75 GiB, spans >=2 domains) now loads and runs natively on the
NPU at pp512=76.75 t/s -- previously an unconditional GGML_ASSERT abort
in ggml-rknpu2.cpp at ~4GiB. Qwen3-4B-Q8_0 (single-domain, the
existing production path) is unaffected: no regression.

Known issue, not yet fixed: Qwen3-30B-A3B-Instruct-2507 UD-Q4_K_XL
(17.69 GB, ~5 domains) maps cleanly (refcnt to 1221 dmabufs, zero
VA-insert failures, zero IOMMU faults -- the domain packing itself is
correct) but prefill deadlocks: the benchmark process goes permanently
D-state on dma_fence_default_wait, drm_sched's TDR does not visibly
recover it, and the NPU has to be recovered by reboot (never rmmod
rkopnu -- silent deadlock). Root cause not yet confirmed (lockdown
blocked /proc/PID/stack); a hardening patch is staged separately
pending live validation. Until resolved, avoid >4GiB NPU-resident
models with rapid cross-domain access patterns (dense models spanning
2-3 domains sequentially, like the 3B-f16 case above, are fine).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EWpfhDgYNA21tETDP9ueBE

rkopnu_ioctl.c

@@ -81,31 +81,47 @@ int rkopnu_ioctl_mem_create(struct drm_device *dev, void *data, struct drm_file
 	gem = &shmem->base;
 	bo = to_rocket_bo(gem);
 	bo->driver_priv = priv;
-	bo->domain = rocket_iommu_domain_get(priv);
+	/*
+	 * rkopnu multi-domain: honor the vendor iommu_domain_id instead of
+	 * collapsing every allocation into one shared domain. This is the
+	 * field librknnrt already fills in from rknn_matmul_info.iommu_domain_id
+	 * (see ggml-rknpu2's IOMMUDomainManager) -- previously read and dropped
+	 * on the floor here, which made every allocation land in domain 0
+	 * regardless of what userspace asked for, capping total NPU-resident
+	 * memory at one domain's ~4GiB IOVA aperture.
+	 */
+	bo->domain = rocket_iommu_domain_get(priv, args->iommu_domain_id);
+	if (IS_ERR(bo->domain)) {
+		ret = PTR_ERR(bo->domain);
+		bo->domain = NULL;
+		goto err;
+	}
 	bo->size = args->size;
 	bo->offset = 0;
 
 	sgt = drm_gem_shmem_get_pages_sgt(shmem);
 	if (IS_ERR(sgt)) {
 		ret = PTR_ERR(sgt);
-		goto err;
+		goto err_put_domain;
 	}
 
-	mutex_lock(&priv->mm_lock);
-	ret = drm_mm_insert_node_generic(&priv->mm, &bo->mm, bo->size,
+	mutex_lock(&bo->domain->mm_lock);
+	ret = drm_mm_insert_node_generic(&bo->domain->mm, &bo->mm, bo->size,
 					 PAGE_SIZE, 0, 0);
-	mutex_unlock(&priv->mm_lock);
+	mutex_unlock(&bo->domain->mm_lock);
 	if (ret) {
 		/* Most likely IOMMU VA-aperture exhaustion/fragmentation over a
 		 * long-lived server (the "works for one, dies after N" mode at
-		 * large contexts). Log the request size so it is diagnosable.
-		 * (Fable review, risk 3.) */
-		dev_warn(dev->dev, "rkopnu: MEM_CREATE VA insert failed (%d) for %llu bytes\n",
-			 ret, (unsigned long long)bo->size);
-		goto err;
+		 * large contexts), or genuine exhaustion of this domain's ~4GiB
+		 * budget if userspace kept packing it past capacity. Log the
+		 * domain id too now that there's more than one. (Fable review,
+		 * risk 3.) */
+		dev_warn(dev->dev, "rkopnu: MEM_CREATE VA insert failed (%d) for %llu bytes (domain %d)\n",
+			 ret, (unsigned long long)bo->size, bo->domain->id);
+		goto err_put_domain;
 	}
 
-	ret = iommu_map_sgtable(priv->domain->domain, bo->mm.start, shmem->sgt,
+	ret = iommu_map_sgtable(bo->domain->domain, bo->mm.start, shmem->sgt,
 				IOMMU_READ | IOMMU_WRITE);
 	if (ret < 0 || (u64)ret < args->size) {
 		ret = -ENOMEM;
@@ -125,11 +141,23 @@ int rkopnu_ioctl_mem_create(struct drm_device *dev, void *data, struct drm_file
 	return 0;
 
 err_unmap:
-	iommu_unmap(priv->domain->domain, bo->mm.start, bo->size);
+	iommu_unmap(bo->domain->domain, bo->mm.start, bo->size);
 err_remove_node:
-	mutex_lock(&priv->mm_lock);
+	mutex_lock(&bo->domain->mm_lock);
 	drm_mm_remove_node(&bo->mm);
-	mutex_unlock(&priv->mm_lock);
+	mutex_unlock(&bo->domain->mm_lock);
+err_put_domain:
+	/*
+	 * Pre-existing single-domain code never dropped this reference on any
+	 * error path (the fd's one domain outlived every failed allocation
+	 * anyway). Now that a domain can be created lazily just for this
+	 * failed BO, leaking it here would leak a whole IOMMU paging domain
+	 * (page tables) per failed MEM_CREATE -- worse under exactly the
+	 * memory-pressure conditions that make MEM_CREATE fail in the first
+	 * place. Drop it explicitly.
+	 */
+	rocket_iommu_domain_put(bo->domain);
+	bo->domain = NULL;
 err:
 	drm_gem_shmem_object_free(gem);
 	return ret;

rocket_drv.c

@@ -57,25 +57,110 @@ rocket_iommu_domain_create(struct device *dev)
 	return domain;
 }
 
+/*
+ * rkopnu multi-domain: get-or-create the domain for domain_id in this fd's
+ * pool. Out-of-range or negative ids (old callers, or a malformed vendor
+ * struct) fall back to domain 0 rather than erroring -- domain 0 always
+ * exists once anything has been allocated, matching the pre-multi-domain
+ * behaviour for callers that never think about domain ids at all.
+ *
+ * iommu_paging_domain_alloc() can sleep, so the allocation itself happens
+ * outside domains_lock; the lock only guards the pool array's slot
+ * install/lookup. That opens a create-create race between two threads
+ * requesting the same not-yet-existing domain id concurrently, handled by
+ * re-checking the slot after allocating and discarding the loser's spare
+ * domain via rocket_iommu_domain_put() (kref drops straight to 0, since
+ * nothing else has seen it yet).
+ */
 struct rocket_iommu_domain *
-rocket_iommu_domain_get(struct rocket_file_priv *rocket_priv)
+rocket_iommu_domain_get(struct rocket_file_priv *rocket_priv, s32 domain_id)
 {
-	kref_get(&rocket_priv->domain->kref);
-	return rocket_priv->domain;
+	struct rocket_iommu_domain *domain, *candidate;
+	u64 start, end;
+
+	if (domain_id < 0 || domain_id >= ROCKET_MAX_IOMMU_DOMAINS)
+		domain_id = 0;
+
+	mutex_lock(&rocket_priv->domains_lock);
+	domain = rocket_priv->domains[domain_id];
+	if (domain)
+		kref_get(&domain->kref);
+	mutex_unlock(&rocket_priv->domains_lock);
+	if (domain)
+		return domain;
+
+	candidate = rocket_iommu_domain_create(rocket_priv->rdev->cores[0].dev);
+	if (IS_ERR(candidate))
+		return candidate;
+
+	candidate->id = domain_id;
+	start = candidate->domain->geometry.aperture_start;
+	end = candidate->domain->geometry.aperture_end;
+	drm_mm_init(&candidate->mm, start, end - start + 1);
+	mutex_init(&candidate->mm_lock);
+
+	mutex_lock(&rocket_priv->domains_lock);
+	domain = rocket_priv->domains[domain_id];
+	if (domain) {
+		/* Lost the race: another thread installed this id first. */
+		kref_get(&domain->kref);
+		mutex_unlock(&rocket_priv->domains_lock);
+		rocket_iommu_domain_put(candidate);
+		return domain;
+	}
+
+	rocket_priv->domains[domain_id] = candidate;
+	/* Caller's reference, on top of the pool's own (kref=1 from create,
+	 * dropped by rocket_iommu_domains_close() at postclose). */
+	kref_get(&candidate->kref);
+	mutex_unlock(&rocket_priv->domains_lock);
+
+	return candidate;
 }
 
 void
 rocket_iommu_domain_put(struct rocket_iommu_domain *domain)
 {
+	/* NULL-tolerant like dma_fence_put()/kfree() elsewhere in this driver:
+	 * rocket_job_cleanup() unconditionally puts job->domain, which is NULL
+	 * on a job that failed before rocket_iommu_domain_get() ever ran
+	 * (e.g. drm_sched_job_init() failure) or that failed the get() itself. */
+	if (!domain)
+		return;
 	kref_put(&domain->kref, rocket_iommu_domain_destroy);
 }
 
+void
+rocket_iommu_domains_close(struct rocket_file_priv *rocket_priv)
+{
+	unsigned int i;
+
+	for (i = 0; i < ROCKET_MAX_IOMMU_DOMAINS; i++) {
+		struct rocket_iommu_domain *domain = rocket_priv->domains[i];
+
+		if (!domain)
+			continue;
+
+		rocket_priv->domains[i] = NULL;
+
+		/*
+		 * Safe only because DRM core has already released this file's
+		 * GEM objects before calling .postclose (rocket_postclose()
+		 * calls us), so every domain's drm_mm is provably empty here --
+		 * same invariant the original single-domain code relied on for
+		 * its drm_mm_takedown() at postclose.
+		 */
+		drm_mm_takedown(&domain->mm);
+		mutex_destroy(&domain->mm_lock);
+		rocket_iommu_domain_put(domain);
+	}
+}
+
 static int
 rocket_open(struct drm_device *dev, struct drm_file *file)
 {
 	struct rocket_device *rdev = to_rocket_device(dev);
 	struct rocket_file_priv *rocket_priv;
-	u64 start, end;
 	int ret;
 
 	if (!try_module_get(THIS_MODULE))
@@ -88,30 +173,20 @@ rocket_open(struct drm_device *dev, struct drm_file *file)
 	}
 
 	rocket_priv->rdev = rdev;
-	rocket_priv->domain = rocket_iommu_domain_create(rdev->cores[0].dev);
-	if (IS_ERR(rocket_priv->domain)) {
-		ret = PTR_ERR(rocket_priv->domain);
-		goto err_free;
-	}
+	mutex_init(&rocket_priv->domains_lock);
+	/* Domains are created lazily on first MEM_CREATE/SUBMIT reference to a
+	 * given iommu_domain_id -- see rocket_iommu_domain_get(). */
 
 	file->driver_priv = rocket_priv;
 
-	start = rocket_priv->domain->domain->geometry.aperture_start;
-	end = rocket_priv->domain->domain->geometry.aperture_end;
-	drm_mm_init(&rocket_priv->mm, start, end - start + 1);
-	mutex_init(&rocket_priv->mm_lock);
-
 	ret = rocket_job_open(rocket_priv);
 	if (ret)
-		goto err_mm_takedown;
+		goto err_destroy_lock;
 
 	return 0;
 
-err_mm_takedown:
-	mutex_destroy(&rocket_priv->mm_lock);
-	drm_mm_takedown(&rocket_priv->mm);
-	rocket_iommu_domain_put(rocket_priv->domain);
-err_free:
+err_destroy_lock:
+	mutex_destroy(&rocket_priv->domains_lock);
 	kfree(rocket_priv);
 err_put_mod:
 	module_put(THIS_MODULE);
@@ -124,9 +199,8 @@ rocket_postclose(struct drm_device *dev, struct drm_file *file)
 	struct rocket_file_priv *rocket_priv = file->driver_priv;
 
 	rocket_job_close(rocket_priv);
-	mutex_destroy(&rocket_priv->mm_lock);
-	drm_mm_takedown(&rocket_priv->mm);
-	rocket_iommu_domain_put(rocket_priv->domain);
+	rocket_iommu_domains_close(rocket_priv);
+	mutex_destroy(&rocket_priv->domains_lock);
 	kfree(rocket_priv);
 	module_put(THIS_MODULE);
 }

rocket_drv.h

@@ -12,17 +12,52 @@
 
 extern const struct dev_pm_ops rocket_pm_ops;
 
+/*
+ * rkopnu multi-domain: the vendor rknpu ABI (rknpu_mem_create.iommu_domain_id,
+ * rknpu_submit.iommu_domain_id -- see rknpu_ioctl.h) lets a client spread
+ * allocations across multiple IOMMU paging domains. Each domain is a fully
+ * independent page-table root, so each gets its own ~4GiB IOVA aperture (the
+ * vendor rknn_matmul API's B-matrix offsets are int32, which is the real
+ * ceiling -- not the IOMMU hardware, which is per-domain). Before this, every
+ * allocation on a given fd landed in the single domain created at open(), so
+ * userspace's domain-spreading (ggml-rknpu2's IOMMUDomainManager) was a
+ * no-op and any model with >~4GiB of NPU-resident weights hit ENOSPC. Each
+ * domain now owns its own drm_mm IOVA allocator (mm/mm_lock below), matching
+ * its own independent address space -- a shared allocator across domains
+ * would double-book identical IOVA ranges in different page tables, which is
+ * fine as far as the IOMMU is concerned but wrong for this driver's
+ * bookkeeping (VA aperture accounting, exhaustion diagnostics).
+ */
 struct rocket_iommu_domain {
 	struct iommu_domain *domain;
 	struct kref kref;
+
+	/* vendor iommu_domain_id this slot was created for (0..ROCKET_MAX_IOMMU_DOMAINS-1) */
+	s32 id;
+
+	struct drm_mm mm;
+	struct mutex mm_lock;
 };
 
+/*
+ * Matches ggml-rknpu2's IOMMUDomainManager, which dynamically assigns
+ * domain ids 0..15 (see assign_domain_memory() in ggml-rknpu2.cpp). Any
+ * out-of-range or unset (old-caller) id falls back to domain 0.
+ */
+#define ROCKET_MAX_IOMMU_DOMAINS 16
+
 struct rocket_file_priv {
 	struct rocket_device *rdev;
 
-	struct rocket_iommu_domain *domain;
-	struct drm_mm mm;
-	struct mutex mm_lock;
+	/*
+	 * Lazily-created pool of per-fd IOMMU domains, indexed by the vendor
+	 * iommu_domain_id. Slots are NULL until first referenced by a
+	 * MEM_CREATE or SUBMIT that names them. domains_lock protects the
+	 * pool array itself (slot install/lookup); each domain's own
+	 * mm_lock protects its drm_mm, and its kref its lifetime.
+	 */
+	struct mutex domains_lock;
+	struct rocket_iommu_domain *domains[ROCKET_MAX_IOMMU_DOMAINS];
 
 	struct drm_sched_entity sched_entity;
 
@@ -34,7 +69,21 @@ struct rocket_file_priv {
 	struct drm_sched_entity core_entity[3];
 };
 
-struct rocket_iommu_domain *rocket_iommu_domain_get(struct rocket_file_priv *rocket_priv);
+/*
+ * Returns a referenced (kref_get'd) domain for domain_id, creating it on
+ * first use. Caller must rocket_iommu_domain_put() when done. Can return
+ * ERR_PTR (e.g. -ENOMEM from iommu_paging_domain_alloc) since domain
+ * creation is now lazy instead of guaranteed-at-open -- callers that were
+ * written against the old infallible single-domain version need updating,
+ * not just recompiling.
+ */
+struct rocket_iommu_domain *rocket_iommu_domain_get(struct rocket_file_priv *rocket_priv, s32 domain_id);
 void rocket_iommu_domain_put(struct rocket_iommu_domain *domain);
 
+/* Tears down every domain this fd ever created. Call from postclose, after
+ * DRM core has already released this file's GEM objects (so every domain's
+ * drm_mm is provably empty -- see rocket_open()/rocket_postclose() in
+ * rocket_drv.c for the ordering this relies on). */
+void rocket_iommu_domains_close(struct rocket_file_priv *rocket_priv);
+
 #endif

rocket_gem.c

@@ -14,7 +14,6 @@
 static void rocket_gem_bo_free(struct drm_gem_object *obj)
 {
 	struct rocket_gem_object *bo = to_rocket_bo(obj);
-	struct rocket_file_priv *rocket_priv = bo->driver_priv;
 	size_t unmapped;
 
 	drm_WARN_ON(obj->dev, refcount_read(&bo->base.pages_use_count) > 1);
@@ -22,9 +21,11 @@ static void rocket_gem_bo_free(struct drm_gem_object *obj)
 	unmapped = iommu_unmap(bo->domain->domain, bo->mm.start, bo->size);
 	drm_WARN_ON(obj->dev, unmapped != bo->size);
 
-	mutex_lock(&rocket_priv->mm_lock);
+	/* rkopnu multi-domain: this BO's own domain owns its drm_mm now, not
+	 * a single per-fd allocator (rocket_drv.h). */
+	mutex_lock(&bo->domain->mm_lock);
 	drm_mm_remove_node(&bo->mm);
-	mutex_unlock(&rocket_priv->mm_lock);
+	mutex_unlock(&bo->domain->mm_lock);
 
 	rocket_iommu_domain_put(bo->domain);
 	bo->domain = NULL;
@@ -75,25 +76,39 @@ int rocket_ioctl_create_bo(struct drm_device *dev, void *data, struct drm_file *
 	rkt_obj = to_rocket_bo(gem_obj);
 
 	rkt_obj->driver_priv = rocket_priv;
-	rkt_obj->domain = rocket_iommu_domain_get(rocket_priv);
+	/*
+	 * rkopnu multi-domain: the native rocket uAPI (this function) has no
+	 * domain-id concept in its own struct -- it predates multi-domain and
+	 * isn't wired into rkopnu_ioctls[] (only the rknpu-compat ABI in
+	 * rkopnu_ioctl.c is registered; this path is currently unreachable
+	 * from userspace). Always domain 0 here, matching old single-domain
+	 * behaviour, so this stays correct (if it's ever re-registered) rather
+	 * than silently wrong.
+	 */
+	rkt_obj->domain = rocket_iommu_domain_get(rocket_priv, 0);
+	if (IS_ERR(rkt_obj->domain)) {
+		ret = PTR_ERR(rkt_obj->domain);
+		rkt_obj->domain = NULL;
+		goto err;
+	}
 	rkt_obj->size = args->size;
 	rkt_obj->offset = 0;
 
 	sgt = drm_gem_shmem_get_pages_sgt(shmem_obj);
 	if (IS_ERR(sgt)) {
 		ret = PTR_ERR(sgt);
-		goto err;
+		goto err_put_domain;
 	}
 
-	mutex_lock(&rocket_priv->mm_lock);
-	ret = drm_mm_insert_node_generic(&rocket_priv->mm, &rkt_obj->mm,
+	mutex_lock(&rkt_obj->domain->mm_lock);
+	ret = drm_mm_insert_node_generic(&rkt_obj->domain->mm, &rkt_obj->mm,
 					 rkt_obj->size, PAGE_SIZE,
 					 0, 0);
-	mutex_unlock(&rocket_priv->mm_lock);
+	mutex_unlock(&rkt_obj->domain->mm_lock);
 	if (ret)
-		goto err;
+		goto err_put_domain;
 
-	ret = iommu_map_sgtable(rocket_priv->domain->domain,
+	ret = iommu_map_sgtable(rkt_obj->domain->domain,
 				rkt_obj->mm.start,
 				shmem_obj->sgt,
 				IOMMU_READ | IOMMU_WRITE);
@@ -118,13 +133,17 @@ int rocket_ioctl_create_bo(struct drm_device *dev, void *data, struct drm_file *
 	return 0;
 
 err_unmap:
-	iommu_unmap(rocket_priv->domain->domain,
+	iommu_unmap(rkt_obj->domain->domain,
 		    rkt_obj->mm.start, rkt_obj->size);
 
 err_remove_node:
-	mutex_lock(&rocket_priv->mm_lock);
+	mutex_lock(&rkt_obj->domain->mm_lock);
 	drm_mm_remove_node(&rkt_obj->mm);
-	mutex_unlock(&rocket_priv->mm_lock);
+	mutex_unlock(&rkt_obj->domain->mm_lock);
+
+err_put_domain:
+	rocket_iommu_domain_put(rkt_obj->domain);
+	rkt_obj->domain = NULL;
 
 err:
 	drm_gem_shmem_object_free(gem_obj);

rocket_job.c

@@ -613,7 +613,7 @@ void rocket_job_close(struct rocket_file_priv *rocket_priv)
 {
 	struct rocket_device *rdev = rocket_priv->rdev;
 	struct drm_sched_entity *entity = &rocket_priv->sched_entity;
-	unsigned int core;
+	unsigned int core, i;
 
 	for (core = 0; core < rdev->num_cores && core < 3; core++)
 		drm_sched_entity_destroy(&rocket_priv->core_entity[core]);
@@ -629,14 +629,23 @@ void rocket_job_close(struct rocket_file_priv *rocket_priv)
 	 * frees it. Skipping this frees the domain while it is still attached
 	 * (rk_iommu_domain_free WARN) and leaves core->attached_domain dangling ->
 	 * use-after-free / panic when the next client's job touches it.
+	 *
+	 * rkopnu multi-domain: a core's attached_domain can now be any one of
+	 * this fd's (up to 16) domains, not the single one it used to have --
+	 * check them all. A core can only ever have one domain attached at a
+	 * time, so at most one of these can match; break once found.
 	 */
 	for (core = 0; core < rdev->num_cores; core++) {
 		struct rocket_core *rcore = &rdev->cores[core];
 
 		scoped_guard(mutex, &rcore->job_lock) {
-			if (rcore->attached_domain == rocket_priv->domain) {
-				iommu_detach_group(NULL, rcore->iommu_group);
-				rcore->attached_domain = NULL;
+			for (i = 0; i < ROCKET_MAX_IOMMU_DOMAINS; i++) {
+				if (rocket_priv->domains[i] &&
+				    rcore->attached_domain == rocket_priv->domains[i]) {
+					iommu_detach_group(NULL, rcore->iommu_group);
+					rcore->attached_domain = NULL;
+					break;
+				}
 			}
 		}
 	}
@@ -694,7 +703,15 @@ static int rocket_ioctl_submit_job(struct drm_device *dev, struct drm_file *file
 
 	rjob->out_bo_count = job->out_bo_handle_count;
 
-	rjob->domain = rocket_iommu_domain_get(file_priv);
+	/* rkopnu multi-domain: native rocket uAPI has no domain-id field;
+	 * always domain 0 (see the matching comment in rocket_gem.c -- this
+	 * ioctl is unregistered/unreachable today, kept correct not deleted). */
+	rjob->domain = rocket_iommu_domain_get(file_priv, 0);
+	if (IS_ERR(rjob->domain)) {
+		ret = PTR_ERR(rjob->domain);
+		rjob->domain = NULL;
+		goto out_cleanup_job;
+	}
 
 	ret = rocket_job_push(rjob);
 	if (ret)
@@ -835,7 +852,27 @@ int rkopnu_ioctl_submit(struct drm_device *dev, void *data, struct drm_file *fil
 
 	rjob->in_bo_count = 0;
 	rjob->out_bo_count = 0;
-	rjob->domain = rocket_iommu_domain_get(file_priv);
+	/*
+	 * rkopnu multi-domain: this is the field librknnrt actually threads
+	 * through from rknn_matmul_info.iommu_domain_id at rknn_matmul_create()
+	 * time, for the *same* matmul context whose weight buffer was created
+	 * against that domain id via MEM_CREATE. Honoring it here is what
+	 * makes rocket_job_run()'s existing "reattach only on domain change"
+	 * logic (rocket_job.c) actually switch domains instead of always
+	 * seeing the same one.
+	 */
+	rjob->domain = rocket_iommu_domain_get(file_priv, (s32)args->iommu_domain_id);
+	if (IS_ERR(rjob->domain)) {
+		/* drm_sched_job_init() already succeeded above (this check runs
+		 * after it), so unlike the two out_put_job jumps elsewhere in
+		 * this function, this one must go through out_cleanup so
+		 * drm_sched_job_cleanup() actually runs -- straight to
+		 * out_put_job here would skip it and leak/imbalance whatever
+		 * job_init set up. */
+		ret = PTR_ERR(rjob->domain);
+		rjob->domain = NULL;
+		goto out_cleanup;
+	}
 
 	ret = rocket_job_push(rjob);
 	if (ret)