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 @@
 		 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 @@
 		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 @@
 			 "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 @@
 		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 @@
 			 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 @@
 			 "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_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 @@
 			 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 @@
 		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 @@
 		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 @@
 		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 @@
 	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 @@
 	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 @@
 	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);