[PATCH v2] drm/amdgpu: Fix NPA-REVOKE racing an in-flight UALink import

[PATCH v2] drm/amdgpu: Fix NPA-REVOKE racing an in-flight UALink import

[David Carlier]

The exporter records an importer when it answers NPA-REQ, so it can send
NPA-REVOKE as soon as the BO is freed, before the importer has finished
building the dma-buf for that handle. The revoke handler assumes a fully
imported node: it dereferences imp_xa_node->dmabuf, which is still NULL
until the import completes, and drops the xarray reference the importing
thread still relies on. The importer then links the node and marks it
READY regardless, so the node can be freed while still on the per-remote
list.

Only tear down a node that is READY. Otherwise mark it for teardown and
send NPA-RELEASE, as nothing has been handed to user-space yet, and wake
the importer if it is still waiting for NPA-RSP so that it fails right
away. A node already in teardown belongs to whoever moved it there, so a
duplicate NPA-REVOKE no longer touches it either. The importer checks for
teardown under the xarray lock before linking the node and marking it
READY, and unwinds otherwise.

Fixes: 7cc82cd90d35 ("drm/amdgpu: Implement mechanism to revoke exported memory")
Assisted-by: LLM
Signed-off-by: David Carlier <devnexen@gmail.com>
---
Changes in v2:
- Tear down only READY nodes, so a duplicate NPA-REVOKE for a node already
  in teardown neither dereferences a NULL dmabuf nor drops the node
  reference twice (Sashiko).
- Complete npa_done when a revoke arrives before NPA-RSP, so the importer
  fails right away instead of timing out into a connection reset (Sashiko).
- Use the current Assisted-by format.

Found by code analysis and compile-tested with W=1. Not tested on hardware,
as it needs two UALink-connected accelerators in a vPod.

v1: https://lore.kernel.org/all/20260926171625.288519-1-devnexen@gmail.com/

 drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c | 34 +++++++++++++++++++---
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c
index 8411ea17172f..cb35026e6eba 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c
@@ -3265,6 +3265,7 @@ static void amdgpu_ualink_process_npa_revoke_msg(struct amdgpu_device *adev,
 {
 	struct amdgpu_ualink_imp_xa_node *imp_xa_node;
 	struct amdgpu_bo *bo;
+	u32 node_state;
 	int r = 0;
 
 	/* Remove the entry from the Xarray. */
@@ -3288,7 +3289,23 @@ static void amdgpu_ualink_process_npa_revoke_msg(struct amdgpu_device *adev,
 		return;
 	}
 
+	node_state = READ_ONCE(imp_xa_node->node_state);
 	WRITE_ONCE(imp_xa_node->node_state, AMDGPU_UALINK_NODE_TEARDOWN);
+
+	/* Only a READY node is torn down here. If the import is still in
+	 * flight, the dmabuf may not exist yet and nothing has been handed
+	 * to user-space: leave the node to the importing thread, which sees
+	 * the teardown state and unwinds, and wake it up if it is still
+	 * waiting for NPA-RSP. A node already in teardown is owned by
+	 * whoever moved it there, e.g. an earlier NPA-REVOKE.
+	 */
+	if (node_state != AMDGPU_UALINK_NODE_READY) {
+		if (node_state == AMDGPU_UALINK_NODE_NOT_READY)
+			complete(&imp_xa_node->npa_done);
+		xa_unlock(&adev->ualink.imp_xa);
+		goto send_release;
+	}
+
 	list_del_init(&imp_xa_node->list);
 	xa_unlock(&adev->ualink.imp_xa);
 
@@ -3299,6 +3316,7 @@ static void amdgpu_ualink_process_npa_revoke_msg(struct amdgpu_device *adev,
 	/* Drop the refcount for the node */
 	amdgpu_ualink_imp_xa_entry_put(imp_xa_node);
 
+send_release:
 	r = amdgpu_ualink_send_npa_release_msg(adev, remote_acc_id, handle);
 	if (r)
 		dev_err(adev->dev,
@@ -3760,9 +3778,20 @@ static int amdgpu_ualink_do_import_handle(struct amdgpu_device *adev,
 		return r;
 	}
 
-	/* Add this node to the imported handles list for the remote GPU */
+	/* Add this node to the imported handles list for the remote GPU,
+	 * unless the exporter revoked the handle while the import was in
+	 * flight. The dmabuf is released with the last node reference.
+	 */
 	xa_lock(&adev->ualink.imp_xa);
+	if (READ_ONCE(imp_xa_node->node_state) == AMDGPU_UALINK_NODE_TEARDOWN) {
+		xa_unlock(&adev->ualink.imp_xa);
+		dev_warn(adev->dev,
+			 "IMPORT: handle:%llx:%llx revoked during import\n",
+			 handle.handle_hi, handle.handle_lo);
+		return -EINVAL;
+	}
 	list_add(&imp_xa_node->list, &adev->ualink.imp_handles_list[remote_acc_id]);
+	WRITE_ONCE(imp_xa_node->node_state, AMDGPU_UALINK_NODE_READY);
 	xa_unlock(&adev->ualink.imp_xa);
 
 	return 0;
@@ -3938,9 +3967,6 @@ int amdgpu_ualink_import_handle(struct drm_device *dev,
 					"IMPORT: XA import failed for handle:%llx:%llx\n",
 					handle.handle_hi, handle.handle_lo);
 			goto cleanup;
-		} else {
-			WRITE_ONCE(imp_xa_node->node_state,
-				   AMDGPU_UALINK_NODE_READY);
 		}
 	}
 
-- 
2.55.0

Re: [PATCH v2] drm/amdgpu: Fix NPA-REVOKE racing an in-flight UALink import

[David CARLIER]

Hi Mukul,

On Tue, 29 Sept 2026 at 20:05, Mukul Joshi <mukul.joshi@amd.com> wrote:
>
> Hi David,
>
> Thanks for the patch. Yes the race is real, however, the patch needs some updations.
>
> More below.
>
>
> On 9/26/2026 1:44 PM, David Carlier wrote:
>
> [You don't often get email from devnexen@gmail.com. Learn why this is important at https://aka.ms/LearnAboutSenderIdentification ]
>
> The exporter records an importer when it answers NPA-REQ, so it can send
> NPA-REVOKE as soon as the BO is freed, before the importer has finished
> building the dma-buf for that handle. The revoke handler assumes a fully
> imported node: it dereferences imp_xa_node->dmabuf, which is still NULL
> until the import completes, and drops the xarray reference the importing
> thread still relies on. The importer then links the node and marks it
> READY regardless, so the node can be freed while still on the per-remote
> list.
>
> Only tear down a node that is READY. Otherwise mark it for teardown and
> send NPA-RELEASE, as nothing has been handed to user-space yet, and wake
> the importer if it is still waiting for NPA-RSP so that it fails right
> away. A node already in teardown belongs to whoever moved it there, so a
> duplicate NPA-REVOKE no longer touches it either. The importer checks for
> teardown under the xarray lock before linking the node and marking it
> READY, and unwinds otherwise.
>
> I think NPA-REVOKE cannot land before a NPA-RSP so we will not hit the condition where
> we have to wake up the importer.
> NPA-REVOKE is sent only when the exporter's XA entry's ref count goes down to 0.
> That will happen at the end of process_npa_req(), by that time, the NPA-RSP is already sent.
>
> Fixes: 7cc82cd90d35 ("drm/amdgpu: Implement mechanism to revoke exported memory")
> Assisted-by: LLM
> Signed-off-by: David Carlier <devnexen@gmail.com>
> ---
> Changes in v2:
> - Tear down only READY nodes, so a duplicate NPA-REVOKE for a node already
>   in teardown neither dereferences a NULL dmabuf nor drops the node
>   reference twice (Sashiko).
> - Complete npa_done when a revoke arrives before NPA-RSP, so the importer
>   fails right away instead of timing out into a connection reset (Sashiko).
> - Use the current Assisted-by format.
>
> Found by code analysis and compile-tested with W=1. Not tested on hardware,
> as it needs two UALink-connected accelerators in a vPod.
>
> v1: https://lore.kernel.org/all/20260926171625.288519-1-devnexen@gmail.com/
>
>  drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c | 34 +++++++++++++++++++---
>  1 file changed, 30 insertions(+), 4 deletions(-)
>
> diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c
> index 8411ea17172f..cb35026e6eba 100644
> --- a/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c
> +++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_ualink.c
> @@ -3265,6 +3265,7 @@ static void amdgpu_ualink_process_npa_revoke_msg(struct amdgpu_device *adev,
>  {
>         struct amdgpu_ualink_imp_xa_node *imp_xa_node;
>         struct amdgpu_bo *bo;
> +       u32 node_state;
>         int r = 0;
>
>         /* Remove the entry from the Xarray. */
> @@ -3288,7 +3289,23 @@ static void amdgpu_ualink_process_npa_revoke_msg(struct amdgpu_device *adev,
>                 return;
>         }
>
> +       node_state = READ_ONCE(imp_xa_node->node_state);
>         WRITE_ONCE(imp_xa_node->node_state, AMDGPU_UALINK_NODE_TEARDOWN);
> +
> +       /* Only a READY node is torn down here. If the import is still in
> +        * flight, the dmabuf may not exist yet and nothing has been handed
> +        * to user-space: leave the node to the importing thread, which sees
> +        * the teardown state and unwinds, and wake it up if it is still
> +        * waiting for NPA-RSP. A node already in teardown is owned by
> +        * whoever moved it there, e.g. an earlier NPA-REVOKE.
> +        */
> +       if (node_state != AMDGPU_UALINK_NODE_READY) {
> +               if (node_state == AMDGPU_UALINK_NODE_NOT_READY)
> +                       complete(&imp_xa_node->npa_done);
> +               xa_unlock(&adev->ualink.imp_xa);
> +               goto send_release;
> +       }
>
> As mentioned above, NPA-REVOKE cannot land before the NPA-RSP is sent by the exporter.
> So, if the node_state is NOT_READY that means its a stale NPA_REVOKE and we should just ignore
> that NPA-REVOKE. Having said that, we should definitely do the teardown when the node_state is READY.
> We should also handle NPA-REVOKE while node_state is in PENDING state.
> So, maybe we can refactor this code to something like this:
>
> switch (READ_ONCE(imp_xa_node->node_state)) {
>     case AMDGPU_UALINK_NODE_READY:
>         /* existing teardown */
>         break;
>     case AMDGPU_UALINK_NODE_PENDING:
>         WRITE_ONCE(imp_xa_node->node_state, AMDGPU_UALINK_NODE_TEARDOWN);
>         xa_unlock(&adev->ualink.imp_xa);
>         break;
>     default:
>         xa_unlock(&adev->ualink.imp_xa);
>         return;
>     }
>
> +
>         list_del_init(&imp_xa_node->list);
>         xa_unlock(&adev->ualink.imp_xa);
>
> @@ -3299,6 +3316,7 @@ static void amdgpu_ualink_process_npa_revoke_msg(struct amdgpu_device *adev,
>         /* Drop the refcount for the node */
>         amdgpu_ualink_imp_xa_entry_put(imp_xa_node);
>
> +send_release:
>         r = amdgpu_ualink_send_npa_release_msg(adev, remote_acc_id, handle);
>         if (r)
>                 dev_err(adev->dev,
> @@ -3760,9 +3778,20 @@ static int amdgpu_ualink_do_import_handle(struct amdgpu_device *adev,
>                 return r;
>         }
>
> -       /* Add this node to the imported handles list for the remote GPU */
> +       /* Add this node to the imported handles list for the remote GPU,
> +        * unless the exporter revoked the handle while the import was in
> +        * flight. The dmabuf is released with the last node reference.
> +        */
>         xa_lock(&adev->ualink.imp_xa);
> +       if (READ_ONCE(imp_xa_node->node_state) == AMDGPU_UALINK_NODE_TEARDOWN) {
> +               xa_unlock(&adev->ualink.imp_xa);
> +               dev_warn(adev->dev,
> +                        "IMPORT: handle:%llx:%llx revoked during import\n",
> +                        handle.handle_hi, handle.handle_lo);
> +               return -EINVAL;
> +       }
>         list_add(&imp_xa_node->list, &adev->ualink.imp_handles_list[remote_acc_id]);
> +       WRITE_ONCE(imp_xa_node->node_state, AMDGPU_UALINK_NODE_READY);
>         xa_unlock(&adev->ualink.imp_xa);
>
> The changes here makes sense. One nit-pick is to change from dev_warn to dev_dbg().
>
>         return 0;
> @@ -3938,9 +3967,6 @@ int amdgpu_ualink_import_handle(struct drm_device *dev,
>                                         "IMPORT: XA import failed for handle:%llx:%llx\n",
>                                         handle.handle_hi, handle.handle_lo);
>                         goto cleanup;
> -               } else {
> -                       WRITE_ONCE(imp_xa_node->node_state,
> -                                  AMDGPU_UALINK_NODE_READY);
>                 }
>
> ACK.
>
>
> Regards,
>
> Mukul
>
>         }
>
> --
> 2.55.0

True, a revoke can't arrive before NPA-RSP, so
v3 uses your switch and dev_dbg():

Cheers.