Break circulare dependency on struct ethosu_device

The 'struct ethosu_device' has been passed as argument to classes.
This creates a circular dependency dependency and it gives all
classes full visibility to all resources in the device struct. This
patch removes the circular dependency.

Using device_lock() and device_unlock() to for synchronization.

Change-Id: I8322e6530c72d7bd67f48f411b4f14b612be2706

15 files changed, 263 insertions, 223 deletions

diff --git a/kernel/ethosu_buffer.c b/kernel/ethosu_buffer.c
index a83a95a..0fcbf3b 100644
--- a/kernel/ethosu_buffer.c
+++ b/kernel/ethosu_buffer.c
@@ -89,21 +89,24 @@ static void ethosu_buffer_destroy(struct kref *kref)
 {
 	struct ethosu_buffer *buf =
 		container_of(kref, struct ethosu_buffer, kref);
+	struct device *dev = buf->dev;
+	struct rproc *rproc = rproc_get_by_child(dev);
 
-	dev_info(buf->edev->dev, "Buffer destroy. buf=0x%pK\n", buf);
+	dev_info(dev, "Buffer destroy. buf=0x%pK\n", buf);
 
-	dma_free_coherent(buf->dev, buf->capacity, buf->cpu_addr,
+	dma_free_coherent(rproc->dev.parent, buf->capacity, buf->cpu_addr,
 			  buf->dma_addr);
 
-	devm_kfree(buf->edev->dev, buf);
+	devm_kfree(dev, buf);
 }
 
 static int ethosu_buffer_release(struct inode *inode,
 				 struct file *file)
 {
 	struct ethosu_buffer *buf = file->private_data;
+	struct device *dev = buf->dev;
 
-	dev_info(buf->edev->dev, "Buffer release. file=0x%pK, buf=0x%pK\n",
+	dev_info(dev, "Buffer release. file=0x%pK, buf=0x%pK\n",
 		 file, buf);
 
 	ethosu_buffer_put(buf);
@@ -115,12 +118,13 @@ static int ethosu_buffer_mmap(struct file *file,
 			      struct vm_area_struct *vma)
 {
 	struct ethosu_buffer *buf = file->private_data;
+	struct device *dev = buf->dev;
 	int ret;
 
-	dev_info(buf->edev->dev, "Buffer mmap. file=0x%pK, buf=0x%pK\n",
+	dev_info(dev, "Buffer mmap. file=0x%pK, buf=0x%pK\n",
 		 file, buf);
 
-	ret = dma_mmap_coherent(buf->edev->dev, vma, buf->cpu_addr,
+	ret = dma_mmap_coherent(dev, vma, buf->cpu_addr,
 				buf->dma_addr, buf->capacity);
 
 	return ret;
@@ -131,14 +135,15 @@ static long ethosu_buffer_ioctl(struct file *file,
 				unsigned long arg)
 {
 	struct ethosu_buffer *buf = file->private_data;
+	struct device *dev = buf->dev;
 	void __user *udata = (void __user *)arg;
 	int ret = -EINVAL;
 
-	ret = mutex_lock_interruptible(&buf->edev->mutex);
+	ret = device_lock_interruptible(dev);
 	if (ret)
 		return ret;
 
-	dev_info(buf->edev->dev,
+	dev_info(dev,
 		 "Buffer ioctl. file=0x%pK, buf=0x%pK, cmd=0x%x, arg=%lu\n",
 		 file, buf, cmd, arg);
 
@@ -149,7 +154,7 @@ static long ethosu_buffer_ioctl(struct file *file,
 		if (copy_from_user(&uapi, udata, sizeof(uapi)))
 			break;
 
-		dev_info(buf->edev->dev,
+		dev_info(dev,
 			 "Buffer ioctl: Buffer set. size=%u, offset=%u\n",
 			 uapi.size, uapi.offset);
 
@@ -162,7 +167,7 @@ static long ethosu_buffer_ioctl(struct file *file,
 		uapi.size = buf->size;
 		uapi.offset = buf->offset;
 
-		dev_info(buf->edev->dev,
+		dev_info(dev,
 			 "Buffer ioctl: Buffer get. size=%u, offset=%u\n",
 			 uapi.size, uapi.offset);
 
@@ -173,40 +178,38 @@ static long ethosu_buffer_ioctl(struct file *file,
 		break;
 	}
 	default: {
-		dev_err(buf->edev->dev, "Invalid ioctl. cmd=%u, arg=%lu",
+		dev_err(dev, "Invalid ioctl. cmd=%u, arg=%lu",
 			cmd, arg);
 		break;
 	}
 	}
 
-	mutex_unlock(&buf->edev->mutex);
+	device_unlock(dev);
 
 	return ret;
 }
 
-int ethosu_buffer_create(struct ethosu_device *edev,
+int ethosu_buffer_create(struct device *dev,
 			 size_t capacity)
 {
-	struct rproc *rproc = rproc_get_by_child(edev->dev);
-	struct device *dev = rproc->dev.parent;
+	struct rproc *rproc = rproc_get_by_child(dev);
 	struct ethosu_buffer *buf;
 	int ret = -ENOMEM;
 
 	if (!capacity)
 		return -EINVAL;
 
-	buf = devm_kzalloc(edev->dev, sizeof(*buf), GFP_KERNEL);
+	buf = devm_kzalloc(dev, sizeof(*buf), GFP_KERNEL);
 	if (!buf)
 		return -ENOMEM;
 
-	buf->edev = edev;
 	buf->dev = dev;
 	buf->capacity = capacity;
 	buf->offset = 0;
 	buf->size = 0;
 	kref_init(&buf->kref);
 
-	buf->cpu_addr = dma_alloc_coherent(dev, capacity,
+	buf->cpu_addr = dma_alloc_coherent(rproc->dev.parent, capacity,
 					   &buf->dma_addr, GFP_KERNEL);
 	if (!buf->cpu_addr)
 		goto free_buf;
@@ -219,7 +222,7 @@ int ethosu_buffer_create(struct ethosu_device *edev,
 	buf->file = fget(ret);
 	fput(buf->file);
 
-	dev_info(buf->edev->dev,
+	dev_info(dev,
 		 "Buffer create. file=0x%pK, fd=%d, buf=0x%pK, capacity=%zu, cpu_addr=0x%pK, dma_addr=0x%llx, phys_addr=0x%llx\n",
 		 buf->file, ret, buf, capacity, buf->cpu_addr, buf->dma_addr,
 		 virt_to_phys(buf->cpu_addr));
@@ -227,11 +230,11 @@ int ethosu_buffer_create(struct ethosu_device *edev,
 	return ret;
 
 free_dma:
-	dma_free_coherent(buf->edev->dev, buf->capacity, buf->cpu_addr,
+	dma_free_coherent(rproc->dev.parent, buf->capacity, buf->cpu_addr,
 			  buf->dma_addr);
 
 free_buf:
-	devm_kfree(buf->edev->dev, buf);
+	devm_kfree(dev, buf);
 
 	return ret;
 }
diff --git a/kernel/ethosu_buffer.h b/kernel/ethosu_buffer.h
index ada8c4c..bc5958c 100644
--- a/kernel/ethosu_buffer.h
+++ b/kernel/ethosu_buffer.h
@@ -50,15 +50,14 @@ struct device;
  * 'offset + size' must not be larger than 'capacity'.
  */
 struct ethosu_buffer {
-	struct ethosu_device *edev;
-	struct device        *dev;
-	struct file          *file;
-	struct kref          kref;
-	size_t               capacity;
-	size_t               offset;
-	size_t               size;
-	void                 *cpu_addr;
-	dma_addr_t           dma_addr;
+	struct device *dev;
+	struct file   *file;
+	struct kref   kref;
+	size_t        capacity;
+	size_t        offset;
+	size_t        size;
+	void          *cpu_addr;
+	dma_addr_t    dma_addr;
 };
 
 /****************************************************************************
@@ -72,7 +71,7 @@ struct ethosu_buffer {
  *
  * Return: fd on success, else error code.
  */
-int ethosu_buffer_create(struct ethosu_device *edev,
+int ethosu_buffer_create(struct device *dev,
 			 size_t capacity);
 
 /**
diff --git a/kernel/ethosu_cancel_inference.c b/kernel/ethosu_cancel_inference.c
index 61f7a9c..ee630f5 100644
--- a/kernel/ethosu_cancel_inference.c
+++ b/kernel/ethosu_cancel_inference.c
@@ -41,9 +41,10 @@
  ****************************************************************************/
 
 static int ethosu_cancel_inference_send(
-	struct ethosu_cancel_inference *cancellation)
+	struct ethosu_cancel_inference *cancellation,
+	struct ethosu_mailbox *mailbox)
 {
-	return ethosu_mailbox_cancel_inference(&cancellation->edev->mailbox,
+	return ethosu_mailbox_cancel_inference(mailbox,
 					       &cancellation->msg,
 					       cancellation->inf->msg.id);
 }
@@ -61,7 +62,9 @@ static void ethosu_cancel_inference_fail(struct ethosu_mailbox_msg *msg)
 	complete(&cancellation->done);
 }
 
-int ethosu_cancel_inference_request(struct ethosu_inference *inf,
+int ethosu_cancel_inference_request(struct device *dev,
+				    struct ethosu_mailbox *mailbox,
+				    struct ethosu_inference *inf,
 				    struct ethosu_uapi_cancel_inference_status *uapi)
 {
 	struct ethosu_cancel_inference *cancellation;
@@ -75,7 +78,7 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	}
 
 	cancellation =
-		devm_kzalloc(inf->edev->dev,
+		devm_kzalloc(dev,
 			     sizeof(struct ethosu_cancel_inference),
 			     GFP_KERNEL);
 	if (!cancellation)
@@ -86,39 +89,39 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	/* mark inference ABORTING to avoid resending the inference message */
 	inf->status = ETHOSU_CORE_STATUS_ABORTING;
 
-	cancellation->edev = inf->edev;
+	cancellation->dev = dev;
 	cancellation->inf = inf;
 	cancellation->uapi = uapi;
 	init_completion(&cancellation->done);
 	cancellation->msg.fail = ethosu_cancel_inference_fail;
 
-	ret = ethosu_mailbox_register(&cancellation->edev->mailbox,
+	ret = ethosu_mailbox_register(mailbox,
 				      &cancellation->msg);
 	if (ret < 0)
 		goto kfree;
 
-	dev_info(cancellation->edev->dev,
-		 "Inference cancellation create. cancel=0x%pK, msg.id=%d\n",
+	dev_info(dev,
+		 "Inference cancellation create. cancel=0x%pK, msg.id=%ddev",
 		 cancellation, cancellation->msg.id);
 
-	ret = ethosu_cancel_inference_send(cancellation);
+	ret = ethosu_cancel_inference_send(cancellation, mailbox);
 	if (0 != ret)
 		goto deregister;
 
 	/* Unlock the mutex before going to block on the condition */
-	mutex_unlock(&cancellation->edev->mutex);
+	device_unlock(dev);
 	/* wait for response to arrive back */
 	timeout = wait_for_completion_timeout(&cancellation->done,
 					      msecs_to_jiffies(
 						      CANCEL_INFERENCE_RESP_TIMEOUT_MS));
 	/* take back the mutex before resuming to do anything */
-	ret = mutex_lock_interruptible(&cancellation->edev->mutex);
+	ret = device_lock_interruptible(dev);
 	if (0 != ret)
 		goto deregister;
 
 	if (0 == timeout /* timed out*/) {
-		dev_warn(inf->edev->dev,
-			 "Msg: Cancel Inference response lost - timeout\n");
+		dev_warn(dev,
+			 "Msg: Cancel Inference response lost - timeoutdev");
 		ret = -EIO;
 		goto deregister;
 	}
@@ -129,30 +132,31 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	}
 
 deregister:
-	ethosu_mailbox_deregister(&cancellation->edev->mailbox,
+	ethosu_mailbox_deregister(mailbox,
 				  &cancellation->msg);
 
 kfree:
-	dev_info(cancellation->edev->dev,
-		 "Cancel inference destroy. cancel=0x%pK\n", cancellation);
+	dev_info(dev,
+		 "Cancel inference destroy. cancel=0x%pK", cancellation);
 	/* decrease the reference on the inference we are refering to */
 	ethosu_inference_put(cancellation->inf);
-	devm_kfree(cancellation->edev->dev, cancellation);
+	devm_kfree(dev, cancellation);
 
 	return ret;
 }
 
-void ethosu_cancel_inference_rsp(struct ethosu_device *edev,
+void ethosu_cancel_inference_rsp(struct ethosu_mailbox *mailbox,
 				 int msg_id,
 				 struct ethosu_core_msg_cancel_inference_rsp *rsp)
 {
+	struct device *dev = mailbox->dev;
 	struct ethosu_mailbox_msg *msg;
 	struct ethosu_cancel_inference *cancellation;
 
-	msg = ethosu_mailbox_find(&edev->mailbox, msg_id);
+	msg = ethosu_mailbox_find(mailbox, msg_id);
 	if (IS_ERR(msg)) {
-		dev_warn(edev->dev,
-			 "Id for cancel inference msg not found. id=%d\n",
+		dev_warn(dev,
+			 "Id for cancel inference msg not found. id=%ddev",
 			 msg_id);
 
 		return;
diff --git a/kernel/ethosu_cancel_inference.h b/kernel/ethosu_cancel_inference.h
index a21b9c0..a3982dd 100644
--- a/kernel/ethosu_cancel_inference.h
+++ b/kernel/ethosu_cancel_inference.h
@@ -36,12 +36,11 @@
  ****************************************************************************/
 
 struct ethosu_core_msg_cancel_inference_rsp;
-struct ethosu_device;
 struct ethosu_inference;
 struct ethosu_uapi_cancel_inference_status;
 
 struct ethosu_cancel_inference {
-	struct ethosu_device                       *edev;
+	struct device                              *dev;
 	struct ethosu_inference                    *inf;
 	struct ethosu_uapi_cancel_inference_status *uapi;
 	struct completion                          done;
@@ -58,13 +57,15 @@ struct ethosu_cancel_inference {
  *
  * Return: 0 on success, error code otherwise.
  */
-int ethosu_cancel_inference_request(struct ethosu_inference *inf,
+int ethosu_cancel_inference_request(struct device *dev,
+				    struct ethosu_mailbox *mailbox,
+				    struct ethosu_inference *inf,
 				    struct ethosu_uapi_cancel_inference_status *uapi);
 
 /**
  * ethosu_cancel_inference_rsp() - Handle cancel inference response
  */
-void ethosu_cancel_inference_rsp(struct ethosu_device *edev,
+void ethosu_cancel_inference_rsp(struct ethosu_mailbox *mailbox,
 				 int msg_id,
 				 struct ethosu_core_msg_cancel_inference_rsp *rsp);
 
diff --git a/kernel/ethosu_capabilities.c b/kernel/ethosu_capabilities.c
index 57db452..8611edf 100644
--- a/kernel/ethosu_capabilities.c
+++ b/kernel/ethosu_capabilities.c
@@ -41,9 +41,10 @@
  * Functions
  ****************************************************************************/
 
-static inline int ethosu_capabilities_send(struct ethosu_capabilities *cap)
+static inline int ethosu_capabilities_send(struct ethosu_capabilities *cap,
+					   struct ethosu_mailbox *mailbox)
 {
-	return ethosu_mailbox_capabilities_request(&cap->edev->mailbox,
+	return ethosu_mailbox_capabilities_request(mailbox,
 						   &cap->msg);
 }
 
@@ -59,16 +60,17 @@ static void ethosu_capabilities_fail(struct ethosu_mailbox_msg *msg)
 	complete(&cap->done);
 }
 
-void ethosu_capability_rsp(struct ethosu_device *edev,
+void ethosu_capability_rsp(struct ethosu_mailbox *mailbox,
 			   int msg_id,
 			   struct ethosu_core_msg_capabilities_rsp *rsp)
 {
+	struct device *dev = mailbox->dev;
 	struct ethosu_mailbox_msg *msg;
 	struct ethosu_capabilities *cap;
 
-	msg = ethosu_mailbox_find(&edev->mailbox, msg_id);
+	msg = ethosu_mailbox_find(mailbox, msg_id);
 	if (IS_ERR(msg)) {
-		dev_warn(edev->dev,
+		dev_warn(dev,
 			 "Id for capabilities msg not found. id=%d\n",
 			 msg_id);
 
@@ -98,36 +100,37 @@ void ethosu_capability_rsp(struct ethosu_device *edev,
 	complete(&cap->done);
 }
 
-int ethosu_capabilities_request(struct ethosu_device *edev,
+int ethosu_capabilities_request(struct device *dev,
+				struct ethosu_mailbox *mailbox,
 				struct ethosu_uapi_device_capabilities *uapi)
 {
 	struct ethosu_capabilities *cap;
 	int ret;
 	int timeout;
 
-	cap = devm_kzalloc(edev->dev, sizeof(struct ethosu_capabilities),
+	cap = devm_kzalloc(dev, sizeof(struct ethosu_capabilities),
 			   GFP_KERNEL);
 	if (!cap)
 		return -ENOMEM;
 
-	cap->edev = edev;
+	cap->dev = dev;
 	cap->uapi = uapi;
 	init_completion(&cap->done);
 	cap->msg.fail = ethosu_capabilities_fail;
 
-	ret = ethosu_mailbox_register(&cap->edev->mailbox, &cap->msg);
+	ret = ethosu_mailbox_register(mailbox, &cap->msg);
 	if (ret < 0)
 		goto kfree;
 
-	dev_info(edev->dev, "Capabilities create. Id=%d, handle=0x%p\n",
+	dev_info(dev, "Capabilities create. Id=%d, handle=0x%p\n",
 		 cap->msg.id, cap);
 
-	ret = ethosu_capabilities_send(cap);
+	ret = ethosu_capabilities_send(cap, mailbox);
 	if (0 != ret)
 		goto deregister;
 
 	/* Unlock the mutex before going to block on the condition */
-	mutex_unlock(&edev->mutex);
+	device_unlock(dev);
 
 	/* wait for response to arrive back */
 	timeout = wait_for_completion_timeout(&cap->done,
@@ -135,10 +138,10 @@ int ethosu_capabilities_request(struct ethosu_device *edev,
 						      CAPABILITIES_RESP_TIMEOUT_MS));
 
 	/* take back the mutex before resuming to do anything */
-	mutex_lock(&edev->mutex);
+	device_lock(dev);
 
 	if (0 == timeout) {
-		dev_warn(edev->dev, "Capabilities response timeout");
+		dev_warn(dev, "Capabilities response timeout");
 		ret = -ETIME;
 		goto deregister;
 	}
@@ -149,12 +152,12 @@ int ethosu_capabilities_request(struct ethosu_device *edev,
 	}
 
 deregister:
-	ethosu_mailbox_deregister(&cap->edev->mailbox, &cap->msg);
+	ethosu_mailbox_deregister(mailbox, &cap->msg);
 
 kfree:
-	dev_info(cap->edev->dev, "Capabilities destroy. Id=%d, handle=0x%p\n",
+	dev_info(dev, "Capabilities destroy. Id=%d, handle=0x%p\n",
 		 cap->msg.id, cap);
-	devm_kfree(cap->edev->dev, cap);
+	devm_kfree(dev, cap);
 
 	return ret;
 }
diff --git a/kernel/ethosu_capabilities.h b/kernel/ethosu_capabilities.h
index 3a79dc5..30bf475 100644
--- a/kernel/ethosu_capabilities.h
+++ b/kernel/ethosu_capabilities.h
@@ -42,7 +42,7 @@ struct ethosu_uapi_device_capabilities;
  * struct ethosu_capabilities - Capabilities internal struct
  */
 struct ethosu_capabilities {
-	struct ethosu_device                   *edev;
+	struct device                          *dev;
 	struct completion                      done;
 	struct ethosu_uapi_device_capabilities *uapi;
 	struct ethosu_mailbox_msg              msg;
@@ -53,10 +53,11 @@ struct ethosu_capabilities {
  * Functions
  ****************************************************************************/
 
-int ethosu_capabilities_request(struct ethosu_device *edev,
+int ethosu_capabilities_request(struct device *dev,
+				struct ethosu_mailbox *mailbox,
 				struct ethosu_uapi_device_capabilities *uapi);
 
-void ethosu_capability_rsp(struct ethosu_device *edev,
+void ethosu_capability_rsp(struct ethosu_mailbox *mailbox,
 			   int msg_id,
 			   struct ethosu_core_msg_capabilities_rsp *rsp);
 
diff --git a/kernel/ethosu_device.c b/kernel/ethosu_device.c
index 004b26f..e99f8c0 100644
--- a/kernel/ethosu_device.c
+++ b/kernel/ethosu_device.c
@@ -36,26 +36,22 @@
 #include <linux/dma-mapping.h>
 #include <linux/errno.h>
 #include <linux/fs.h>
-#include <linux/interrupt.h>
 #include <linux/io.h>
 #include <linux/of_reserved_mem.h>
+#include <linux/remoteproc.h>
 #include <linux/slab.h>
 #include <linux/uaccess.h>
 
 /****************************************************************************
- * Defines
- ****************************************************************************/
-
-/****************************************************************************
  * Functions
  ****************************************************************************/
 
 /* Incoming messages */
-static int ethosu_handle_msg(struct rpmsg_device *rpdev,
-			     void *data,
-			     int len,
-			     void *priv,
-			     u32 src)
+static int ethosu_handle_rpmsg(struct rpmsg_device *rpdev,
+			       void *data,
+			       int len,
+			       void *priv,
+			       u32 src)
 {
 	struct ethosu_device *edev = dev_get_drvdata(&rpdev->dev);
 	struct device *dev = &rpdev->dev;
@@ -68,7 +64,7 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 		 rpmsg->header.magic, rpmsg->header.type, rpmsg->header.msg_id);
 
 	switch (rpmsg->header.type) {
-	case ETHOSU_CORE_MSG_ERR:
+	case ETHOSU_CORE_MSG_ERR: {
 		if (length != sizeof(rpmsg->error)) {
 			dev_warn(dev,
 				 "Msg: Error message of incorrect size. size=%u, expected=%zu", length,
@@ -80,8 +76,8 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 		rpmsg->error.msg[sizeof(rpmsg->error.msg) - 1] = '\0';
 		dev_warn(dev, "Msg: Error. type=%u, msg=\"%s\"",
 			 rpmsg->error.type, rpmsg->error.msg);
-		ret = -EBADMSG;
 		break;
+	}
 	case ETHOSU_CORE_MSG_PING:
 		dev_info(dev, "Msg: Ping");
 		ret = ethosu_mailbox_pong(&edev->mailbox);
@@ -102,7 +98,7 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 			 "Msg: Inference response. ofm_count=%u, status=%u",
 			 rpmsg->inf_rsp.ofm_count, rpmsg->inf_rsp.status);
 
-		ethosu_inference_rsp(edev, rpmsg->header.msg_id,
+		ethosu_inference_rsp(&edev->mailbox, rpmsg->header.msg_id,
 				     &rpmsg->inf_rsp);
 		break;
 	case ETHOSU_CORE_MSG_CANCEL_INFERENCE_RSP:
@@ -117,7 +113,8 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 		dev_info(dev,
 			 "Msg: Cancel Inference response. status=%u",
 			 rpmsg->cancel_rsp.status);
-		ethosu_cancel_inference_rsp(edev, rpmsg->header.msg_id,
+		ethosu_cancel_inference_rsp(&edev->mailbox,
+					    rpmsg->header.msg_id,
 					    &rpmsg->cancel_rsp);
 		break;
 	case ETHOSU_CORE_MSG_VERSION_RSP:
@@ -169,7 +166,7 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 			 rpmsg->cap_rsp.cmd_stream_version,
 			 rpmsg->cap_rsp.custom_dma);
 
-		ethosu_capability_rsp(edev, rpmsg->header.msg_id,
+		ethosu_capability_rsp(&edev->mailbox, rpmsg->header.msg_id,
 				      &rpmsg->cap_rsp);
 		break;
 	case ETHOSU_CORE_MSG_NETWORK_INFO_RSP:
@@ -185,7 +182,8 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 			 "Msg: Network info response. status=%u",
 			 rpmsg->net_info_rsp.status);
 
-		ethosu_network_info_rsp(edev, rpmsg->header.msg_id,
+		ethosu_network_info_rsp(&edev->mailbox,
+					rpmsg->header.msg_id,
 					&rpmsg->net_info_rsp);
 
 		break;
@@ -203,12 +201,15 @@ static int ethosu_handle_msg(struct rpmsg_device *rpdev,
 static int ethosu_open(struct inode *inode,
 		       struct file *file)
 {
-	struct ethosu_device *edev =
-		container_of(inode->i_cdev, struct ethosu_device, cdev);
+	struct cdev *cdev = inode->i_cdev;
+	struct ethosu_device *edev = container_of(cdev, struct ethosu_device,
+						  cdev);
+	struct rpmsg_device *rpdev = edev->rpdev;
+	struct device *dev = &rpdev->dev;
 
-	file->private_data = edev;
+	dev_info(dev, "Device open. file=0x%pK", file);
 
-	dev_info(edev->dev, "Device open. file=0x%pK", file);
+	file->private_data = rpdev;
 
 	return nonseekable_open(inode, file);
 }
@@ -217,29 +218,31 @@ static long ethosu_ioctl(struct file *file,
 			 unsigned int cmd,
 			 unsigned long arg)
 {
-	struct ethosu_device *edev = file->private_data;
+	struct rpmsg_device *rpdev = file->private_data;
+	struct device *dev = &rpdev->dev;
+	struct ethosu_device *edev = dev_get_drvdata(dev);
 	void __user *udata = (void __user *)arg;
 	int ret = -EINVAL;
 
-	ret = mutex_lock_interruptible(&edev->mutex);
+	ret = device_lock_interruptible(dev);
 	if (ret)
 		return ret;
 
-	dev_info(edev->dev, "Device ioctl. file=0x%pK, cmd=0x%x, arg=0x%lx",
+	dev_info(dev, "Device ioctl. file=0x%pK, cmd=0x%x, arg=0x%lx",
 		 file, cmd, arg);
 
 	switch (cmd) {
 	case ETHOSU_IOCTL_VERSION_REQ:
-		dev_info(edev->dev, "Device ioctl: Send version request");
+		dev_info(dev, "Device ioctl: Send version request");
 		ret = ethosu_mailbox_version_request(&edev->mailbox);
 		break;
 	case ETHOSU_IOCTL_CAPABILITIES_REQ: {
 		struct ethosu_uapi_device_capabilities uapi;
 
-		dev_info(edev->dev,
+		dev_info(dev,
 			 "Device ioctl: Send capabilities request");
 
-		ret = ethosu_capabilities_request(edev, &uapi);
+		ret = ethosu_capabilities_request(dev, &edev->mailbox, &uapi);
 		if (ret)
 			break;
 
@@ -247,7 +250,7 @@ static long ethosu_ioctl(struct file *file,
 		break;
 	}
 	case ETHOSU_IOCTL_PING: {
-		dev_info(edev->dev, "Device ioctl: Send ping");
+		dev_info(dev, "Device ioctl: Send ping");
 		ret = ethosu_mailbox_ping(&edev->mailbox);
 		break;
 	}
@@ -257,11 +260,11 @@ static long ethosu_ioctl(struct file *file,
 		if (copy_from_user(&uapi, udata, sizeof(uapi)))
 			break;
 
-		dev_info(edev->dev,
+		dev_info(dev,
 			 "Device ioctl: Buffer create. capacity=%u",
 			 uapi.capacity);
 
-		ret = ethosu_buffer_create(edev, uapi.capacity);
+		ret = ethosu_buffer_create(dev, uapi.capacity);
 		break;
 	}
 	case ETHOSU_IOCTL_NETWORK_CREATE: {
@@ -270,21 +273,21 @@ static long ethosu_ioctl(struct file *file,
 		if (copy_from_user(&uapi, udata, sizeof(uapi)))
 			break;
 
-		dev_info(edev->dev,
+		dev_info(dev,
 			 "Device ioctl: Network create. type=%u, fd/index=%u",
 			 uapi.type, uapi.fd);
 
-		ret = ethosu_network_create(edev, &uapi);
+		ret = ethosu_network_create(dev, &edev->mailbox, &uapi);
 		break;
 	}
 	default: {
-		dev_err(edev->dev, "Invalid ioctl. cmd=%u, arg=%lu",
+		dev_err(dev, "Invalid ioctl. cmd=%u, arg=%lu",
 			cmd, arg);
 		break;
 	}
 	}
 
-	mutex_unlock(&edev->mutex);
+	device_unlock(dev);
 
 	return ret;
 }
@@ -303,7 +306,7 @@ static struct rpmsg_endpoint *ethosu_create_ept(struct rpmsg_device *rpdev)
 	dev_info(dev, "Creating rpmsg endpoint. name=%s, src=%u, dst=%u",
 		 info.name, info.src, info.dst);
 
-	ept = rpmsg_create_ept(rpdev, ethosu_handle_msg, NULL, info);
+	ept = rpmsg_create_ept(rpdev, ethosu_handle_rpmsg, NULL, info);
 	if (!ept) {
 		dev_err(&rpdev->dev, "Failed to create endpoint");
 
@@ -313,54 +316,65 @@ static struct rpmsg_endpoint *ethosu_create_ept(struct rpmsg_device *rpdev)
 	return ept;
 }
 
-int ethosu_dev_init(struct ethosu_device *edev,
-		    struct rpmsg_device *rpdev,
+static const struct file_operations fops = {
+	.owner          = THIS_MODULE,
+	.open           = &ethosu_open,
+	.unlocked_ioctl = &ethosu_ioctl,
+#ifdef CONFIG_COMPAT
+	.compat_ioctl   = &ethosu_ioctl,
+#endif
+};
+
+int ethosu_dev_init(struct rpmsg_device *rpdev,
 		    struct class *class,
 		    dev_t devt)
 {
-	static const struct file_operations fops = {
-		.owner          = THIS_MODULE,
-		.open           = &ethosu_open,
-		.unlocked_ioctl = &ethosu_ioctl,
-#ifdef CONFIG_COMPAT
-		.compat_ioctl   = &ethosu_ioctl,
-#endif
-	};
+	struct device *dev = &rpdev->dev;
+	struct ethosu_device *edev;
 	struct device *sysdev;
 	int ret;
 
-	edev->dev = &rpdev->dev;
+	dev_info(dev, "%s", __FUNCTION__);
+
+	/* Allocate and create Ethos-U device */
+	edev = devm_kzalloc(dev, sizeof(*edev), GFP_KERNEL);
+	if (!edev)
+		return -ENOMEM;
+
+	dev_set_drvdata(dev, edev);
+
 	edev->rpdev = rpdev;
 	edev->class = class;
 	edev->devt = devt;
-	mutex_init(&edev->mutex);
 
+	/* Create RPMsg endpoint */
 	edev->ept = ethosu_create_ept(rpdev);
 	if (IS_ERR(edev->ept))
 		return PTR_ERR(edev->ept);
 
-	ret = ethosu_mailbox_init(&edev->mailbox, edev->dev, edev->ept);
+	ret = ethosu_mailbox_init(&edev->mailbox, dev, edev->ept);
 	if (ret)
 		return ret;
 
+	/* Create device node */
 	cdev_init(&edev->cdev, &fops);
 	edev->cdev.owner = THIS_MODULE;
 
 	ret = cdev_add(&edev->cdev, edev->devt, 1);
 	if (ret) {
-		dev_err(edev->dev, "Failed to add character device.");
+		dev_err(dev, "Failed to add character device.");
 		goto deinit_mailbox;
 	}
 
-	sysdev = device_create(edev->class, NULL, edev->devt, edev,
+	sysdev = device_create(edev->class, NULL, edev->devt, rpdev,
 			       "ethosu%d", MINOR(edev->devt));
 	if (IS_ERR(sysdev)) {
-		dev_err(edev->dev, "Failed to create device.");
+		dev_err(dev, "Failed to create device.");
 		ret = PTR_ERR(sysdev);
 		goto del_cdev;
 	}
 
-	dev_info(edev->dev,
+	dev_info(dev,
 		 "Created Arm Ethos-U device. name=%s, major=%d, minor=%d",
 		 dev_name(sysdev), MAJOR(edev->devt), MINOR(edev->devt));
 
@@ -377,9 +391,12 @@ deinit_mailbox:
 	return ret;
 }
 
-void ethosu_dev_deinit(struct ethosu_device *edev)
+void ethosu_dev_deinit(struct rpmsg_device *rpdev)
 {
-	dev_info(edev->dev, "%s\n", __FUNCTION__);
+	struct device *dev = &rpdev->dev;
+	struct ethosu_device *edev = dev_get_drvdata(dev);
+
+	dev_info(dev, "%s", __FUNCTION__);
 
 	ethosu_mailbox_deinit(&edev->mailbox);
 	rpmsg_destroy_ept(edev->ept);
diff --git a/kernel/ethosu_device.h b/kernel/ethosu_device.h
index cc4271d..d02f29d 100644
--- a/kernel/ethosu_device.h
+++ b/kernel/ethosu_device.h
@@ -42,13 +42,11 @@
  * struct ethosu_device - Device structure
  */
 struct ethosu_device {
-	struct device         *dev;
 	struct rpmsg_device   *rpdev;
 	struct rpmsg_endpoint *ept;
 	struct cdev           cdev;
 	struct                class *class;
 	dev_t                 devt;
-	struct mutex          mutex;
 	struct ethosu_mailbox mailbox;
 };
 
@@ -61,14 +59,13 @@ struct ethosu_device {
  *
  * Return: 0 on success, else error code.
  */
-int ethosu_dev_init(struct ethosu_device *edev,
-		    struct rpmsg_device *rpdev,
+int ethosu_dev_init(struct rpmsg_device *rpdev,
 		    struct class *class,
 		    dev_t devt);
 
 /**
  * ethosu_dev_deinit() - Initialize the device
  */
-void ethosu_dev_deinit(struct ethosu_device *edev);
+void ethosu_dev_deinit(struct rpmsg_device *rpdev);
 
 #endif /* ETHOSU_DEVICE_H */
diff --git a/kernel/ethosu_driver.c b/kernel/ethosu_driver.c
index c6fc8cd..6e5bfb9 100644
--- a/kernel/ethosu_driver.c
+++ b/kernel/ethosu_driver.c
@@ -57,7 +57,6 @@ static DECLARE_BITMAP(minors, MINOR_COUNT);
 static int ethosu_rpmsg_probe(struct rpmsg_device *rpdev)
 {
 	struct device *dev = &rpdev->dev;
-	struct ethosu_device *edev;
 	int minor;
 	int ret;
 
@@ -69,15 +68,8 @@ static int ethosu_rpmsg_probe(struct rpmsg_device *rpdev)
 		return -ENOMEM;
 	}
 
-	edev = devm_kzalloc(dev, sizeof(*edev), GFP_KERNEL);
-	if (!edev)
-		return -ENOMEM;
-
-	dev_set_drvdata(dev, edev);
-
 	/* Initialize device */
-	ret = ethosu_dev_init(edev, rpdev, ethosu_class,
-			      MKDEV(MAJOR(devt), minor));
+	ret = ethosu_dev_init(rpdev, ethosu_class, MKDEV(MAJOR(devt), minor));
 	if (ret)
 		return ret;
 
@@ -88,12 +80,13 @@ static int ethosu_rpmsg_probe(struct rpmsg_device *rpdev)
 
 static void ethosu_rpmsg_remove(struct rpmsg_device *rpdev)
 {
-	struct ethosu_device *edev = dev_get_drvdata(&rpdev->dev);
+	struct device *dev = &rpdev->dev;
+	struct ethosu_device *edev = dev_get_drvdata(dev);
 
-	dev_info(&rpdev->dev, "%s", __FUNCTION__);
+	dev_info(dev, "%s", __FUNCTION__);
 
 	clear_bit(MINOR(edev->devt), minors);
-	ethosu_dev_deinit(edev);
+	ethosu_dev_deinit(rpdev);
 }
 
 static int ethosu_rpmsg_cb(struct rpmsg_device *rpdev,
diff --git a/kernel/ethosu_inference.c b/kernel/ethosu_inference.c
index 4f56126..6befd3a 100644
--- a/kernel/ethosu_inference.c
+++ b/kernel/ethosu_inference.c
@@ -91,11 +91,12 @@ static const char *status_to_string(const enum ethosu_uapi_status status)
 
 static int ethosu_inference_send(struct ethosu_inference *inf)
 {
+	struct device *dev = inf->dev;
 	int ret;
 
 	inf->status = ETHOSU_UAPI_STATUS_ERROR;
 
-	ret = ethosu_mailbox_inference(&inf->edev->mailbox, &inf->msg,
+	ret = ethosu_mailbox_inference(inf->mailbox, &inf->msg,
 				       inf->ifm_count, inf->ifm,
 				       inf->ofm_count, inf->ofm,
 				       inf->net->buf,
@@ -104,7 +105,7 @@ static int ethosu_inference_send(struct ethosu_inference *inf)
 				       ETHOSU_PMU_EVENT_MAX,
 				       inf->pmu_cycle_counter_enable);
 	if (ret) {
-		dev_warn(inf->edev->dev,
+		dev_warn(dev,
 			 "Failed to send inference request. inf=0x%pK, ret=%d",
 			 inf, ret);
 
@@ -150,12 +151,13 @@ static void ethosu_inference_kref_destroy(struct kref *kref)
 {
 	struct ethosu_inference *inf =
 		container_of(kref, struct ethosu_inference, kref);
+	struct device *dev = inf->dev;
 
-	dev_info(inf->edev->dev,
+	dev_info(dev,
 		 "Inference destroy. inf=0x%pK, status=%d, ifm_count=%u, ofm_count=%u",
 		 inf, inf->status, inf->ifm_count, inf->ofm_count);
 
-	ethosu_mailbox_deregister(&inf->edev->mailbox, &inf->msg);
+	ethosu_mailbox_deregister(inf->mailbox, &inf->msg);
 
 	while (inf->ifm_count-- > 0)
 		ethosu_buffer_put(inf->ifm[inf->ifm_count]);
@@ -164,15 +166,16 @@ static void ethosu_inference_kref_destroy(struct kref *kref)
 		ethosu_buffer_put(inf->ofm[inf->ofm_count]);
 
 	ethosu_network_put(inf->net);
-	devm_kfree(inf->edev->dev, inf);
+	devm_kfree(dev, inf);
 }
 
 static int ethosu_inference_release(struct inode *inode,
 				    struct file *file)
 {
 	struct ethosu_inference *inf = file->private_data;
+	struct device *dev = inf->dev;
 
-	dev_info(inf->edev->dev,
+	dev_info(dev,
 		 "Inference release. file=0x%pK, inf=0x%pK",
 		 file, inf);
 
@@ -200,14 +203,15 @@ static long ethosu_inference_ioctl(struct file *file,
 				   unsigned long arg)
 {
 	struct ethosu_inference *inf = file->private_data;
+	struct device *dev = inf->dev;
 	void __user *udata = (void __user *)arg;
 	int ret;
 
-	ret = mutex_lock_interruptible(&inf->edev->mutex);
+	ret = device_lock_interruptible(dev);
 	if (ret)
 		return ret;
 
-	dev_info(inf->edev->dev,
+	dev_info(dev,
 		 "Inference ioctl: file=0x%pK, inf=0x%pK, cmd=0x%x, arg=%lu",
 		 file, inf, cmd, arg);
 
@@ -228,7 +232,7 @@ static long ethosu_inference_ioctl(struct file *file,
 		uapi.pmu_config.cycle_count = inf->pmu_cycle_counter_enable;
 		uapi.pmu_count.cycle_count = inf->pmu_cycle_counter_count;
 
-		dev_info(inf->edev->dev,
+		dev_info(dev,
 			 "Inference ioctl: Inference status. status=%s (%d)\n",
 			 status_to_string(uapi.status), uapi.status);
 
@@ -239,11 +243,12 @@ static long ethosu_inference_ioctl(struct file *file,
 	case ETHOSU_IOCTL_INFERENCE_CANCEL: {
 		struct ethosu_uapi_cancel_inference_status uapi;
 
-		dev_info(inf->edev->dev,
+		dev_info(dev,
 			 "Inference ioctl: Cancel Inference. Handle=%p\n",
 			 inf);
 
-		ret = ethosu_cancel_inference_request(inf, &uapi);
+		ret = ethosu_cancel_inference_request(dev, inf->mailbox, inf,
+						      &uapi);
 		if (ret)
 			break;
 
@@ -252,18 +257,19 @@ static long ethosu_inference_ioctl(struct file *file,
 		break;
 	}
 	default: {
-		dev_err(inf->edev->dev, "Invalid ioctl. cmd=%u, arg=%lu\n",
+		dev_err(dev, "Invalid ioctl. cmd=%u, arg=%lu\n",
 			cmd, arg);
 		break;
 	}
 	}
 
-	mutex_unlock(&inf->edev->mutex);
+	device_unlock(dev);
 
 	return ret;
 }
 
-int ethosu_inference_create(struct ethosu_device *edev,
+int ethosu_inference_create(struct device *dev,
+			    struct ethosu_mailbox *mailbox,
 			    struct ethosu_network *net,
 			    struct ethosu_uapi_inference_create *uapi)
 {
@@ -274,18 +280,19 @@ int ethosu_inference_create(struct ethosu_device *edev,
 
 	if (uapi->ifm_count > ETHOSU_FD_MAX ||
 	    uapi->ofm_count > ETHOSU_FD_MAX) {
-		dev_warn(edev->dev,
+		dev_warn(dev,
 			 "Too many IFM and/or OFM buffers for inference. ifm_count=%u, ofm_count=%u",
 			 uapi->ifm_count, uapi->ofm_count);
 
 		return -EFAULT;
 	}
 
-	inf = devm_kzalloc(edev->dev, sizeof(*inf), GFP_KERNEL);
+	inf = devm_kzalloc(dev, sizeof(*inf), GFP_KERNEL);
 	if (!inf)
 		return -ENOMEM;
 
-	inf->edev = edev;
+	inf->dev = dev;
+	inf->mailbox = mailbox;
 	inf->net = net;
 	inf->done = false;
 	inf->status = ETHOSU_UAPI_STATUS_ERROR;
@@ -294,7 +301,7 @@ int ethosu_inference_create(struct ethosu_device *edev,
 	inf->msg.fail = ethosu_inference_fail;
 
 	/* Add inference to pending list */
-	ret = ethosu_mailbox_register(&edev->mailbox, &inf->msg);
+	ret = ethosu_mailbox_register(mailbox, &inf->msg);
 	if (ret < 0)
 		goto kfree;
 
@@ -321,7 +328,7 @@ int ethosu_inference_create(struct ethosu_device *edev,
 	}
 
 	/* Configure PMU and cycle counter */
-	dev_info(inf->edev->dev,
+	dev_info(dev,
 		 "Configuring events for PMU. events=[%u, %u, %u, %u]\n",
 		 uapi->pmu_config.events[0], uapi->pmu_config.events[1],
 		 uapi->pmu_config.events[2], uapi->pmu_config.events[3]);
@@ -333,7 +340,7 @@ int ethosu_inference_create(struct ethosu_device *edev,
 	}
 
 	if (uapi->pmu_config.cycle_count)
-		dev_info(inf->edev->dev, "Enabling cycle counter\n");
+		dev_info(dev, "Enabling cycle counter\n");
 
 	/* Configure cycle counter and reset any previous count */
 	inf->pmu_cycle_counter_enable = uapi->pmu_config.cycle_count;
@@ -357,7 +364,7 @@ int ethosu_inference_create(struct ethosu_device *edev,
 	inf->file = fget(ret);
 	fput(inf->file);
 
-	dev_info(edev->dev,
+	dev_info(dev,
 		 "Inference create. file=0x%pK, fd=%d, inf=0x%p, net=0x%pK, msg.id=0x%x",
 		 inf->file, fd, inf, inf->net, inf->msg.id);
 
@@ -375,7 +382,7 @@ put_ifm:
 		ethosu_buffer_put(inf->ifm[inf->ifm_count]);
 
 kfree:
-	devm_kfree(edev->dev, inf);
+	devm_kfree(dev, inf);
 
 	return ret;
 }
@@ -412,18 +419,19 @@ int ethosu_inference_put(struct ethosu_inference *inf)
 	return kref_put(&inf->kref, &ethosu_inference_kref_destroy);
 }
 
-void ethosu_inference_rsp(struct ethosu_device *edev,
+void ethosu_inference_rsp(struct ethosu_mailbox *mailbox,
 			  int msg_id,
 			  struct ethosu_core_msg_inference_rsp *rsp)
 {
+	struct device *dev = mailbox->dev;
 	struct ethosu_mailbox_msg *msg;
 	struct ethosu_inference *inf;
 	int ret;
 	int i;
 
-	msg = ethosu_mailbox_find(&edev->mailbox, msg_id);
+	msg = ethosu_mailbox_find(mailbox, msg_id);
 	if (IS_ERR(msg)) {
-		dev_warn(edev->dev,
+		dev_warn(dev,
 			 "Id for inference msg not found. Id=%d\n",
 			 msg_id);
 
@@ -464,14 +472,14 @@ void ethosu_inference_rsp(struct ethosu_device *edev,
 		inf->pmu_cycle_counter_enable = rsp->pmu_cycle_counter_enable;
 		inf->pmu_cycle_counter_count = rsp->pmu_cycle_counter_count;
 
-		dev_info(edev->dev,
+		dev_info(dev,
 			 "PMU events. config=[%u, %u, %u, %u], count=[%u, %u, %u, %u]\n",
 			 inf->pmu_event_config[0], inf->pmu_event_config[1],
 			 inf->pmu_event_config[2], inf->pmu_event_config[3],
 			 inf->pmu_event_count[0], inf->pmu_event_count[1],
 			 inf->pmu_event_count[2], inf->pmu_event_count[3]);
 
-		dev_info(edev->dev,
+		dev_info(dev,
 			 "PMU cycle counter. enable=%u, count=%llu\n",
 			 inf->pmu_cycle_counter_enable,
 			 inf->pmu_cycle_counter_count);
diff --git a/kernel/ethosu_inference.h b/kernel/ethosu_inference.h
index 5a3b860..9fa0a4a 100644
--- a/kernel/ethosu_inference.h
+++ b/kernel/ethosu_inference.h
@@ -37,7 +37,6 @@
 
 struct ethosu_buffer;
 struct ethosu_core_msg_inference_rsp;
-struct ethosu_device;
 struct ethosu_network;
 struct ethosu_uapi_inference_create;
 struct file;
@@ -60,7 +59,8 @@ struct file;
  * @msg:			Mailbox message
  */
 struct ethosu_inference {
-	struct ethosu_device      *edev;
+	struct device             *dev;
+	struct ethosu_mailbox     *mailbox;
 	struct file               *file;
 	struct kref               kref;
 	wait_queue_head_t         waitq;
@@ -89,7 +89,8 @@ struct ethosu_inference {
  *
  * Return: fd on success, else error code.
  */
-int ethosu_inference_create(struct ethosu_device *edev,
+int ethosu_inference_create(struct device *dev,
+			    struct ethosu_mailbox *mailbox,
 			    struct ethosu_network *net,
 			    struct ethosu_uapi_inference_create *uapi);
 
@@ -117,7 +118,7 @@ int ethosu_inference_put(struct ethosu_inference *inf);
 /**
  * ethosu_inference_rsp() - Handle inference response
  */
-void ethosu_inference_rsp(struct ethosu_device *edev,
+void ethosu_inference_rsp(struct ethosu_mailbox *mailbox,
 			  int msg_id,
 			  struct ethosu_core_msg_inference_rsp *rsp);
 
diff --git a/kernel/ethosu_network.c b/kernel/ethosu_network.c
index 251b181..d64d6d7 100644
--- a/kernel/ethosu_network.c
+++ b/kernel/ethosu_network.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020,2022 Arm Limited.
+ * Copyright 2020,2022-2023 Arm Limited and/or its affiliates
  *
  * This program is free software and is provided to you under the terms of the
  * GNU General Public License version 2 as published by the Free Software
@@ -67,21 +67,23 @@ static void ethosu_network_destroy(struct kref *kref)
 {
 	struct ethosu_network *net =
 		container_of(kref, struct ethosu_network, kref);
+	struct device *dev = net->dev;
 
-	dev_info(net->edev->dev, "Network destroy. net=0x%pK\n", net);
+	dev_info(dev, "Network destroy. net=0x%pK\n", net);
 
 	if (net->buf != NULL)
 		ethosu_buffer_put(net->buf);
 
-	devm_kfree(net->edev->dev, net);
+	devm_kfree(dev, net);
 }
 
 static int ethosu_network_release(struct inode *inode,
 				  struct file *file)
 {
 	struct ethosu_network *net = file->private_data;
+	struct device *dev = net->dev;
 
-	dev_info(net->edev->dev, "Network release. file=0x%pK, net=0x%pK\n",
+	dev_info(dev, "Network release. file=0x%pK, net=0x%pK\n",
 		 file, net);
 
 	ethosu_network_put(net);
@@ -94,14 +96,15 @@ static long ethosu_network_ioctl(struct file *file,
 				 unsigned long arg)
 {
 	struct ethosu_network *net = file->private_data;
+	struct device *dev = net->dev;
 	void __user *udata = (void __user *)arg;
 	int ret = -EINVAL;
 
-	ret = mutex_lock_interruptible(&net->edev->mutex);
+	ret = device_lock_interruptible(net->dev);
 	if (ret)
 		return ret;
 
-	dev_info(net->edev->dev,
+	dev_info(dev,
 		 "Network ioctl: file=0x%pK, net=0x%pK, cmd=0x%x, arg=0x%lx\n",
 		 file, net, cmd, arg);
 
@@ -112,11 +115,13 @@ static long ethosu_network_ioctl(struct file *file,
 		if (copy_from_user(&uapi, udata, sizeof(uapi)))
 			break;
 
-		dev_info(net->edev->dev,
+		dev_info(dev,
 			 "Network ioctl: Network info. net=0x%pK\n",
 			 net);
 
-		ret = ethosu_network_info_request(net, &uapi);
+		ret =
+			ethosu_network_info_request(dev, net->mailbox, net,
+						    &uapi);
 		if (ret)
 			break;
 
@@ -129,36 +134,38 @@ static long ethosu_network_ioctl(struct file *file,
 		if (copy_from_user(&uapi, udata, sizeof(uapi)))
 			break;
 
-		dev_info(net->edev->dev,
+		dev_info(dev,
 			 "Network ioctl: Inference. ifm_fd=%u, ofm_fd=%u\n",
 			 uapi.ifm_fd[0], uapi.ofm_fd[0]);
 
-		ret = ethosu_inference_create(net->edev, net, &uapi);
+		ret = ethosu_inference_create(dev, net->mailbox, net, &uapi);
 		break;
 	}
 	default: {
-		dev_err(net->edev->dev, "Invalid ioctl. cmd=%u, arg=%lu",
+		dev_err(dev, "Invalid ioctl. cmd=%u, arg=%lu",
 			cmd, arg);
 		break;
 	}
 	}
 
-	mutex_unlock(&net->edev->mutex);
+	device_unlock(net->dev);
 
 	return ret;
 }
 
-int ethosu_network_create(struct ethosu_device *edev,
+int ethosu_network_create(struct device *dev,
+			  struct ethosu_mailbox *mailbox,
 			  struct ethosu_uapi_network_create *uapi)
 {
 	struct ethosu_network *net;
 	int ret = -ENOMEM;
 
-	net = devm_kzalloc(edev->dev, sizeof(*net), GFP_KERNEL);
+	net = devm_kzalloc(dev, sizeof(*net), GFP_KERNEL);
 	if (!net)
 		return -ENOMEM;
 
-	net->edev = edev;
+	net->dev = dev;
+	net->mailbox = mailbox;
 	net->buf = NULL;
 	kref_init(&net->kref);
 
@@ -180,7 +187,7 @@ int ethosu_network_create(struct ethosu_device *edev,
 	net->file = fget(ret);
 	fput(net->file);
 
-	dev_info(edev->dev,
+	dev_info(dev,
 		 "Network create. file=0x%pK, fd=%d, net=0x%pK, buf=0x%pK, index=%u",
 		 net->file, ret, net, net->buf, net->index);
 
@@ -191,7 +198,7 @@ put_buf:
 		ethosu_buffer_put(net->buf);
 
 free_net:
-	devm_kfree(edev->dev, net);
+	devm_kfree(dev, net);
 
 	return ret;
 }
diff --git a/kernel/ethosu_network.h b/kernel/ethosu_network.h
index 8ee6818..e0d41b1 100644
--- a/kernel/ethosu_network.h
+++ b/kernel/ethosu_network.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020,2022 Arm Limited.
+ * Copyright 2020,2022-2023 Arm Limited and/or its affiliates
  *
  * This program is free software and is provided to you under the terms of the
  * GNU General Public License version 2 as published by the Free Software
@@ -33,17 +33,17 @@
  ****************************************************************************/
 
 struct ethosu_buffer;
-struct ethosu_device;
 struct ethosu_uapi_network_create;
 struct device;
 struct file;
 
 struct ethosu_network {
-	struct ethosu_device *edev;
-	struct file          *file;
-	struct kref          kref;
-	struct ethosu_buffer *buf;
-	uint32_t             index;
+	struct device         *dev;
+	struct ethosu_mailbox *mailbox;
+	struct file           *file;
+	struct kref           kref;
+	struct ethosu_buffer  *buf;
+	uint32_t              index;
 };
 
 /****************************************************************************
@@ -57,7 +57,8 @@ struct ethosu_network {
  *
  * Return: fd on success, else error code.
  */
-int ethosu_network_create(struct ethosu_device *edev,
+int ethosu_network_create(struct device *dev,
+			  struct ethosu_mailbox *mailbox,
 			  struct ethosu_uapi_network_create *uapi);
 
 /**
diff --git a/kernel/ethosu_network_info.c b/kernel/ethosu_network_info.c
index 5e7a1b9..0e205db 100644
--- a/kernel/ethosu_network_info.c
+++ b/kernel/ethosu_network_info.c
@@ -31,10 +31,11 @@
 
 #define NETWORK_INFO_RESP_TIMEOUT_MS 3000
 
-static inline int ethosu_network_info_send(struct ethosu_network_info *info)
+static inline int ethosu_network_info_send(struct ethosu_network_info *info,
+					   struct ethosu_mailbox *mailbox)
 {
 	/* Send network info request to firmware */
-	return ethosu_mailbox_network_info_request(&info->edev->mailbox,
+	return ethosu_mailbox_network_info_request(mailbox,
 						   &info->msg,
 						   info->net->buf,
 						   info->net->index);
@@ -52,47 +53,49 @@ static void ethosu_network_info_fail(struct ethosu_mailbox_msg *msg)
 	complete(&info->done);
 }
 
-int ethosu_network_info_request(struct ethosu_network *net,
+int ethosu_network_info_request(struct device *dev,
+				struct ethosu_mailbox *mailbox,
+				struct ethosu_network *net,
 				struct ethosu_uapi_network_info *uapi)
 {
 	struct ethosu_network_info *info;
 	int ret;
 	int timeout;
 
-	info = devm_kzalloc(net->edev->dev, sizeof(*info), GFP_KERNEL);
+	info = devm_kzalloc(dev, sizeof(*info), GFP_KERNEL);
 	if (!info)
 		return -ENOMEM;
 
-	info->edev = net->edev;
+	info->dev = dev;
 	info->net = net;
 	info->uapi = uapi;
 	init_completion(&info->done);
 	info->msg.fail = ethosu_network_info_fail;
 
-	ret = ethosu_mailbox_register(&info->edev->mailbox, &info->msg);
+	ret = ethosu_mailbox_register(mailbox, &info->msg);
 	if (ret < 0)
 		goto kfree;
 
 	/* Get reference to network */
 	ethosu_network_get(info->net);
 
-	ret = ethosu_network_info_send(info);
+	ret = ethosu_network_info_send(info, mailbox);
 	if (ret)
 		goto deregister;
 
-	dev_info(info->edev->dev,
+	dev_info(dev,
 		 "Network info create. info=0x%pK, net=0x%pK, msg.id=0x%x\n",
 		 info, info->net, info->msg.id);
 
 	/* Unlock the device mutex and wait for completion */
-	mutex_unlock(&info->edev->mutex);
+	device_unlock(dev);
 	timeout = wait_for_completion_timeout(&info->done,
 					      msecs_to_jiffies(
 						      NETWORK_INFO_RESP_TIMEOUT_MS));
-	mutex_lock(&info->edev->mutex);
+	device_lock(dev);
 
 	if (0 == timeout) {
-		dev_warn(info->edev->dev, "Network info timed out. info=0x%pK",
+		dev_warn(dev, "Network info timed out. info=0x%pK",
 			 info);
 
 		ret = -ETIME;
@@ -102,30 +105,31 @@ int ethosu_network_info_request(struct ethosu_network *net,
 	ret = info->errno;
 
 deregister:
-	ethosu_mailbox_deregister(&info->edev->mailbox, &info->msg);
+	ethosu_mailbox_deregister(mailbox, &info->msg);
 	ethosu_network_put(info->net);
 
 kfree:
-	dev_info(info->edev->dev,
+	dev_info(dev,
 		 "Network info destroy. info=0x%pK, msg.id=0x%x\n",
 		 info, info->msg.id);
-	devm_kfree(info->edev->dev, info);
+	devm_kfree(dev, info);
 
 	return ret;
 }
 
-void ethosu_network_info_rsp(struct ethosu_device *edev,
+void ethosu_network_info_rsp(struct ethosu_mailbox *mailbox,
 			     int msg_id,
 			     struct ethosu_core_msg_network_info_rsp *rsp)
 {
 	int ret;
+	struct device *dev = mailbox->dev;
 	struct ethosu_mailbox_msg *msg;
 	struct ethosu_network_info *info;
 	uint32_t i;
 
-	msg = ethosu_mailbox_find(&edev->mailbox, msg_id);
+	msg = ethosu_mailbox_find(mailbox, msg_id);
 	if (IS_ERR(msg)) {
-		dev_warn(edev->dev,
+		dev_warn(dev,
 			 "Id for network info msg not found. msg.id=0x%x\n",
 			 msg_id);
 
diff --git a/kernel/ethosu_network_info.h b/kernel/ethosu_network_info.h
index a809d35..8c2e659 100644
--- a/kernel/ethosu_network_info.h
+++ b/kernel/ethosu_network_info.h
@@ -35,12 +35,11 @@
  ****************************************************************************/
 
 struct ethosu_core_msg_network_info_rsp;
-struct ethosu_device;
 struct ethosu_network;
 struct ethosu_uapi_network_info;
 
 struct ethosu_network_info {
-	struct ethosu_device            *edev;
+	struct device                   *dev;
 	struct ethosu_network           *net;
 	struct ethosu_uapi_network_info *uapi;
 	struct completion               done;
@@ -59,13 +58,15 @@ struct ethosu_network_info {
  *
  * Return: 0 on success, .
  */
-int ethosu_network_info_request(struct ethosu_network *net,
+int ethosu_network_info_request(struct device *dev,
+				struct ethosu_mailbox *mailbox,
+				struct ethosu_network *net,
 				struct ethosu_uapi_network_info *uapi);
 
 /**
  * ethosu_network_info_rsp() - Handle network info response.
  */
-void ethosu_network_info_rsp(struct ethosu_device *edev,
+void ethosu_network_info_rsp(struct ethosu_mailbox *mailbox,
 			     int msg_id,
 			     struct ethosu_core_msg_network_info_rsp *rsp);