Files
rkopnu/rocket_gem.c
Markus Fritsche b831a2d931 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
2026-07-20 05:18:30 +02:00

211 lines
5.5 KiB
C

// SPDX-License-Identifier: GPL-2.0-only
/* Copyright 2024-2025 Tomeu Vizoso <tomeu@tomeuvizoso.net> */
#include <drm/drm_device.h>
#include <drm/drm_print.h>
#include <drm/drm_utils.h>
#include <drm/rocket_accel.h>
#include <linux/dma-mapping.h>
#include <linux/iommu.h>
#include "rocket_drv.h"
#include "rocket_gem.h"
static void rocket_gem_bo_free(struct drm_gem_object *obj)
{
struct rocket_gem_object *bo = to_rocket_bo(obj);
size_t unmapped;
drm_WARN_ON(obj->dev, refcount_read(&bo->base.pages_use_count) > 1);
unmapped = iommu_unmap(bo->domain->domain, bo->mm.start, bo->size);
drm_WARN_ON(obj->dev, unmapped != bo->size);
/* 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(&bo->domain->mm_lock);
rocket_iommu_domain_put(bo->domain);
bo->domain = NULL;
drm_gem_shmem_free(&bo->base);
}
static const struct drm_gem_object_funcs rocket_gem_funcs = {
.free = rocket_gem_bo_free,
.print_info = drm_gem_shmem_object_print_info,
.pin = drm_gem_shmem_object_pin,
.unpin = drm_gem_shmem_object_unpin,
.get_sg_table = drm_gem_shmem_object_get_sg_table,
.vmap = drm_gem_shmem_object_vmap,
.vunmap = drm_gem_shmem_object_vunmap,
.mmap = drm_gem_shmem_object_mmap,
.vm_ops = &drm_gem_shmem_vm_ops,
};
struct drm_gem_object *rocket_gem_create_object(struct drm_device *dev, size_t size)
{
struct rocket_gem_object *obj;
obj = kzalloc_obj(*obj);
if (!obj)
return ERR_PTR(-ENOMEM);
obj->base.base.funcs = &rocket_gem_funcs;
return &obj->base.base;
}
int rocket_ioctl_create_bo(struct drm_device *dev, void *data, struct drm_file *file)
{
struct rocket_file_priv *rocket_priv = file->driver_priv;
struct drm_rocket_create_bo *args = data;
struct drm_gem_shmem_object *shmem_obj;
struct rocket_gem_object *rkt_obj;
struct drm_gem_object *gem_obj;
struct sg_table *sgt;
int ret;
shmem_obj = drm_gem_shmem_create(dev, args->size);
if (IS_ERR(shmem_obj))
return PTR_ERR(shmem_obj);
gem_obj = &shmem_obj->base;
rkt_obj = to_rocket_bo(gem_obj);
rkt_obj->driver_priv = 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_put_domain;
}
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(&rkt_obj->domain->mm_lock);
if (ret)
goto err_put_domain;
ret = iommu_map_sgtable(rkt_obj->domain->domain,
rkt_obj->mm.start,
shmem_obj->sgt,
IOMMU_READ | IOMMU_WRITE);
if (ret < 0 || ret < args->size) {
drm_err(dev, "failed to map buffer: size=%d request_size=%u\n",
ret, args->size);
ret = -ENOMEM;
goto err_remove_node;
}
/* iommu_map_sgtable might have aligned the size */
rkt_obj->size = ret;
args->offset = drm_vma_node_offset_addr(&gem_obj->vma_node);
args->dma_address = rkt_obj->mm.start;
ret = drm_gem_handle_create(file, gem_obj, &args->handle);
if (ret)
goto err_unmap;
drm_gem_object_put(gem_obj);
return 0;
err_unmap:
iommu_unmap(rkt_obj->domain->domain,
rkt_obj->mm.start, rkt_obj->size);
err_remove_node:
mutex_lock(&rkt_obj->domain->mm_lock);
drm_mm_remove_node(&rkt_obj->mm);
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);
return ret;
}
int rocket_ioctl_prep_bo(struct drm_device *dev, void *data, struct drm_file *file)
{
struct drm_rocket_prep_bo *args = data;
unsigned long timeout = drm_timeout_abs_to_jiffies(args->timeout_ns);
struct drm_gem_object *gem_obj;
struct drm_gem_shmem_object *shmem_obj;
long ret = 0;
if (args->reserved != 0) {
drm_dbg(dev, "Reserved field in drm_rocket_prep_bo struct should be 0.\n");
return -EINVAL;
}
gem_obj = drm_gem_object_lookup(file, args->handle);
if (!gem_obj)
return -ENOENT;
ret = dma_resv_wait_timeout(gem_obj->resv, DMA_RESV_USAGE_WRITE, true, timeout);
if (!ret)
ret = timeout ? -ETIMEDOUT : -EBUSY;
else if (ret > 0)
ret = 0;
shmem_obj = &to_rocket_bo(gem_obj)->base;
dma_sync_sgtable_for_cpu(dev->dev, shmem_obj->sgt, DMA_BIDIRECTIONAL);
drm_gem_object_put(gem_obj);
return ret;
}
int rocket_ioctl_fini_bo(struct drm_device *dev, void *data, struct drm_file *file)
{
struct drm_rocket_fini_bo *args = data;
struct drm_gem_shmem_object *shmem_obj;
struct rocket_gem_object *rkt_obj;
struct drm_gem_object *gem_obj;
if (args->reserved != 0) {
drm_dbg(dev, "Reserved field in drm_rocket_fini_bo struct should be 0.\n");
return -EINVAL;
}
gem_obj = drm_gem_object_lookup(file, args->handle);
if (!gem_obj)
return -ENOENT;
rkt_obj = to_rocket_bo(gem_obj);
shmem_obj = &rkt_obj->base;
dma_sync_sgtable_for_device(dev->dev, shmem_obj->sgt, DMA_BIDIRECTIONAL);
drm_gem_object_put(gem_obj);
return 0;
}