diff --git a/drivers/accel/amdxdna/aie2_message.c b/drivers/accel/amdxdna/aie2_message.c
index 27fccce2655f..3249b810bf7c 100644
--- a/drivers/accel/amdxdna/aie2_message.c
+++ b/drivers/accel/amdxdna/aie2_message.c
@@ -31,6 +31,7 @@ void *aie2_alloc_msg_buffer(struct amdxdna_dev_hdl *ndev, u32 *size,
 			    dma_addr_t *dma_addr)
 {
 	struct amdxdna_dev *xdna = ndev->aie.xdna;
+	void *vaddr;
 	int order;
 
 	*size = max(*size, SZ_8K);
@@ -42,8 +43,12 @@ void *aie2_alloc_msg_buffer(struct amdxdna_dev_hdl *ndev, u32 *size,
 	if (amdxdna_iova_on(xdna))
 		return amdxdna_iommu_alloc(xdna, *size, dma_addr);
 
-	return dma_alloc_noncoherent(xdna->ddev.dev, *size, dma_addr,
+	vaddr = dma_alloc_noncoherent(xdna->ddev.dev, *size, dma_addr,
 				      DMA_FROM_DEVICE, GFP_KERNEL);
+	if (!vaddr)
+		return ERR_PTR(-ENOMEM);
+
+	return vaddr;
 }
 
 void aie2_free_msg_buffer(struct amdxdna_dev_hdl *ndev, size_t size,