Claude review: accel/amdxdna: Add expandable device heap support

[Claude Code Review Bot <claude-review@example.com>]

Patch Review

**Locking: `amdxdna_hwctx_expand_heap` drops and re-takes `mm_lock` mid-iteration**

In `amdxdna_ctx.c`, `amdxdna_hwctx_expand_heap()` iterates the `dev_heap_xa` under `mm_lock`, but drops the lock to call `hwctx_heap_expand`:

```c
xa_for_each_range(&client->dev_heap_xa, heap_id, heap,
                  nid, client->dev_heap_nid) {
    drm_gem_object_get(to_gobj(heap));
    ret = amdxdna_gem_pin(heap);
    ...
    mutex_unlock(&client->mm_lock);
    ret = xdna->dev_info->ops->hwctx_heap_expand(hwctx, heap);
    mutex_lock(&client->mm_lock);
    ...
```

Dropping `mm_lock` while iterating the xarray is dangerous. `client->dev_heap_nid` could change while the lock is dropped (another thread could add a new heap chunk). The `xa_for_each_range` macro caches its iteration state, but the upper bound `client->dev_heap_nid` is evaluated once. More importantly, if another concurrent call to `amdxdna_update_heap(client, NULL)` runs, you could get concurrent iteration + expansion on the same hwctx. The `dev_lock` is held in the `amdxdna_drm_create_bo_ioctl` path but not in the `aie2_hwctx_init` path, so this is a real concern.

**Off-by-one in `xa_for_each_range` upper bound**

In `aie2_hwctx_restart`:
```c
xa_for_each_range(&hwctx->client->dev_heap_xa, heap_id, heap, 1,
                  hwctx->last_attached_heap) {
```

And in `amdxdna_hwctx_release_expanded_heap`:
```c
xa_for_each_range(&client->dev_heap_xa, heap_id, heap, 1,
                  hwctx->last_attached_heap) {
```

These iterate from index 1 to `last_attached_heap` inclusive. But in `amdxdna_hwctx_expand_heap`, `last_attached_heap` is set to `heap_id` on success:
```c
hwctx->last_attached_heap = heap_id;
```

This looks correct for the restart/release paths — they replay from 1 (skipping the initial heap at 0) to the last expanded heap inclusive. However, `last_attached_heap` is initialized to 0 (by kzalloc), and the conditional in `amdxdna_hwctx_release_expanded_heap` checks `if (hwctx->last_attached_heap)` — this works but is fragile. If `last_attached_heap` were ever set to 0 for a different reason, the release path would skip cleanup.

**`amdxdna_update_heap` called from `amdxdna_drm_create_bo_ioctl` with wrong locking**

```c
case AMDXDNA_BO_DEV:
    abo = amdxdna_drm_create_dev_bo(dev, args, filp);
    if (!IS_ERR(abo)) {
        mutex_lock(&xdna->dev_lock);
        ret = amdxdna_update_heap(client, NULL);
        mutex_unlock(&xdna->dev_lock);
        if (ret)
            goto put_obj;
    }
```

`amdxdna_update_heap(client, NULL)` iterates all hwctxs with `amdxdna_for_each_hwctx`, which is `xa_for_each(&client->hwctx_xa, ...)`. In other callsites this xarray is iterated under SRCU read lock (`client->hwctx_srcu`), but here it's only under `dev_lock`. While `dev_lock` may prevent concurrent hwctx creation/destruction, the existing code consistently uses SRCU for hwctx xarray iteration. This is at minimum inconsistent and may be incorrect if other paths remove hwctxs without holding `dev_lock`.

**Missing pin in sync BO path for dev BOs**

The original `amdxdna_drm_sync_bo_ioctl` pinned the BO before flushing. In the new code, the AMDXDNA_BO_DEV path skips pinning entirely:

```c
if (abo->type == AMDXDNA_BO_DEV) {
    ...
    xa_for_each_range(&client->dev_heap_xa, heap_id, heap, ...) {
        ...
        ret = amdxdna_flush_bo(heap, start - heap_start, end - start);
```

The heap chunks are flushed without pinning them first. The original code pinned the dev BO (which redirected to pinning the heap). Without pinning, `amdxdna_gem_vmap()` or `abo->base.pages` may not be valid. The flush helper `amdxdna_flush_bo` tries `amdxdna_gem_vmap()` and `abo->base.pages`, but these may be NULL if the heap isn't pinned.

**`amdxdna_flush_bo` silently clamps size but logs no warning**

```c
static int amdxdna_flush_bo(struct amdxdna_gem_obj *abo, u64 offset, u64 size)
{
    ...
    size = min(abo->mem.size, end) - offset;
    if (is_import_bo(abo))
        drm_clflush_sg(abo->base.sgt);
    else if (amdxdna_gem_vmap(abo))
        drm_clflush_virt_range(amdxdna_gem_vmap(abo) + offset, size);
    else if (abo->base.pages)
        drm_clflush_pages(abo->base.pages, abo->mem.size >> PAGE_SHIFT);
```

When the `pages` path is taken, it flushes `abo->mem.size >> PAGE_SHIFT` pages — the entire BO, ignoring `offset` and `size`. This was the same in the original code and is technically a pre-existing issue, but it's now been refactored into a helper that receives offset/size parameters suggesting it should respect them.

Also, the original code had `drm_WARN` for the "can not flush" case. The new `amdxdna_flush_bo` returns `-EINVAL` for that case, which changes the behavior from a warning to an ioctl error. The caller in the non-dev-BO path then issues `drm_WARN` on the error, which is fine. But the dev-BO path doesn't issue a `drm_WARN`, just logs an error — inconsistent.

**`aie2_send_host_buf_msgs` reuses `DECLARE_AIE_MSG` stack variables across loop iterations**

```c
static int aie2_send_host_buf_msgs(...)
{
    DECLARE_AIE_MSG(map_host_buffer, MSG_OP_MAP_HOST_BUFFER);
    ...
    msg.opcode = initial_opcode;
    do {
        req.context_id = context_id;
        req.buf_addr = addr;
        req.buf_size = chunk_size;
        ret = aie_send_mgmt_msg_wait(&ndev->aie, &msg);
        ...
        msg.opcode = MSG_OP_ADD_HOST_BUFFER;
    } while (size);
```

The `DECLARE_AIE_MSG` macro declares `req` and `resp` as stack-initialized structures. This should be fine since `req` fields are fully re-assigned each iteration. However, the `msg` struct's `opcode` is overridden to `initial_opcode` before the loop, and set to `MSG_OP_ADD_HOST_BUFFER` after the first iteration. The `resp.status` is initialized to -1 by the macro, but after the first `aie_send_mgmt_msg_wait`, it may contain a different value — this should be re-initialized for subsequent iterations if the message handling checks it. Worth verifying.

**Error log removed from `amdxdna_drm_create_dev_bo`**

```c
ret = amdxdna_gem_heap_alloc(abo);
if (ret) {
-    XDNA_ERR(xdna, "Failed to alloc dev bo memory, ret %d", ret);
     amdxdna_gem_destroy_obj(abo);
     return ERR_PTR(ret);
}
```

The error log was removed. Since `-EAGAIN` is now an expected return (signaling the need for heap expansion), this makes sense — but it means real allocation failures are also silent. Consider logging at `XDNA_DBG` level or only suppressing for `-EAGAIN`.

**`amdxdna_gem_uva` changed from function to inline but semantics changed**

The old `amdxdna_gem_uva` computed the UVA for `AMDXDNA_BO_DEV` objects dynamically from the heap's UVA + offset. The new version is a simple inline returning `abo->mem.uva`:

```c
static inline u64 amdxdna_gem_uva(struct amdxdna_gem_obj *abo)
{
    return abo->mem.uva;
}
```

This means `mem.uva` must be correctly set for DEV BOs during `amdxdna_init_dev_bo`. Looking at that function:
```c
dev_bo->mem.uva = dev_bo->mm_node.start - heap_addr + exp_heap_uva;
```

This is set once at allocation time. If the heap's UVA changes after the dev BO is created (e.g., remap), the dev BO's UVA would be stale. In the old code, it was always computed live. This may be acceptable if heap UVA is immutable after mmap, but it's a subtle semantic change worth noting.

**`amdxdna_init_dev_bo` has a subtle underflow risk**

```c
for (heap_id = 0; heap_id < client->dev_heap_nid; heap_id++) {
    heap = xa_load(&client->dev_heap_xa, heap_id);
    ...
    heap_addr = amdxdna_gem_dev_addr(heap);
    if (heap_addr > dev_bo->mm_node.start)
        break;
}

heap_id--;
heap = xa_load(&client->dev_heap_xa, heap_id);
```

If `heap_id` is 0 and the first heap's `dev_addr` is already greater than `mm_node.start`, `heap_id` will be 0 when it's decremented, wrapping to `UINT_MAX` (it's `u32`). This would be a bug. This scenario shouldn't happen in practice since the first heap starts at `dev_mem_base` and the mm allocator also starts at `dev_mem_base`, but it's not defended against.

**`amdxdna_gem_pin_nolock` for DEV BOs doesn't hold `abo->lock`**

The new `amdxdna_gem_pin_nolock` for `AMDXDNA_BO_DEV` type iterates heap chunks and calls `amdxdna_bo_pin` on each. But `amdxdna_gem_pin` (the caller) takes `abo->lock` — the lock of the dev BO, not the heap BO. The underlying `drm_gem_shmem_pin` on the heap BO is called without the heap BO's lock. The old code redirected to `client->dev_heap` which was always the same object, and `amdxdna_gem_pin` would take its lock. Now the lock discipline has changed — the dev BO's lock is held but the individual heap BOs' locks are not.

**npu1 gets `dev_heap_max_size = AIE2_DEVM_SIZE` (64M) while others get `AIE2_DEVM_MAX_SIZE` (512M)**

This is intentional — npu1 is an older device that doesn't support heap expansion. But the code path that creates additional heap BOs doesn't check whether the firmware supports `MSG_OP_ADD_HOST_BUFFER` until `aie2_add_host_buf` is called. On npu1, `dev_heap_max_size` is 64M, same as `dev_mem_size`, so only one 64M chunk fits. This means expansion is effectively blocked by size, not by feature check. This works but is fragile — if someone changed the sizes without understanding this coupling, npu1 could try unsupported firmware commands.

**Summary**: The overall design is sound, but the locking around `amdxdna_hwctx_expand_heap` dropping `mm_lock` mid-xarray iteration needs careful scrutiny. The missing pin in the sync-BO path for dev BOs is likely a bug. The `heap_id--` underflow in `amdxdna_init_dev_bo` should be defended. The inconsistent lock discipline between `amdxdna_gem_pin_nolock` (holds dev BO lock but not heap BO locks) and the old code (which held the heap's lock) is a correctness concern.

---
Generated by Claude Code Patch Reviewer