accel/amdxdna: add DRAM work buffer support for AIE4 devices

NPU firmware requires a host-allocated DRAM work buffer before any
hardware contexts can be created. Allocate a 4 MB (power-of-2
aligned) buffer during device init via the common
amdxdna_alloc_msg_buffer() helper and attach it to firmware after
mailbox init. The firmware releases the buffer internally during its
suspend sequence, so no explicit detach is needed from the driver.

Add the attach_work_buffer mailbox message, using GENMASK/FIELD_PREP
for PASID register encoding.

Signed-off-by: Nishad Saraf <nishads@amd.com>

Author: NishadSaraf

drivers/accel/amdxdna/aie4_message.c

@@ -25,3 +25,23 @@ int aie4_suspend_fw(struct amdxdna_dev_hdl *ndev)
 
 	return ret;
 }
+
+int aie4_attach_work_buffer(struct amdxdna_dev_hdl *ndev, u32 pasid,
+			    dma_addr_t addr, u32 size)
+{
+	DECLARE_AIE_MSG(aie4_msg_attach_work_buffer, AIE4_MSG_OP_ATTACH_WORK_BUFFER);
+	struct amdxdna_dev *xdna = ndev->aie.xdna;
+	int ret;
+
+	req.buff_addr = addr;
+	req.buff_size = size;
+	req.pasid = FIELD_PREP(AIE4_MSG_PASID_MASK, pasid);
+	if (pasid)
+		req.pasid |= AIE4_MSG_PASID_VLD;
+
+	ret = aie_send_mgmt_msg_wait(&ndev->aie, &msg);
+	if (ret)
+		XDNA_ERR(xdna, "Failed to attach work buffer, ret %d", ret);
+
+	return ret;
+}

drivers/accel/amdxdna/aie4_msg_priv.h

@@ -6,11 +6,13 @@
 #ifndef _AIE4_MSG_PRIV_H_
 #define _AIE4_MSG_PRIV_H_
 
+#include <linux/bitfield.h>
+#include <linux/sizes.h>
 #include <linux/types.h>
 
 enum aie4_msg_opcode {
 	AIE4_MSG_OP_SUSPEND                          = 0x10003,
-
+	AIE4_MSG_OP_ATTACH_WORK_BUFFER               = 0x1000D,
 	AIE4_MSG_OP_CREATE_VFS                       = 0x20001,
 	AIE4_MSG_OP_DESTROY_VFS                      = 0x20002,
 };
@@ -46,4 +48,19 @@ struct aie4_msg_destroy_vfs_resp {
 	enum aie4_msg_status status;
 } __packed;
 
+#define AIE4_MSG_PASID_MASK		GENMASK(19, 0)
+#define AIE4_MSG_PASID_VLD		BIT(31)
+
+#define AIE4_WORK_BUFFER_MIN_SIZE	SZ_4M
+
+struct aie4_msg_attach_work_buffer_req {
+	__u64 buff_addr;
+	__u32 pasid;
+	__u32 buff_size;
+} __packed;
+
+struct aie4_msg_attach_work_buffer_resp {
+	enum aie4_msg_status status;
+} __packed;
+
 #endif /* _AIE4_MSG_PRIV_H_ */

drivers/accel/amdxdna/aie4_pci.c

@@ -6,9 +6,11 @@
 #include "drm/amdxdna_accel.h"
 #include <drm/drm_managed.h>
 #include <drm/drm_print.h>
+#include <linux/dma-mapping.h>
 #include <linux/firmware.h>
 #include <linux/sizes.h>
 
+#include "aie4_msg_priv.h"
 #include "aie4_pci.h"
 #include "amdxdna_pci_drv.h"
 
@@ -234,6 +236,40 @@ static int aie4_fw_load(struct amdxdna_dev_hdl *ndev)
 	return ret;
 }
 
+static int aie4_alloc_work_buffer(struct amdxdna_dev_hdl *ndev)
+{
+	struct amdxdna_dev *xdna = ndev->aie.xdna;
+	u32 buf_size = AIE4_WORK_BUFFER_MIN_SIZE;
+
+	ndev->work_buf = amdxdna_alloc_msg_buffer(xdna, &buf_size,
+						  &ndev->work_buf_addr);
+	if (IS_ERR(ndev->work_buf)) {
+		int ret = PTR_ERR(ndev->work_buf);
+
+		XDNA_ERR(xdna, "Failed to alloc work buffer, size 0x%x",
+			 AIE4_WORK_BUFFER_MIN_SIZE);
+		ndev->work_buf = NULL;
+		return ret;
+	}
+
+	ndev->work_buf_size = buf_size;
+	XDNA_DBG(xdna, "Work buffer allocated: size 0x%x", buf_size);
+
+	return 0;
+}
+
+static void aie4_free_work_buffer(struct amdxdna_dev_hdl *ndev)
+{
+	struct amdxdna_dev *xdna = ndev->aie.xdna;
+
+	if (!ndev->work_buf)
+		return;
+
+	amdxdna_free_msg_buffer(xdna, ndev->work_buf_size,
+				ndev->work_buf, ndev->work_buf_addr);
+	ndev->work_buf = NULL;
+}
+
 static int aie4_hw_start(struct amdxdna_dev *xdna)
 {
 	struct amdxdna_dev_hdl *ndev = xdna->dev_handle;
@@ -247,8 +283,18 @@ static int aie4_hw_start(struct amdxdna_dev *xdna)
 	if (ret)
 		goto fw_unload;
 
+	if (ndev->work_buf) {
+		/* Firmware releases the DRAM work buffer internally during suspend */
+		ret = aie4_attach_work_buffer(ndev, 0, ndev->work_buf_addr,
+					      ndev->work_buf_size);
+		if (ret)
+			goto mbox_fini;
+	}
+
 	return 0;
 
+mbox_fini:
+	aie4_mailbox_fini(ndev);
 fw_unload:
 	aie4_fw_unload(ndev);
 
@@ -459,6 +505,7 @@ static void aie4_fini(struct amdxdna_dev *xdna)
 
 	aie4_sriov_stop(ndev);
 	aie4_pcidev_fini(ndev);
+	aie4_free_work_buffer(ndev);
 }
 
 static int aie4_init(struct amdxdna_dev *xdna)
@@ -474,15 +521,23 @@ static int aie4_init(struct amdxdna_dev *xdna)
 	ndev->aie.xdna = xdna;
 	xdna->dev_handle = ndev;
 
+	ret = aie4_alloc_work_buffer(ndev);
+	if (ret)
+		return ret;
+
 	ret = aie4_pcidev_init(ndev);
 	if (ret) {
 		XDNA_ERR(xdna, "Setup PCI device failed, ret %d", ret);
-		return ret;
+		goto free_work_buf;
 	}
 
 	amdxdna_vbnv_init(xdna);
 	XDNA_DBG(xdna, "aie4 init finished");
 	return 0;
+
+free_work_buf:
+	aie4_free_work_buffer(ndev);
+	return ret;
 }
 
 const struct amdxdna_dev_ops aie4_ops = {

drivers/accel/amdxdna/aie4_pci.h

@@ -31,10 +31,15 @@ struct amdxdna_dev_hdl {
 	void			__iomem *rbuf_base;
 
 	struct mailbox			*mbox;
+	void				*work_buf;
+	dma_addr_t			work_buf_addr;
+	u32				work_buf_size;
 };
 
 /* aie4_message.c */
 int aie4_suspend_fw(struct amdxdna_dev_hdl *ndev);
+int aie4_attach_work_buffer(struct amdxdna_dev_hdl *ndev, u32 pasid,
+			    dma_addr_t addr, u32 size);
 
 /* aie4_sriov.c */
 #if IS_ENABLED(CONFIG_PCI_IOV)