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

[Lizhi Hou]

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....@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