diff --git a/rkopnu_ioctl.c b/rkopnu_ioctl.c index f5b3faa..90e4e6c 100644 --- a/rkopnu_ioctl.c +++ b/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; diff --git a/rocket_drv.c b/rocket_drv.c index a6c87cd..2f58d62 100644 --- a/rocket_drv.c +++ b/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); } diff --git a/rocket_drv.h b/rocket_drv.h index 0f12073..90ed573 100644 --- a/rocket_drv.h +++ b/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 diff --git a/rocket_gem.c b/rocket_gem.c index a5fffa5..351267e 100644 --- a/rocket_gem.c +++ b/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); diff --git a/rocket_job.c b/rocket_job.c index 23aeb5c..eb7f691 100644 --- a/rocket_job.c +++ b/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)