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 @@
 {
 	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 @@
 				 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 @@
 		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 @@
 		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 @@
 	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 @@
 		ethosu_buffer_put(net->buf);
 
 free_net:
-	devm_kfree(edev->dev, net);
+	devm_kfree(dev, net);
 
 	return ret;
 }