rknpu: take the MMU force-reset out of the domain-switch path (supersedes #388)

drivers/iommu/rockchip-iommu.c

@@ -1323,6 +1323,125 @@ static void rk_iommu_detach_device(struct iommu_domain *domain,
 	}
 }
 
+/*
+ * Switch a master to a different domain by reprogramming the page table base in
+ * place, without going through the IOMMU core.
+ *
+ * The core path for this is iommu_detach_device() + iommu_attach_device(), and
+ * for rockchip-iommu the attach runs rk_iommu_enable(), i.e. a full
+ * rk_iommu_force_reset() of every MMU bank. With work in flight on other cores
+ * that reset intermittently fails its DTE_ADDR readback ("Error during raw
+ * reset. MMU_DTE_ADDR is not functioning"), the attach then fails, and the core
+ * leaves the master attached to no domain at all.
+ *
+ * The hardware does not need any of that to change page tables. On an already
+ * enabled MMU it is:
+ *
+ *	enable_stall -> write DTE_ADDR on every bank -> ZAP_CACHE -> disable_stall
+ *
+ * with paging left on throughout. No reset, no paging off/on.
+ *
+ * Because the core is not involved, iommu_get_domain_for_dev() does NOT track
+ * this switch; a caller using it must track the live domain itself. The caller
+ * is also responsible for ensuring no DMA is in flight across the swap.
+ */
+int rk_iommu_switch_domain(struct device *dev, struct iommu_domain *domain)
+{
+	struct rk_iommu *iommu = rk_iommu_from_dev(dev);
+	struct rk_iommu_domain *rk_domain;
+	unsigned long flags;
+	int ret, i;
+
+	if (!iommu || !domain)
+		return -ENODEV;
+
+	rk_domain = to_rk_domain(domain);
+
+	/* the third-party ops wrapper owns its own attach path */
+	if (rk_domain->opt_ops)
+		return -EOPNOTSUPP;
+
+	if (iommu->domain == domain)
+		return 0;
+
+	/*
+	 * Move this iommu between the two domains' iommus lists: rk_iommu_zap_iova()
+	 * and the TLB flush paths walk that list, so it has to name the live domain.
+	 */
+	if (iommu->domain) {
+		struct rk_iommu_domain *old = to_rk_domain(iommu->domain);
+
+		spin_lock_irqsave(&old->iommus_lock, flags);
+		list_del_init(&iommu->node);
+		spin_unlock_irqrestore(&old->iommus_lock, flags);
+	}
+	iommu->domain = domain;
+	spin_lock_irqsave(&rk_domain->iommus_lock, flags);
+	list_add_tail(&iommu->node, &rk_domain->iommus);
+	spin_unlock_irqrestore(&rk_domain->iommus_lock, flags);
+	rk_domain->shootdown_entire = iommu->shootdown_entire;
+
+	ret = pm_runtime_get_if_in_use(iommu->dev);
+	if (!ret || WARN_ON_ONCE(ret < 0)) {
+		/* not runtime-active: rk_iommu_resume() programs DTE_ADDR from iommu->domain */
+		return 0;
+	}
+
+	ret = clk_bulk_enable(iommu->num_clocks, iommu->clocks);
+	if (ret)
+		goto out_pm_put;
+
+	ret = rk_iommu_enable_stall(iommu);
+	if (ret)
+		goto out_disable_clocks;
+
+	for (i = 0; i < iommu->num_mmu; i++) {
+		rk_iommu_write(iommu->bases[i], RK_MMU_DTE_ADDR,
+			       rk_ops->mk_dtentries(rk_domain->dt_dma));
+		rk_iommu_base_command(iommu->bases[i], RK_MMU_CMD_ZAP_CACHE);
+	}
+
+	rk_iommu_disable_stall(iommu);
+
+out_disable_clocks:
+	clk_bulk_disable(iommu->num_clocks, iommu->clocks);
+out_pm_put:
+	pm_runtime_put(iommu->dev);
+	return ret;
+}
+EXPORT_SYMBOL(rk_iommu_switch_domain);
+
+/*
+ * Re-establish the MMU for the currently attached domain, unconditionally.
+ *
+ * A hardware soft reset wipes the MMU (DTE_ADDR, paging), so the page table has
+ * to be reprogrammed afterwards. rk_iommu_switch_domain() deliberately
+ * short-circuits when the domain has not changed, which is exactly wrong here:
+ * the domain is the same, the hardware is not.
+ *
+ * iommu->domain is the live domain, so reprogramming from it is correct by
+ * construction and does not depend on what the IOMMU core believes.
+ */
+int rk_iommu_reprogram(struct device *dev)
+{
+	struct rk_iommu *iommu = rk_iommu_from_dev(dev);
+	int ret;
+
+	if (!iommu || !iommu->domain)
+		return -ENODEV;
+	if (to_rk_domain(iommu->domain)->opt_ops)
+		return -EOPNOTSUPP;
+
+	ret = pm_runtime_get_if_in_use(iommu->dev);
+	if (!ret || WARN_ON_ONCE(ret < 0))
+		return 0;   /* not runtime-active: rk_iommu_resume() will program it */
+
+	ret = rk_iommu_enable(iommu);
+	pm_runtime_put(iommu->dev);
+	return ret;
+}
+EXPORT_SYMBOL(rk_iommu_reprogram);
+
 static int rk_iommu_attach_device(struct iommu_domain *domain,
 		struct device *dev)
 {

drivers/rknpu/include/rknpu_ioctl.h

@@ -122,6 +122,16 @@ enum e_rknpu_action {
 	RKNPU_GET_BW_TW = 11,
 	RKNPU_SET_BW_TW = 12,
 	RKNPU_ACT_CLR_TOTAL_RW_AMOUNT = 13,
+	/*
+	 * Make an IOMMU domain live, with no allocation. `value` is the domain id.
+	 *
+	 * A dma-buf import maps its sg during PRIME_FD_TO_HANDLE, into whichever
+	 * domain is live at that instant, because the ioctl carries no domain. The
+	 * target is named only by the later MEM_CREATE, which does not re-map an
+	 * already-imported handle. Userspace issues this first so the mapping lands
+	 * in the domain the buffer will actually be used in.
+	 */
+	RKNPU_ACT_SET_DOMAIN = 0x100,
 	RKNPU_GET_DT_WR_AMOUNT = 14,
 	RKNPU_GET_DT_RD_AMOUNT = 15,
 	RKNPU_GET_WT_RD_AMOUNT = 16,

drivers/rknpu/include/rknpu_iommu.h

@@ -53,6 +53,30 @@ int rknpu_iommu_domain_get_and_switch(struct rknpu_device *rknpu_dev,
 				      int domain_id);
 int rknpu_iommu_domain_put(struct rknpu_device *rknpu_dev);
 
+/*
+ * The domain the NPU is actually running in, from the driver's own state.
+ *
+ * iommu_get_domain_for_dev() returns what the IOMMU core last attached, which is
+ * only the same thing while domain switching goes through the core. Everything in
+ * this driver that must target the live domain uses this instead.
+ */
+struct iommu_domain *rknpu_iommu_live_domain(struct device *dev);
+
+/*
+ * Provided by drivers/iommu/rockchip-iommu.c. Switching with these keeps the IOMMU
+ * core out of the path, which is what avoids the rk_iommu_force_reset() that an
+ * iommu_attach_device() would perform on every domain change.
+ */
+extern int rk_iommu_switch_domain(struct device *dev, struct iommu_domain *domain);
+extern int rk_iommu_reprogram(struct device *dev);
+
+/*
+ * Scoped override of the core's default domain, for the dma-buf attachment map and
+ * its matching unmap only. Returns the previous value; the caller restores it.
+ */
+struct iommu_domain *rknpu_iommu_default_swap(struct device *dev,
+					      struct iommu_domain *dom);
+
 #if KERNEL_VERSION(5, 10, 0) < LINUX_VERSION_CODE
 int iommu_get_dma_cookie(struct iommu_domain *domain);
 #endif

drivers/rknpu/include/rknpu_job.h

@@ -32,6 +32,17 @@ struct rknpu_job {
 	struct work_struct cleanup_work;
 	bool irq_entry[RKNPU_MAX_CORES];
 	unsigned int flags;
+	/*
+	 * Does this job currently hold the IOMMU domain reference?
+	 *
+	 * The reference is acquired exactly once, in rknpu_job_commit(), but three
+	 * teardown paths release it -- the completion path, rknpu_job_abort() and
+	 * rknpu_job_timeout_clean() -- with nothing recording whether this particular
+	 * job still holds one. Its own word rather than a bit in ->flags because the
+	 * releases run from both interrupt and process context, so the test-and-clear
+	 * has to be atomic.
+	 */
+	unsigned long dom_held;
 	int ret;
 	struct rknpu_submit *args;
 	bool args_owner;

drivers/rknpu/rknpu_drv.c

@@ -466,6 +466,14 @@ static int rknpu_action(struct rknpu_device *rknpu_dev,
 	case RKNPU_SET_BW_TW:
 		ret = rknpu_set_bw_priority(rknpu_dev, 0, 0, args->value);
 		break;
+	case RKNPU_ACT_SET_DOMAIN: {
+		int _id = (int)args->value;
+
+		ret = rknpu_iommu_domain_get_and_switch(rknpu_dev, _id);
+		if (!ret)
+			rknpu_iommu_domain_put(rknpu_dev);   /* switch only, hold no reference */
+		break;
+	}
 	case RKNPU_ACT_CLR_TOTAL_RW_AMOUNT:
 		ret = rknpu_clear_rw_amount(rknpu_dev);
 		break;

drivers/rknpu/rknpu_gem.c

@@ -72,7 +72,7 @@ static int rknpu_gem_get_pages(struct rknpu_gem_object *rknpu_obj)
 			      rknpu_obj->size);
 		goto free_sgt;
 	}
-	iommu_flush_iotlb_all(iommu_get_domain_for_dev(drm->dev));
+	iommu_flush_iotlb_all(rknpu_iommu_live_domain(drm->dev));
 
 	if (rknpu_obj->flags & RKNPU_MEM_KERNEL_MAPPING) {
 		rknpu_obj->cookie = vmap(rknpu_obj->pages, rknpu_obj->num_pages,
@@ -501,7 +501,7 @@ static int rknpu_gem_alloc_buf_with_cache(struct rknpu_gem_object *rknpu_obj,
 	}
 
 	/* iova map to cache */
-	domain = iommu_get_domain_for_dev(rknpu_dev->dev);
+	domain = rknpu_iommu_live_domain(rknpu_dev->dev);
 	if (!domain) {
 		LOG_ERROR("failed to get iommu domain!");
 		return -EINVAL;
@@ -654,7 +654,7 @@ static void rknpu_gem_free_buf_with_cache(struct rknpu_gem_object *rknpu_obj,
 		return;
 	}
 
-	domain = iommu_get_domain_for_dev(rknpu_dev->dev);
+	domain = rknpu_iommu_live_domain(rknpu_dev->dev);
 	if (domain) {
 		iommu_unmap(domain, rknpu_obj->iova_start, cache_size);
 		if (rknpu_obj->size > 0)
@@ -850,7 +850,22 @@ void rknpu_gem_object_destroy(struct rknpu_gem_object *rknpu_obj)
 	 * once dmabuf's refcount becomes 0.
 	 */
 	if (obj->import_attach) {
-		drm_prime_gem_destroy(obj, rknpu_obj->sgt);
+		{
+			/*
+			 * dma_buf_unmap_attachment() unmaps through the DMA API
+			 * too, so it must see the same default domain the map saw.
+			 * The caller has already switched to this object's recorded
+			 * domain, so the live domain is the right one.
+			 */
+			struct iommu_domain *l = rknpu_dev->iommu_en ?
+				rknpu_iommu_live_domain(rknpu_dev->dev) : NULL;
+			struct iommu_domain *sv = l ?
+				rknpu_iommu_default_swap(rknpu_dev->dev, l) : NULL;
+
+			drm_prime_gem_destroy(obj, rknpu_obj->sgt);
+			if (sv)
+				rknpu_iommu_default_swap(rknpu_dev->dev, sv);
+		}
 		rknpu_gem_free_page(rknpu_obj->pages);
 	} else {
 		if (IS_ENABLED(CONFIG_ROCKCHIP_RKNPU_SRAM) &&
@@ -1398,7 +1413,30 @@ int rknpu_gem_mmap(struct file *filp, struct vm_area_struct *vma)
 struct drm_gem_object *rknpu_gem_prime_import(struct drm_device *dev,
 					      struct dma_buf *dma_buf)
 {
-	return drm_gem_prime_import_dev(dev, dma_buf, dev->dev);
+	struct rknpu_device *rknpu_dev = dev->dev_private;
+	struct drm_gem_object *obj;
+	struct iommu_domain *live, *saved = NULL;
+
+	if (!rknpu_dev->iommu_en)
+		return drm_gem_prime_import_dev(dev, dma_buf, dev->dev);
+
+	/*
+	 * DRM core maps the attachment through the DMA API here, which always
+	 * targets group->default_domain. Point that at the live domain for the
+	 * duration of the import so the sg is mapped where the NPU will run, then
+	 * restore it. domain_lock keeps a concurrent switch from moving the live
+	 * domain underneath the map.
+	 */
+	mutex_lock(&rknpu_dev->domain_lock);
+	live = rknpu_iommu_live_domain(rknpu_dev->dev);
+	if (live)
+		saved = rknpu_iommu_default_swap(rknpu_dev->dev, live);
+	obj = drm_gem_prime_import_dev(dev, dma_buf, dev->dev);
+	if (saved)
+		rknpu_iommu_default_swap(rknpu_dev->dev, saved);
+	mutex_unlock(&rknpu_dev->domain_lock);
+
+	return obj;
 }
 #endif