rknpu: rknpu_job_abort() can zero iommu_domain_refcount while other cores have jobs in flight, allowing a domain switch under live work

[mafischer]

Summary

In drivers/rknpu, rknpu_job_abort() releases the IOMMU domain reference unconditionally. When
several cores have jobs in flight and one of them is aborted, this can drive
rknpu_dev->iommu_domain_refcount to zero while the other cores are still executing.

rknpu_iommu_domain_get_and_switch() decides it is safe to switch domains using only:

if (atomic_read(&rknpu_dev->iommu_domain_refcount) == 0) {
        ret = rknpu_iommu_switch_domain(rknpu_dev, domain_id);

so once the count reaches zero prematurely, a domain switch is performed underneath live work.
The result is a cascade that only a reboot clears:

RKNPU: mismatch domain get from iommu_get_domain_for_dev
RKNPU: failed to switch iommu domain, id: 1, ret: -22
RKNPU: rknpu_gem_get_pages: dma map 2097152 fail
  -> userspace MEM_CREATE returns -EINVAL for every subsequent allocation

This only affects users of multiple IOMMU domains (rknpu_mem_create's iommu_domain_id), so a
single-domain workload will not see it.

Evidence

We instrumented rknpu_iommu_domain_put() with __builtin_return_address(0) and reported two
conditions separately — an underflow (count below zero) and a premature zero (count reaches exactly
0 while subcore_datas[i].job != NULL for some core). Over one run:

PREMATURE ZERO (17x)  rknpu_job_abort+0x40/0x250
                (1x)  rknpu_gem_sync_ioctl+0xfc/0x234
UNDERFLOW       (1x)  rknpu_job_timeout_clean+0x11c/0x178

The premature zeros produce no warning of any kind — the count never goes negative, so a
WARN_ON-style underflow check does not catch them. Sampling across six consecutive runs on one boot
shows it accumulating monotonically (0 -> 29 -> 57 ...) until allocation fails outright.

Note the counter is a bare atomic_dec():

int rknpu_iommu_domain_put(struct rknpu_device *rknpu_dev)
{
	atomic_dec(&rknpu_dev->iommu_domain_refcount);
	return 0;
}

so an unbalanced put is also unbounded downward; we added a clamp locally, which prevents the count
going negative but does not prevent the damaging premature-zero window.

Reproducing

Any workload that (a) uses more than one IOMMU domain and (b) aborts jobs while other cores are busy.
Anything that raises the abort rate makes it appear much faster — in our case a watchdog that fails
stalled submits early took it from occasional to reproducible within two runs.

Suggested fix

Gating the switch on the refcount alone is not sufficient, because the count does not describe whether
the device is busy. Either:

1. gate on refcount == 0 and no core having a job (subcore_datas[i].job == NULL for all i), or
2. make the domain reference per-job — take it when the job acquires the domain and release it exactly
   once on the job's terminal transition — so abort cannot release a reference the job does not hold.

(1) is the smaller change and closes the window directly.

Context

Found while debugging a separate NPU stall whose signature matches
rockchip-linux/rknn-toolkit2#320 — details of that investigation are in the
comment we left there. We are an independent clean-room project
(ork-driver) driving the NPU through the in-tree rknpu DRM
uABI without librknnrt, so the code path above is entirely vendor code. Not affiliated with Rockchip.

Kernel: 6.1.115-vendor-rk35xx, RK3588 (Radxa ROCK 5B).