[PATCH V1] accel/amdxdna: Fix hardware context race in amdxdna_update_heap()

[Lizhi Hou <lizhi.hou@amd.com>]

amdxdna_update_heap() iterates over hardware contexts while holding
xdna->dev_lock. During the iteration, amdxdna_pm_resume_get_locked() may
temporarily release and reacquire the lock, allowing hardware contexts to
be modified concurrently.

Fix the race by calling amdxdna_pm_resume_get_locked() before
iterating over hardware contexts.

Fixes: dbc8fd7a03cb ("accel/amdxdna: Add expandable device heap support")
Signed-off-by: Lizhi Hou <lizhi.hou@amd.com>
---
 drivers/accel/amdxdna/aie2_ctx.c    |  6 ------
 drivers/accel/amdxdna/amdxdna_ctx.c | 24 +++++++++++++++---------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/drivers/accel/amdxdna/aie2_ctx.c b/drivers/accel/amdxdna/aie2_ctx.c
index 30ccb8d5e23d..101f324ee178 100644
--- a/drivers/accel/amdxdna/aie2_ctx.c
+++ b/drivers/accel/amdxdna/aie2_ctx.c
@@ -1222,10 +1222,6 @@ int aie2_hwctx_heap_expand(struct amdxdna_hwctx *hwctx,
 	u64 addr;
 	int ret;
 
-	ret = amdxdna_pm_resume_get_locked(xdna);
-	if (ret)
-		return ret;
-
 	addr = amdxdna_obj_dma_addr(heap);
 	ret = aie2_add_host_buf(xdna->dev_handle, hwctx->fw_ctx_id,
 				addr, heap->mem.size);
@@ -1234,7 +1230,5 @@ int aie2_hwctx_heap_expand(struct amdxdna_hwctx *hwctx,
 			 hwctx->name, heap->mem.size, ret);
 	}
 
-	amdxdna_pm_suspend_put(xdna);
-
 	return ret;
 }
diff --git a/drivers/accel/amdxdna/amdxdna_ctx.c b/drivers/accel/amdxdna/amdxdna_ctx.c
index 9ae19393e488..8f8df9d04ec5 100644
--- a/drivers/accel/amdxdna/amdxdna_ctx.c
+++ b/drivers/accel/amdxdna/amdxdna_ctx.c
@@ -472,9 +472,7 @@ static int amdxdna_hwctx_expand_heap(struct amdxdna_hwctx *hwctx)
 			break;
 		}
 
-		mutex_unlock(&client->mm_lock);
 		ret = xdna->dev_info->ops->hwctx_heap_expand(hwctx, heap);
-		mutex_lock(&client->mm_lock);
 		if (ret) {
 			amdxdna_gem_unpin(heap);
 			drm_gem_object_put(to_gobj(heap));
@@ -493,18 +491,26 @@ int amdxdna_update_heap(struct amdxdna_client *client, struct amdxdna_hwctx *hwc
 	unsigned long hwctx_id;
 	int ret;
 
-	guard(mutex)(&client->mm_lock);
+	ret = amdxdna_pm_resume_get_locked(client->xdna);
+	if (ret)
+		return ret;
 
-	if (hwctx)
-		return amdxdna_hwctx_expand_heap(hwctx);
+	mutex_lock(&client->mm_lock);
 
-	amdxdna_for_each_hwctx(client, hwctx_id, hwctx) {
+	if (hwctx) {
 		ret = amdxdna_hwctx_expand_heap(hwctx);
-		if (ret)
-			return ret;
+	} else {
+		amdxdna_for_each_hwctx(client, hwctx_id, hwctx) {
+			ret = amdxdna_hwctx_expand_heap(hwctx);
+			if (ret)
+				break;
+		}
 	}
+	mutex_unlock(&client->mm_lock);
 
-	return 0;
+	amdxdna_pm_suspend_put(client->xdna);
+
+	return ret;
 }
 
 static void
-- 
2.34.1