[PATCH 2/2] accel/rocket: keep core slots stable across unbind and rebind

[Igor Paunovic <royalnet026@gmail.com>]

The shared rocket_device tracks bound cores with a single counter and
uses it for three different jobs at once: find_core_for_dev() searches
[0, num_cores), rocket_probe() inserts the new core at index num_cores,
and rocket_remove() only decrements the counter without clearing the
slot.

This bookkeeping falls apart as soon as cores are unbound in any order
other than strict reverse bind order:

 - unbinding core 0 shrinks the search range, so the still-bound core
   at the highest index can no longer be found: its runtime PM
   callbacks start failing with -ENODEV and a later unbind of it is
   silently ignored, skipping rocket_core_fini() entirely;

 - a subsequent bind then reuses the index of that still-live core and
   overwrites its slot while its IRQ handler (dev_id points into
   cores[]) and its DRM scheduler are still active;

 - rocket_open() unconditionally uses cores[0].dev, which after an
   unbind of core 0 is a stale pointer to an unbound device.

Give the array a fixed capacity (max_cores, the DT core count already
used to size the allocation) and make .dev the slot-liveness marker:
probe takes the first free slot, remove clears .dev after
rocket_core_fini() and warns if the core cannot be found, lookups
iterate the full capacity, and rocket_open() and rocket_job_open() use
only live slots. num_cores keeps counting bound cores for the
last-core teardown check.

Fixes: ed98261b4168 ("accel/rocket: Add a new driver for Rockchip's NPU")
Signed-off-by: Igor Paunovic <royalnet026@gmail.com>
---
Unbinding a core that still has jobs in flight has further
pre-existing issues (scheduler and open-file lifetime) that are out of
scope for this bookkeeping fix.

Verified on RK3588 (Orange Pi 5 Plus): out-of-order unbind/rebind
sequences (including the previously corrupting unbind of core 0 with
cores 1 and 2 still bound, followed by rebind) now keep all three
cores findable and functional, with a MobileNetV1 inference run via
the Teflon TFLite delegate bit-identical to the stock driver
afterwards.

 drivers/accel/rocket/rocket_device.c |  2 ++
 drivers/accel/rocket/rocket_device.h |  3 +++
 drivers/accel/rocket/rocket_drv.c    | 32 ++++++++++++++++++++++++++++----
 drivers/accel/rocket/rocket_job.c    | 11 ++++++-----
 4 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/drivers/accel/rocket/rocket_device.c b/drivers/accel/rocket/rocket_device.c
index 46e6ee1..8303f05 100644
--- a/drivers/accel/rocket/rocket_device.c
+++ b/drivers/accel/rocket/rocket_device.c
@@ -35,6 +35,8 @@ struct rocket_device *rocket_device_init(struct platform_device *pdev,
 	if (!rdev->cores)
 		return ERR_PTR(-ENOMEM);
 
+	rdev->max_cores = num_cores;
+
 	dma_set_max_seg_size(dev, UINT_MAX);
 
 	err = dma_set_mask_and_coherent(dev, DMA_BIT_MASK(40));
diff --git a/drivers/accel/rocket/rocket_device.h b/drivers/accel/rocket/rocket_device.h
index ce662ab..7fb6a9d 100644
--- a/drivers/accel/rocket/rocket_device.h
+++ b/drivers/accel/rocket/rocket_device.h
@@ -18,6 +18,9 @@ struct rocket_device {
 	struct mutex sched_lock;
 
 	struct rocket_core *cores;
+	/* Slot capacity (DT core count); slots with a NULL .dev are free. */
+	unsigned int max_cores;
+	/* Number of currently bound cores. */
 	unsigned int num_cores;
 };
 
diff --git a/drivers/accel/rocket/rocket_drv.c b/drivers/accel/rocket/rocket_drv.c
index d29c5ee..04c1e47 100644
--- a/drivers/accel/rocket/rocket_drv.c
+++ b/drivers/accel/rocket/rocket_drv.c
@@ -69,11 +69,21 @@ rocket_iommu_domain_put(struct rocket_iommu_domain *domain)
 	kref_put(&domain->kref, rocket_iommu_domain_destroy);
 }
 
+static struct rocket_core *rocket_first_live_core(struct rocket_device *rdev)
+{
+	for (unsigned int core = 0; core < rdev->max_cores; core++)
+		if (rdev->cores[core].dev)
+			return &rdev->cores[core];
+
+	return NULL;
+}
+
 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;
+	struct rocket_core *core;
 	u64 start, end;
 	int ret;
 
@@ -86,8 +96,14 @@ rocket_open(struct drm_device *dev, struct drm_file *file)
 		goto err_put_mod;
 	}
 
+	core = rocket_first_live_core(rdev);
+	if (!core) {
+		ret = -ENODEV;
+		goto err_free;
+	}
+
 	rocket_priv->rdev = rdev;
-	rocket_priv->domain = rocket_iommu_domain_create(rdev->cores[0].dev);
+	rocket_priv->domain = rocket_iommu_domain_create(core->dev);
 	if (IS_ERR(rocket_priv->domain)) {
 		ret = PTR_ERR(rocket_priv->domain);
 		goto err_free;
@@ -179,10 +195,17 @@ static int rocket_probe(struct platform_device *pdev)
 		devres_close_group(&drm_dev->dev, rdev_group);
 	}
 
-	unsigned int core = rdev->num_cores;
+	unsigned int core;
 
 	dev_set_drvdata(&pdev->dev, rdev);
 
+	/* Take the first free slot: cores can unbind and rebind in any order. */
+	for (core = 0; core < rdev->max_cores; core++)
+		if (!rdev->cores[core].dev)
+			break;
+	if (WARN_ON(core == rdev->max_cores))
+		return -ENXIO;
+
 	rdev->cores[core].rdev = rdev;
 	rdev->cores[core].dev = &pdev->dev;
 	rdev->cores[core].index = core;
@@ -210,10 +233,11 @@ static void rocket_remove(struct platform_device *pdev)
 	struct device *dev = &pdev->dev;
 	int core = find_core_for_dev(dev);
 
-	if (core < 0)
+	if (WARN_ON(core < 0))
 		return;
 
 	rocket_core_fini(&rdev->cores[core]);
+	rdev->cores[core].dev = NULL;
 	rdev->num_cores--;
 
 	if (rdev->num_cores == 0) {
@@ -234,7 +258,7 @@ static int find_core_for_dev(struct device *dev)
 {
 	struct rocket_device *rdev = dev_get_drvdata(dev);
 
-	for (unsigned int core = 0; core < rdev->num_cores; core++) {
+	for (unsigned int core = 0; core < rdev->max_cores; core++) {
 		if (dev == rdev->cores[core].dev)
 			return core;
 	}
diff --git a/drivers/accel/rocket/rocket_job.c b/drivers/accel/rocket/rocket_job.c
index ac51bff..680f943 100644
--- a/drivers/accel/rocket/rocket_job.c
+++ b/drivers/accel/rocket/rocket_job.c
@@ -276,7 +276,7 @@ static struct rocket_core *sched_to_core(struct rocket_device *rdev,
 {
 	unsigned int core;
 
-	for (core = 0; core < rdev->num_cores; core++) {
+	for (core = 0; core < rdev->max_cores; core++) {
 		if (&rdev->cores[core].sched == sched)
 			return &rdev->cores[core];
 	}
@@ -498,16 +498,17 @@ int rocket_job_open(struct rocket_file_priv *rocket_priv)
 	struct rocket_device *rdev = rocket_priv->rdev;
 	struct drm_gpu_scheduler **scheds = kmalloc_objs(*scheds,
 							 rdev->num_cores);
-	unsigned int core;
+	unsigned int core, n = 0;
 	int ret;
 
-	for (core = 0; core < rdev->num_cores; core++)
-		scheds[core] = &rdev->cores[core].sched;
+	for (core = 0; core < rdev->max_cores; core++)
+		if (rdev->cores[core].dev)
+			scheds[n++] = &rdev->cores[core].sched;
 
 	ret = drm_sched_entity_init(&rocket_priv->sched_entity,
 				    DRM_SCHED_PRIORITY_NORMAL,
 				    scheds,
-				    rdev->num_cores, NULL);
+				    n, NULL);
 	if (WARN_ON(ret))
 		return ret;