}
 
 static int
-aie2_sched_resp_handler(void *handle, const u32 *data, size_t size)
+aie2_sched_resp_handler(void *handle, void __iomem *data, size_t size)
 {
        struct amdxdna_sched_job *job = handle;
        struct amdxdna_gem_obj *cmd_abo;
                goto out;
        }
 
-       status = *data;
+       status = readl(data);
        XDNA_DBG(job->hwctx->client->xdna, "Resp status 0x%x", status);
        if (status == AIE2_STATUS_SUCCESS)
                amdxdna_cmd_set_state(cmd_abo, ERT_CMD_STATE_COMPLETED);
 }
 
 static int
-aie2_sched_nocmd_resp_handler(void *handle, const u32 *data, size_t size)
+aie2_sched_nocmd_resp_handler(void *handle, void __iomem *data, size_t size)
 {
        struct amdxdna_sched_job *job = handle;
        u32 ret = 0;
                goto out;
        }
 
-       status = *data;
+       status = readl(data);
        XDNA_DBG(job->hwctx->client->xdna, "Resp status 0x%x", status);
 
 out:
 }
 
 static int
-aie2_sched_cmdlist_resp_handler(void *handle, const u32 *data, size_t size)
+aie2_sched_cmdlist_resp_handler(void *handle, void __iomem *data, size_t size)
 {
        struct amdxdna_sched_job *job = handle;
        struct amdxdna_gem_obj *cmd_abo;
-       struct cmd_chain_resp *resp;
        struct amdxdna_dev *xdna;
        u32 fail_cmd_status;
        u32 fail_cmd_idx;
+       u32 cmd_status;
        u32 ret = 0;
 
        cmd_abo = job->cmd_bo;
                goto out;
        }
 
-       resp = (struct cmd_chain_resp *)data;
+       cmd_status = readl(data + offsetof(struct cmd_chain_resp, status));
        xdna = job->hwctx->client->xdna;
-       XDNA_DBG(xdna, "Status 0x%x", resp->status);
-       if (resp->status == AIE2_STATUS_SUCCESS) {
+       XDNA_DBG(xdna, "Status 0x%x", cmd_status);
+       if (cmd_status == AIE2_STATUS_SUCCESS) {
                amdxdna_cmd_set_state(cmd_abo, ERT_CMD_STATE_COMPLETED);
                goto out;
        }
 
        /* Slow path to handle error, read from ringbuf on BAR */
-       fail_cmd_idx = resp->fail_cmd_idx;
-       fail_cmd_status = resp->fail_cmd_status;
+       fail_cmd_idx = readl(data + offsetof(struct cmd_chain_resp, fail_cmd_idx));
+       fail_cmd_status = readl(data + offsetof(struct cmd_chain_resp, fail_cmd_status));
        XDNA_DBG(xdna, "Failed cmd idx %d, status 0x%x",
                 fail_cmd_idx, fail_cmd_status);
 
 
        return err_col;
 }
 
-static int aie2_error_async_cb(void *handle, const u32 *data, size_t size)
+static int aie2_error_async_cb(void *handle, void __iomem *data, size_t size)
 {
-       struct async_event_msg_resp *resp;
        struct async_event *e = handle;
 
        if (data) {
-               resp = (struct async_event_msg_resp *)data;
-               e->resp.type = resp->type;
+               e->resp.type = readl(data + offsetof(struct async_event_msg_resp, type));
                wmb(); /* Update status in the end, so that no lock for here */
-               e->resp.status = resp->status;
+               e->resp.status = readl(data + offsetof(struct async_event_msg_resp, status));
        }
        queue_work(e->wq, &e->work);
        return 0;
 
 }
 
 int aie2_register_asyn_event_msg(struct amdxdna_dev_hdl *ndev, dma_addr_t addr, u32 size,
-                                void *handle, int (*cb)(void*, const u32 *, size_t))
+                                void *handle, int (*cb)(void*, void __iomem *, size_t))
 {
        struct async_event_msg_req req = { 0 };
        struct xdna_mailbox_msg msg = {
 }
 
 int aie2_execbuf(struct amdxdna_hwctx *hwctx, struct amdxdna_sched_job *job,
-                int (*notify_cb)(void *, const u32 *, size_t))
+                int (*notify_cb)(void *, void __iomem *, size_t))
 {
        struct mailbox_channel *chann = hwctx->priv->mbox_chann;
        struct amdxdna_dev *xdna = hwctx->client->xdna;
 
 int aie2_cmdlist_multi_execbuf(struct amdxdna_hwctx *hwctx,
                               struct amdxdna_sched_job *job,
-                              int (*notify_cb)(void *, const u32 *, size_t))
+                              int (*notify_cb)(void *, void __iomem *, size_t))
 {
        struct amdxdna_gem_obj *cmdbuf_abo = aie2_cmdlist_get_cmd_buf(job);
        struct mailbox_channel *chann = hwctx->priv->mbox_chann;
 
 int aie2_cmdlist_single_execbuf(struct amdxdna_hwctx *hwctx,
                                struct amdxdna_sched_job *job,
-                               int (*notify_cb)(void *, const u32 *, size_t))
+                               int (*notify_cb)(void *, void __iomem *, size_t))
 {
        struct amdxdna_gem_obj *cmdbuf_abo = aie2_cmdlist_get_cmd_buf(job);
        struct mailbox_channel *chann = hwctx->priv->mbox_chann;
 }
 
 int aie2_sync_bo(struct amdxdna_hwctx *hwctx, struct amdxdna_sched_job *job,
-                int (*notify_cb)(void *, const u32 *, size_t))
+                int (*notify_cb)(void *, void __iomem *, size_t))
 {
        struct mailbox_channel *chann = hwctx->priv->mbox_chann;
        struct amdxdna_gem_obj *abo = to_xdna_obj(job->bos[0]);
 
 int aie2_map_host_buf(struct amdxdna_dev_hdl *ndev, u32 context_id, u64 addr, u64 size);
 int aie2_query_status(struct amdxdna_dev_hdl *ndev, char __user *buf, u32 size, u32 *cols_filled);
 int aie2_register_asyn_event_msg(struct amdxdna_dev_hdl *ndev, dma_addr_t addr, u32 size,
-                                void *handle, int (*cb)(void*, const u32 *, size_t));
+                                void *handle, int (*cb)(void*, void __iomem *, size_t));
 int aie2_config_cu(struct amdxdna_hwctx *hwctx);
 int aie2_execbuf(struct amdxdna_hwctx *hwctx, struct amdxdna_sched_job *job,
-                int (*notify_cb)(void *, const u32 *, size_t));
+                int (*notify_cb)(void *, void __iomem *, size_t));
 int aie2_cmdlist_single_execbuf(struct amdxdna_hwctx *hwctx,
                                struct amdxdna_sched_job *job,
-                               int (*notify_cb)(void *, const u32 *, size_t));
+                               int (*notify_cb)(void *, void __iomem *, size_t));
 int aie2_cmdlist_multi_execbuf(struct amdxdna_hwctx *hwctx,
                               struct amdxdna_sched_job *job,
-                              int (*notify_cb)(void *, const u32 *, size_t));
+                              int (*notify_cb)(void *, void __iomem *, size_t));
 int aie2_sync_bo(struct amdxdna_hwctx *hwctx, struct amdxdna_sched_job *job,
-                int (*notify_cb)(void *, const u32 *, size_t));
+                int (*notify_cb)(void *, void __iomem *, size_t));
 
 /* aie2_hwctx.c */
 int aie2_hwctx_init(struct amdxdna_hwctx *hwctx);
 
 
 struct mailbox_msg {
        void                    *handle;
-       int                     (*notify_cb)(void *handle, const u32 *data, size_t size);
+       int                     (*notify_cb)(void *handle, void __iomem *data, size_t size);
        size_t                  pkg_size; /* package size in bytes */
        struct mailbox_pkg      pkg;
 };
 
 static int
 mailbox_get_resp(struct mailbox_channel *mb_chann, struct xdna_msg_header *header,
-                void *data)
+                void __iomem *data)
 {
        struct mailbox_msg *mb_msg;
        int msg_id;
        memcpy_fromio((u32 *)&header + 1, read_addr, rest);
        read_addr += rest;
 
-       ret = mailbox_get_resp(mb_chann, &header, (u32 *)read_addr);
+       ret = mailbox_get_resp(mb_chann, &header, read_addr);
 
        mailbox_set_headptr(mb_chann, head + msg_size);
        /* After update head, it can equal to ringbuf_size. This is expected. */
 
 struct xdna_mailbox_msg {
        u32             opcode;
        void            *handle;
-       int             (*notify_cb)(void *handle, const u32 *data, size_t size);
+       int             (*notify_cb)(void *handle, void __iomem *data, size_t size);
        u8              *send_data;
        size_t          send_size;
 };
 
 #include "amdxdna_mailbox_helper.h"
 #include "amdxdna_pci_drv.h"
 
-int xdna_msg_cb(void *handle, const u32 *data, size_t size)
+int xdna_msg_cb(void *handle, void __iomem *data, size_t size)
 {
        struct xdna_notify *cb_arg = handle;
        int ret;
                goto out;
        }
 
+       memcpy_fromio(cb_arg->data, data, cb_arg->size);
        print_hex_dump_debug("resp data: ", DUMP_PREFIX_OFFSET,
-                            16, 4, data, cb_arg->size, true);
-       memcpy(cb_arg->data, data, cb_arg->size);
+                            16, 4, cb_arg->data, cb_arg->size, true);
 out:
        ret = cb_arg->error;
        complete(&cb_arg->comp);
 
                .notify_cb = xdna_msg_cb,                               \
        }
 
-int xdna_msg_cb(void *handle, const u32 *data, size_t size);
+int xdna_msg_cb(void *handle, void __iomem *data, size_t size);
 int xdna_send_msg_wait(struct amdxdna_dev *xdna, struct mailbox_channel *chann,
                       struct xdna_mailbox_msg *msg);