13 files changed, 418 insertions, 374 deletions

diff --git a/kernel/Kbuild b/kernel/Kbuild
index d54ec2c..d4610bc 100644
--- a/kernel/Kbuild
+++ b/kernel/Kbuild
@@ -28,4 +28,5 @@ ethosu-objs := ethosu_driver.o \
                ethosu_network.o \
                ethosu_network_info.o \
                ethosu_watchdog.o \
+               ethosu_capabilities.o \
                ethosu_cancel_inference.o
diff --git a/kernel/ethosu_cancel_inference.c b/kernel/ethosu_cancel_inference.c
index befdd2f..6d93cf1 100644
--- a/kernel/ethosu_cancel_inference.c
+++ b/kernel/ethosu_cancel_inference.c
@@ -40,24 +40,12 @@
  * Functions
  ****************************************************************************/
 
-static void ethosu_cancel_inference_destroy(struct kref *kref)
-{
-	struct ethosu_cancel_inference *cancellation =
-		container_of(kref, struct ethosu_cancel_inference, kref);
-
-	dev_info(cancellation->edev->dev,
-		 "Cancel inference destroy. handle=0x%p\n", cancellation);
-	list_del(&cancellation->msg.list);
-	/* decrease the reference on the inference we are refering to */
-	ethosu_inference_put(cancellation->inf);
-	devm_kfree(cancellation->edev->dev, cancellation);
-}
-
 static int ethosu_cancel_inference_send(
 	struct ethosu_cancel_inference *cancellation)
 {
 	return ethosu_mailbox_cancel_inference(&cancellation->edev->mailbox,
-					       cancellation, cancellation->inf);
+					       &cancellation->msg,
+					       cancellation->inf->msg.id);
 }
 
 static void ethosu_cancel_inference_fail(struct ethosu_mailbox_msg *msg)
@@ -120,7 +108,6 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	cancellation->edev = inf->edev;
 	cancellation->inf = inf;
 	cancellation->uapi = uapi;
-	kref_init(&cancellation->kref);
 	init_completion(&cancellation->done);
 	cancellation->msg.fail = ethosu_cancel_inference_fail;
 
@@ -128,13 +115,18 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	 * whole firmware and marked the inference as aborted */
 	cancellation->msg.resend = ethosu_cancel_inference_complete;
 
-	/* Add cancel inference to pending list */
-	list_add(&cancellation->msg.list,
-		 &cancellation->edev->mailbox.pending_list);
+	ret = ethosu_mailbox_register(&cancellation->edev->mailbox,
+				      &cancellation->msg);
+	if (ret < 0)
+		goto kfree;
+
+	dev_info(cancellation->edev->dev,
+		 "Inference cancellation create. Id=%d, handle=0x%p\n",
+		 cancellation->msg.id, cancellation);
 
 	ret = ethosu_cancel_inference_send(cancellation);
 	if (0 != ret)
-		goto put_kref;
+		goto deregister;
 
 	/* Unlock the mutex before going to block on the condition */
 	mutex_unlock(&cancellation->edev->mutex);
@@ -145,18 +137,18 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	/* take back the mutex before resuming to do anything */
 	ret = mutex_lock_interruptible(&cancellation->edev->mutex);
 	if (0 != ret)
-		goto put_kref;
+		goto deregister;
 
 	if (0 == timeout /* timed out*/) {
 		dev_warn(inf->edev->dev,
 			 "Msg: Cancel Inference response lost - timeout\n");
 		ret = -EIO;
-		goto put_kref;
+		goto deregister;
 	}
 
 	if (cancellation->errno) {
 		ret = cancellation->errno;
-		goto put_kref;
+		goto deregister;
 	}
 
 	/* if cancellation failed and the inference did not complete then reset
@@ -165,11 +157,19 @@ int ethosu_cancel_inference_request(struct ethosu_inference *inf,
 	    !cancellation->inf->done) {
 		ret = ethosu_firmware_reset(cancellation->edev);
 		if (ret)
-			goto put_kref;
+			goto deregister;
 	}
 
-put_kref:
-	kref_put(&cancellation->kref, &ethosu_cancel_inference_destroy);
+deregister:
+	ethosu_mailbox_deregister(&cancellation->edev->mailbox,
+				  &cancellation->msg);
+
+kfree:
+	dev_info(cancellation->edev->dev,
+		 "Cancel inference destroy. handle=0x%p\n", cancellation);
+	/* decrease the reference on the inference we are refering to */
+	ethosu_inference_put(cancellation->inf);
+	devm_kfree(cancellation->edev->dev, cancellation);
 
 	return ret;
 }
@@ -177,12 +177,12 @@ put_kref:
 void ethosu_cancel_inference_rsp(struct ethosu_device *edev,
 				 struct ethosu_core_cancel_inference_rsp *rsp)
 {
-	struct ethosu_cancel_inference *cancellation =
-		(struct ethosu_cancel_inference *)rsp->user_arg;
-	int ret;
+	int id = (int)rsp->user_arg;
+	struct ethosu_mailbox_msg *msg;
+	struct ethosu_cancel_inference *cancellation;
 
-	ret = ethosu_mailbox_find(&edev->mailbox, &cancellation->msg);
-	if (ret) {
+	msg = ethosu_mailbox_find(&edev->mailbox, id);
+	if (IS_ERR(msg)) {
 		dev_warn(edev->dev,
 			 "Handle not found in cancel inference list. handle=0x%p\n",
 			 rsp);
@@ -190,6 +190,8 @@ void ethosu_cancel_inference_rsp(struct ethosu_device *edev,
 		return;
 	}
 
+	cancellation = container_of(msg, typeof(*cancellation), msg);
+
 	if (completion_done(&cancellation->done))
 		return;
 
diff --git a/kernel/ethosu_cancel_inference.h b/kernel/ethosu_cancel_inference.h
index 94d9fe1..044692c 100644
--- a/kernel/ethosu_cancel_inference.h
+++ b/kernel/ethosu_cancel_inference.h
@@ -28,7 +28,6 @@
 #include "ethosu_mailbox.h"
 #include "uapi/ethosu.h"
 
-#include <linux/kref.h>
 #include <linux/types.h>
 #include <linux/completion.h>
 
@@ -45,7 +44,6 @@ struct ethosu_cancel_inference {
 	struct ethosu_device                       *edev;
 	struct ethosu_inference                    *inf;
 	struct ethosu_uapi_cancel_inference_status *uapi;
-	struct kref                                kref;
 	struct completion                          done;
 	struct ethosu_mailbox_msg                  msg;
 	int                                        errno;
diff --git a/kernel/ethosu_capabilities.c b/kernel/ethosu_capabilities.c
new file mode 100644
index 0000000..d5f77f2
--- /dev/null
+++ b/kernel/ethosu_capabilities.c
@@ -0,0 +1,173 @@
+/*
+ * Copyright (c) 2022 Arm Limited.
+ *
+ * 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
+ * Foundation, and any use by you of this program is subject to the terms
+ * of such GNU licence.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, you can access it online at
+ * http://www.gnu.org/licenses/gpl-2.0.html.
+ *
+ * SPDX-License-Identifier: GPL-2.0-only
+ */
+
+/****************************************************************************
+ * Includes
+ ****************************************************************************/
+
+#include "ethosu_capabilities.h"
+
+#include "ethosu_device.h"
+
+#include <linux/delay.h>
+#include <linux/dma-mapping.h>
+#include <linux/errno.h>
+
+/****************************************************************************
+ * Defines
+ ****************************************************************************/
+
+#define CAPABILITIES_RESP_TIMEOUT_MS 2000
+
+/****************************************************************************
+ * Functions
+ ****************************************************************************/
+
+static inline int ethosu_capabilities_send(struct ethosu_capabilities *cap)
+{
+	return ethosu_mailbox_capabilities_request(&cap->edev->mailbox,
+						   &cap->msg);
+}
+
+static void ethosu_capabilities_fail(struct ethosu_mailbox_msg *msg)
+{
+	struct ethosu_capabilities *cap =
+		container_of(msg, typeof(*cap), msg);
+
+	if (completion_done(&cap->done))
+		return;
+
+	cap->errno = -EFAULT;
+	complete(&cap->done);
+}
+
+static int ethosu_capabilities_resend(struct ethosu_mailbox_msg *msg)
+{
+	struct ethosu_capabilities *cap =
+		container_of(msg, typeof(*cap), msg);
+
+	/* Don't resend request if response has already been received */
+	if (completion_done(&cap->done))
+		return 0;
+
+	/* Resend request */
+	return ethosu_capabilities_send(cap);
+}
+
+void ethosu_capability_rsp(struct ethosu_device *edev,
+			   struct ethosu_core_msg_capabilities_rsp *rsp)
+{
+	int id = (int)rsp->user_arg;
+	struct ethosu_mailbox_msg *msg;
+	struct ethosu_capabilities *cap;
+
+	msg = ethosu_mailbox_find(&edev->mailbox, id);
+	if (IS_ERR(msg)) {
+		dev_warn(edev->dev,
+			 "Id for capabilities msg not found. id=%d\n",
+			 id);
+
+		return;
+	}
+
+	cap = container_of(msg, typeof(*cap), msg);
+
+	if (completion_done(&cap->done))
+		return;
+
+	cap->uapi->hw_id.version_status = rsp->version_status;
+	cap->uapi->hw_id.version_minor = rsp->version_minor;
+	cap->uapi->hw_id.version_major = rsp->version_major;
+	cap->uapi->hw_id.product_major = rsp->product_major;
+	cap->uapi->hw_id.arch_patch_rev = rsp->arch_patch_rev;
+	cap->uapi->hw_id.arch_minor_rev = rsp->arch_minor_rev;
+	cap->uapi->hw_id.arch_major_rev = rsp->arch_major_rev;
+	cap->uapi->driver_patch_rev = rsp->driver_patch_rev;
+	cap->uapi->driver_minor_rev = rsp->driver_minor_rev;
+	cap->uapi->driver_major_rev = rsp->driver_major_rev;
+	cap->uapi->hw_cfg.macs_per_cc = rsp->macs_per_cc;
+	cap->uapi->hw_cfg.cmd_stream_version = rsp->cmd_stream_version;
+	cap->uapi->hw_cfg.custom_dma = rsp->custom_dma;
+
+	cap->errno = 0;
+	complete(&cap->done);
+}
+
+int ethosu_capabilities_request(struct ethosu_device *edev,
+				struct ethosu_uapi_device_capabilities *uapi)
+{
+	struct ethosu_capabilities *cap;
+	int ret;
+	int timeout;
+
+	cap = devm_kzalloc(edev->dev, sizeof(struct ethosu_capabilities),
+			   GFP_KERNEL);
+	if (!cap)
+		return -ENOMEM;
+
+	cap->edev = edev;
+	cap->uapi = uapi;
+	init_completion(&cap->done);
+	cap->msg.fail = ethosu_capabilities_fail;
+	cap->msg.resend = ethosu_capabilities_resend;
+
+	ret = ethosu_mailbox_register(&cap->edev->mailbox, &cap->msg);
+	if (ret < 0)
+		goto kfree;
+
+	dev_info(edev->dev, "Capabilities create. Id=%d, handle=0x%p\n",
+		 cap->msg.id, cap);
+
+	ret = ethosu_capabilities_send(cap);
+	if (0 != ret)
+		goto deregister;
+
+	/* Unlock the mutex before going to block on the condition */
+	mutex_unlock(&edev->mutex);
+
+	/* wait for response to arrive back */
+	timeout = wait_for_completion_timeout(&cap->done,
+					      msecs_to_jiffies(
+						      CAPABILITIES_RESP_TIMEOUT_MS));
+
+	/* take back the mutex before resuming to do anything */
+	mutex_lock(&edev->mutex);
+
+	if (0 == timeout) {
+		dev_warn(edev->dev, "Capabilities response timeout");
+		ret = -ETIME;
+		goto deregister;
+	}
+
+	if (cap->errno) {
+		ret = cap->errno;
+		goto deregister;
+	}
+
+deregister:
+	ethosu_mailbox_deregister(&cap->edev->mailbox, &cap->msg);
+
+kfree:
+	dev_info(cap->edev->dev, "Capabilities destroy. Id=%d, handle=0x%p\n",
+		 cap->msg.id, cap);
+	devm_kfree(cap->edev->dev, cap);
+
+	return ret;
+}
diff --git a/kernel/ethosu_capabilities.h b/kernel/ethosu_capabilities.h
new file mode 100644
index 0000000..87cd168
--- /dev/null
+++ b/kernel/ethosu_capabilities.h
@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2022 Arm Limited.
+ *
+ * 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
+ * Foundation, and any use by you of this program is subject to the terms
+ * of such GNU licence.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, you can access it online at
+ * http://www.gnu.org/licenses/gpl-2.0.html.
+ *
+ * SPDX-License-Identifier: GPL-2.0-only
+ */
+
+#ifndef ETHOSU_CAPABILITIES_H
+#define ETHOSU_CAPABILITIES_H
+
+/****************************************************************************
+ * Includes
+ ****************************************************************************/
+
+#include "ethosu_core_interface.h"
+#include "ethosu_mailbox.h"
+
+#include <linux/types.h>
+#include <linux/completion.h>
+
+/****************************************************************************
+ * Types
+ ****************************************************************************/
+
+struct ethosu_device;
+struct ethosu_uapi_device_capabilities;
+
+/**
+ * struct ethosu_capabilities - Capabilities internal struct
+ */
+struct ethosu_capabilities {
+	struct ethosu_device                   *edev;
+	struct completion                      done;
+	struct ethosu_uapi_device_capabilities *uapi;
+	struct ethosu_mailbox_msg              msg;
+	int                                    errno;
+};
+
+/****************************************************************************
+ * Functions
+ ****************************************************************************/
+
+int ethosu_capabilities_request(struct ethosu_device *edev,
+				struct ethosu_uapi_device_capabilities *uapi);
+
+void ethosu_capability_rsp(struct ethosu_device *edev,
+			   struct ethosu_core_msg_capabilities_rsp *rsp);
+
+#endif
diff --git a/kernel/ethosu_device.c b/kernel/ethosu_device.c
index 6b911ca..6bc3c28 100644
--- a/kernel/ethosu_device.c
+++ b/kernel/ethosu_device.c
@@ -26,13 +26,13 @@
 
 #include "ethosu_buffer.h"
 #include "ethosu_core_interface.h"
+#include "ethosu_capabilities.h"
 #include "ethosu_inference.h"
 #include "ethosu_cancel_inference.h"
 #include "ethosu_network.h"
 #include "ethosu_network_info.h"
 #include "uapi/ethosu.h"
 
-#include <linux/delay.h>
 #include <linux/dma-mapping.h>
 #include <linux/errno.h>
 #include <linux/fs.h>
@@ -49,158 +49,10 @@
 
 #define DMA_ADDR_BITS 32 /* Number of address bits */
 
-#define CAPABILITIES_RESP_TIMEOUT_MS 2000
-
 /****************************************************************************
  * Functions
  ****************************************************************************/
 
-static void ethosu_capabilities_destroy(struct kref *kref)
-{
-	struct ethosu_capabilities *cap =
-		container_of(kref, struct ethosu_capabilities, refcount);
-
-	dev_info(cap->edev->dev, "Capabilities destroy. handle=0x%pK\n", cap);
-
-	list_del(&cap->msg.list);
-
-	devm_kfree(cap->edev->dev, cap);
-}
-
-static int ethosu_capabilities_send(struct ethosu_capabilities *cap)
-{
-	int ret;
-
-	ret = ethosu_mailbox_capabilities_request(&cap->edev->mailbox, cap);
-	if (ret)
-		return ret;
-
-	return 0;
-}
-
-static void ethosu_capabilities_fail(struct ethosu_mailbox_msg *msg)
-{
-	struct ethosu_capabilities *cap =
-		container_of(msg, typeof(*cap), msg);
-
-	if (completion_done(&cap->done))
-		return;
-
-	cap->errno = -EFAULT;
-	complete(&cap->done);
-}
-
-static int ethosu_capabilities_resend(struct ethosu_mailbox_msg *msg)
-{
-	struct ethosu_capabilities *cap =
-		container_of(msg, typeof(*cap), msg);
-	int ret;
-
-	/* Don't resend request if response has already been received */
-	if (completion_done(&cap->done))
-		return 0;
-
-	/* Resend request */
-	ret = ethosu_capabilities_send(cap);
-	if (ret)
-		return ret;
-
-	return 0;
-}
-
-static void ethosu_capability_rsp(struct ethosu_device *edev,
-				  struct ethosu_core_msg_capabilities_rsp *msg)
-{
-	struct ethosu_capabilities *cap =
-		(struct ethosu_capabilities *)msg->user_arg;
-	struct ethosu_uapi_device_capabilities *capabilities;
-	int ret;
-
-	ret = ethosu_mailbox_find(&edev->mailbox, &cap->msg);
-	if (ret) {
-		dev_warn(edev->dev,
-			 "Capabilities not found in pending list. handle=0x%p\n",
-			 cap);
-
-		return;
-	}
-
-	if (completion_done(&cap->done))
-		return;
-
-	capabilities = cap->capabilities;
-
-	capabilities->hw_id.version_status = msg->version_status;
-	capabilities->hw_id.version_minor = msg->version_minor;
-	capabilities->hw_id.version_major = msg->version_major;
-	capabilities->hw_id.product_major = msg->product_major;
-	capabilities->hw_id.arch_patch_rev = msg->arch_patch_rev;
-	capabilities->hw_id.arch_minor_rev = msg->arch_minor_rev;
-	capabilities->hw_id.arch_major_rev = msg->arch_major_rev;
-	capabilities->driver_patch_rev = msg->driver_patch_rev;
-	capabilities->driver_minor_rev = msg->driver_minor_rev;
-	capabilities->driver_major_rev = msg->driver_major_rev;
-	capabilities->hw_cfg.macs_per_cc = msg->macs_per_cc;
-	capabilities->hw_cfg.cmd_stream_version = msg->cmd_stream_version;
-	capabilities->hw_cfg.custom_dma = msg->custom_dma;
-
-	cap->errno = 0;
-	complete(&cap->done);
-}
-
-static int ethosu_capabilities_request(struct ethosu_device *edev,
-				       void __user *udata)
-{
-	struct ethosu_uapi_device_capabilities uapi;
-	struct ethosu_capabilities *cap;
-	int ret;
-	int timeout;
-
-	cap = devm_kzalloc(edev->dev, sizeof(struct ethosu_capabilities),
-			   GFP_KERNEL);
-	if (!cap)
-		return -ENOMEM;
-
-	cap->edev = edev;
-	cap->capabilities = &uapi;
-	kref_init(&cap->refcount);
-	init_completion(&cap->done);
-	list_add(&cap->msg.list, &edev->mailbox.pending_list);
-	cap->msg.fail = ethosu_capabilities_fail;
-	cap->msg.resend = ethosu_capabilities_resend;
-
-	ret = ethosu_capabilities_send(cap);
-	if (0 != ret)
-		goto put_kref;
-
-	/* Unlock the mutex before going to block on the condition */
-	mutex_unlock(&edev->mutex);
-
-	/* wait for response to arrive back */
-	timeout = wait_for_completion_timeout(&cap->done,
-					      msecs_to_jiffies(
-						      CAPABILITIES_RESP_TIMEOUT_MS));
-
-	/* take back the mutex before resuming to do anything */
-	mutex_lock(&edev->mutex);
-
-	if (0 == timeout) {
-		dev_warn(edev->dev, "Capabilities response timeout");
-		ret = -EIO;
-		goto put_kref;
-	}
-
-	if (cap->errno)
-		goto put_kref;
-
-	ret = copy_to_user(udata, &uapi, sizeof(uapi)) ? -EFAULT : 0;
-
-put_kref:
-	kref_put(&cap->refcount, ethosu_capabilities_destroy);
-
-	return ret;
-}
-
 /* Incoming messages */
 static int ethosu_handle_msg(struct ethosu_device *edev)
 {
@@ -463,10 +315,16 @@ static long ethosu_ioctl(struct file *file,
 		dev_info(edev->dev, "Ioctl: Send version request\n");
 		ret = ethosu_mailbox_version_request(&edev->mailbox);
 		break;
-	case ETHOSU_IOCTL_CAPABILITIES_REQ:
+	case ETHOSU_IOCTL_CAPABILITIES_REQ: {
+		struct ethosu_uapi_device_capabilities uapi;
 		dev_info(edev->dev, "Ioctl: Send capabilities request\n");
-		ret = ethosu_capabilities_request(edev, udata);
+		ret = ethosu_capabilities_request(edev, &uapi);
+		if (ret)
+			break;
+
+		ret = copy_to_user(udata, &uapi, sizeof(uapi)) ? -EFAULT : 0;
 		break;
+	}
 	case ETHOSU_IOCTL_PING: {
 		dev_info(edev->dev, "Ioctl: Send ping\n");
 		ret = ethosu_mailbox_ping(&edev->mailbox);
diff --git a/kernel/ethosu_device.h b/kernel/ethosu_device.h
index 132dff2..d1e4334 100644
--- a/kernel/ethosu_device.h
+++ b/kernel/ethosu_device.h
@@ -33,8 +33,6 @@
 #include <linux/cdev.h>
 #include <linux/io.h>
 #include <linux/mutex.h>
-#include <linux/workqueue.h>
-#include <linux/completion.h>
 
 /****************************************************************************
  * Types
@@ -56,18 +54,6 @@ struct ethosu_device {
 	struct reset_control   *reset;
 };
 
-/**
- * struct ethosu_capabilities - Capabilities internal struct
- */
-struct ethosu_capabilities {
-	struct ethosu_device                   *edev;
-	struct completion                      done;
-	struct kref                            refcount;
-	struct ethosu_uapi_device_capabilities *capabilities;
-	struct ethosu_mailbox_msg              msg;
-	int                                    errno;
-};
-
 /****************************************************************************
  * Functions
  ****************************************************************************/
diff --git a/kernel/ethosu_inference.c b/kernel/ethosu_inference.c
index 0599b53..55cabc3 100644
--- a/kernel/ethosu_inference.c
+++ b/kernel/ethosu_inference.c
@@ -95,7 +95,7 @@ static int ethosu_inference_send(struct ethosu_inference *inf)
 
 	inf->status = ETHOSU_UAPI_STATUS_ERROR;
 
-	ret = ethosu_mailbox_inference(&inf->edev->mailbox, inf,
+	ret = ethosu_mailbox_inference(&inf->edev->mailbox, &inf->msg,
 				       inf->ifm_count, inf->ifm,
 				       inf->ofm_count, inf->ofm,
 				       inf->net->buf,
@@ -180,7 +180,7 @@ static void ethosu_inference_kref_destroy(struct kref *kref)
 		 "Inference destroy. handle=0x%pK, status=%d\n",
 		 inf, inf->status);
 
-	list_del(&inf->msg.list);
+	ethosu_mailbox_deregister(&inf->edev->mailbox, &inf->msg);
 
 	while (inf->ifm_count-- > 0)
 		ethosu_buffer_put(inf->ifm[inf->ifm_count]);
@@ -307,6 +307,11 @@ int ethosu_inference_create(struct ethosu_device *edev,
 	inf->msg.fail = ethosu_inference_fail;
 	inf->msg.resend = ethosu_inference_resend;
 
+	/* Add inference to pending list */
+	ret = ethosu_mailbox_register(&edev->mailbox, &inf->msg);
+	if (ret < 0)
+		goto kfree;
+
 	/* Get pointer to IFM buffers */
 	for (i = 0; i < uapi->ifm_count; i++) {
 		inf->ifm[i] = ethosu_buffer_get_from_fd(uapi->ifm_fd[i]);
@@ -361,14 +366,13 @@ int ethosu_inference_create(struct ethosu_device *edev,
 	inf->file = fget(ret);
 	fput(inf->file);
 
-	/* Add inference to pending list */
-	list_add(&inf->msg.list, &edev->mailbox.pending_list);
-
 	/* Send inference request to Arm Ethos-U subsystem */
-	(void)ethosu_inference_send(inf);
+	ret = ethosu_inference_send(inf);
+	if (ret)
+		goto put_net;
 
-	dev_info(edev->dev, "Inference create. handle=0x%pK, fd=%d",
-		 inf, fd);
+	dev_info(edev->dev, "Inference create. Id=%d, handle=0x%p, fd=%d",
+		 inf->msg.id, inf, fd);
 
 	return fd;
 
@@ -383,6 +387,7 @@ put_ifm:
 	while (inf->ifm_count-- > 0)
 		ethosu_buffer_put(inf->ifm[inf->ifm_count]);
 
+kfree:
 	devm_kfree(edev->dev, inf);
 
 	return ret;
@@ -423,20 +428,23 @@ int ethosu_inference_put(struct ethosu_inference *inf)
 void ethosu_inference_rsp(struct ethosu_device *edev,
 			  struct ethosu_core_inference_rsp *rsp)
 {
-	struct ethosu_inference *inf =
-		(struct ethosu_inference *)rsp->user_arg;
+	int id = (int)rsp->user_arg;
+	struct ethosu_mailbox_msg *msg;
+	struct ethosu_inference *inf;
 	int ret;
 	int i;
 
-	ret = ethosu_mailbox_find(&edev->mailbox, &inf->msg);
-	if (ret) {
+	msg = ethosu_mailbox_find(&edev->mailbox, id);
+	if (IS_ERR(msg)) {
 		dev_warn(edev->dev,
-			 "Handle not found in inference list. handle=0x%p\n",
-			 rsp);
+			 "Id for inference msg not found. Id=%d\n",
+			 id);
 
 		return;
 	}
 
+	inf = container_of(msg, typeof(*inf), msg);
+
 	if (rsp->status == ETHOSU_CORE_STATUS_OK &&
 	    inf->ofm_count <= ETHOSU_CORE_BUFFER_MAX) {
 		uint32_t i;
diff --git a/kernel/ethosu_mailbox.c b/kernel/ethosu_mailbox.c
index 843cb58..014ceb8 100644
--- a/kernel/ethosu_mailbox.c
+++ b/kernel/ethosu_mailbox.c
@@ -282,34 +282,51 @@ int ethosu_mailbox_read(struct ethosu_mailbox *mbox,
 	return 0;
 }
 
-int ethosu_mailbox_find(struct ethosu_mailbox *mbox,
-			struct ethosu_mailbox_msg *msg)
+int ethosu_mailbox_register(struct ethosu_mailbox *mbox,
+			    struct ethosu_mailbox_msg *msg)
 {
-	struct ethosu_mailbox_msg *cur;
+	msg->id = idr_alloc_cyclic(&mbox->msg_idr, msg, 0, INT_MAX, GFP_KERNEL);
+	if (msg->id < 0)
+		return msg->id;
 
-	list_for_each_entry(cur, &mbox->pending_list, list) {
-		if (cur == msg)
-			return 0;
-	}
+	return 0;
+}
+
+void ethosu_mailbox_deregister(struct ethosu_mailbox *mbox,
+			       struct ethosu_mailbox_msg *msg)
+{
+	idr_remove(&mbox->msg_idr, msg->id);
+}
+
+struct ethosu_mailbox_msg *ethosu_mailbox_find(struct ethosu_mailbox *mbox,
+					       int msg_id)
+{
+	struct ethosu_mailbox_msg *ptr = (struct ethosu_mailbox_msg *)idr_find(
+		&mbox->msg_idr, msg_id);
 
-	return -EINVAL;
+	if (ptr == NULL)
+		return ERR_PTR(-EINVAL);
+
+	return ptr;
 }
 
 void ethosu_mailbox_fail(struct ethosu_mailbox *mbox)
 {
-	struct ethosu_mailbox_msg *cur, *cur_tmp;
+	struct ethosu_mailbox_msg *cur;
+	int id;
 
-	list_for_each_entry_safe(cur, cur_tmp, &mbox->pending_list, list) {
+	idr_for_each_entry(&mbox->msg_idr, cur, id) {
 		cur->fail(cur);
 	}
 }
 
 void ethosu_mailbox_resend(struct ethosu_mailbox *mbox)
 {
-	struct ethosu_mailbox_msg *cur, *cur_tmp;
+	struct ethosu_mailbox_msg *cur;
+	int id;
 	int ret;
 
-	list_for_each_entry_safe(cur, cur_tmp, &mbox->pending_list, list) {
+	idr_for_each_entry(&mbox->msg_idr, cur, id) {
 		ret = cur->resend(cur);
 		if (ret) {
 			dev_warn(mbox->dev, "Failed to resend msg. ret=%d",
@@ -336,10 +353,10 @@ int ethosu_mailbox_version_request(struct ethosu_mailbox *mbox)
 }
 
 int ethosu_mailbox_capabilities_request(struct ethosu_mailbox *mbox,
-					void *user_arg)
+					struct ethosu_mailbox_msg *msg)
 {
 	struct ethosu_core_capabilities_req req = {
-		.user_arg = (ptrdiff_t)user_arg
+		.user_arg = msg->id
 	};
 
 	return ethosu_queue_write_msg(mbox, ETHOSU_CORE_MSG_CAPABILITIES_REQ,
@@ -348,7 +365,7 @@ int ethosu_mailbox_capabilities_request(struct ethosu_mailbox *mbox,
 }
 
 int ethosu_mailbox_inference(struct ethosu_mailbox *mbox,
-			     void *user_arg,
+			     struct ethosu_mailbox_msg *msg,
 			     uint32_t ifm_count,
 			     struct ethosu_buffer **ifm,
 			     uint32_t ofm_count,
@@ -369,7 +386,7 @@ int ethosu_mailbox_inference(struct ethosu_mailbox *mbox,
 		return -EINVAL;
 	}
 
-	inf.user_arg = (ptrdiff_t)user_arg;
+	inf.user_arg = msg->id;
 	inf.ifm_count = ifm_count;
 	inf.ofm_count = ofm_count;
 	inf.pmu_cycle_counter_enable = pmu_cycle_counter_enable;
@@ -396,13 +413,13 @@ int ethosu_mailbox_inference(struct ethosu_mailbox *mbox,
 }
 
 int ethosu_mailbox_network_info_request(struct ethosu_mailbox *mbox,
-					void *user_arg,
+					struct ethosu_mailbox_msg *msg,
 					struct ethosu_buffer *network,
 					uint32_t network_index)
 {
 	struct ethosu_core_network_info_req info;
 
-	info.user_arg = (ptrdiff_t)user_arg;
+	info.user_arg = msg->id;
 
 	if (network != NULL) {
 		info.network.type = ETHOSU_CORE_NETWORK_BUFFER;
@@ -417,13 +434,13 @@ int ethosu_mailbox_network_info_request(struct ethosu_mailbox *mbox,
 }
 
 int ethosu_mailbox_cancel_inference(struct ethosu_mailbox *mbox,
-				    void *user_arg,
-				    void *inference_handle)
+				    struct ethosu_mailbox_msg *msg,
+				    int inference_handle)
 {
 	struct ethosu_core_cancel_inference_req req;
 
-	req.user_arg = (ptrdiff_t)user_arg;
-	req.inference_handle = (ptrdiff_t)inference_handle;
+	req.user_arg = msg->id;
+	req.inference_handle = inference_handle;
 
 	return ethosu_queue_write_msg(mbox,
 				      ETHOSU_CORE_MSG_CANCEL_INFERENCE_REQ,
@@ -473,7 +490,7 @@ int ethosu_mailbox_init(struct ethosu_mailbox *mbox,
 	mbox->user_arg = user_arg;
 	mbox->wdog = wdog;
 	mbox->ping_count = 0;
-	INIT_LIST_HEAD(&mbox->pending_list);
+	idr_init(&mbox->msg_idr);
 
 	mbox->client.dev = dev;
 	mbox->client.rx_callback = ethosu_mailbox_rx_callback;
diff --git a/kernel/ethosu_mailbox.h b/kernel/ethosu_mailbox.h
index 07276f6..26367f6 100644
--- a/kernel/ethosu_mailbox.h
+++ b/kernel/ethosu_mailbox.h
@@ -29,6 +29,7 @@
 #include <linux/types.h>
 #include <linux/mailbox_client.h>
 #include <linux/workqueue.h>
+#include <linux/idr.h>
 
 /****************************************************************************
  * Types
@@ -55,15 +56,15 @@ struct ethosu_mailbox {
 	struct mbox_chan         *tx;
 	ethosu_mailbox_cb        callback;
 	void                     *user_arg;
-	struct list_head         pending_list;
+	struct idr               msg_idr;
 	struct ethosu_watchdog   *wdog;
 	unsigned                 ping_count;
 };
 
 struct ethosu_mailbox_msg {
-	struct list_head list;
-	void             (*fail)(struct ethosu_mailbox_msg *msg);
-	int              (*resend)(struct ethosu_mailbox_msg *msg);
+	int  id;
+	void (*fail)(struct ethosu_mailbox_msg *msg);
+	int  (*resend)(struct ethosu_mailbox_msg *msg);
 };
 
 /****************************************************************************
@@ -119,12 +120,26 @@ int ethosu_mailbox_read(struct ethosu_mailbox *mbox,
 			size_t length);
 
 /**
- * ethosu_mailbox_find() - Find mailbox message
+ * ethosu_mailbox_register() - Register the ethosu_mailbox_msg in ethosu_mailbox
  *
  * Return: 0 on success, else error code.
  */
-int ethosu_mailbox_find(struct ethosu_mailbox *mbox,
-			struct ethosu_mailbox_msg *msg);
+int ethosu_mailbox_register(struct ethosu_mailbox *mbox,
+			    struct ethosu_mailbox_msg *msg);
+
+/**
+ * ethosu_mailbox_free_id() - Free the id of the ethosu_mailbox_msg
+ */
+void ethosu_mailbox_deregister(struct ethosu_mailbox *mbox,
+			       struct ethosu_mailbox_msg *msg);
+
+/**
+ * ethosu_mailbox_find() - Find mailbox message
+ *
+ * Return: a valid pointer on success, otherwise an error ptr.
+ */
+struct ethosu_mailbox_msg *ethosu_mailbox_find(struct ethosu_mailbox *mbox,
+					       int msg_id);
 
 /**
  * ethosu_mailbox_fail() - Fail mailbox messages
@@ -137,7 +152,6 @@ void ethosu_mailbox_fail(struct ethosu_mailbox *mbox);
  * ethosu_mailbox_resend() - Resend mailbox messages
  *
  * Call resend() callback on all messages in pending list.
- *
  */
 void ethosu_mailbox_resend(struct ethosu_mailbox *mbox);
 
@@ -173,7 +187,7 @@ int ethosu_mailbox_version_request(struct ethosu_mailbox *mbox);
  * Return: 0 on success, else error code.
  */
 int ethosu_mailbox_capabilities_request(struct ethosu_mailbox *mbox,
-					void *user_arg);
+					struct ethosu_mailbox_msg *msg);
 
 /**
  * ethosu_mailbox_inference() - Send inference
@@ -181,7 +195,7 @@ int ethosu_mailbox_capabilities_request(struct ethosu_mailbox *mbox,
  * Return: 0 on success, else error code.
  */
 int ethosu_mailbox_inference(struct ethosu_mailbox *mbox,
-			     void *user_arg,
+			     struct ethosu_mailbox_msg *msg,
 			     uint32_t ifm_count,
 			     struct ethosu_buffer **ifm,
 			     uint32_t ofm_count,
@@ -198,7 +212,7 @@ int ethosu_mailbox_inference(struct ethosu_mailbox *mbox,
  * Return: 0 on success, else error code.
  */
 int ethosu_mailbox_network_info_request(struct ethosu_mailbox *mbox,
-					void *user_arg,
+					struct ethosu_mailbox_msg *msg,
 					struct ethosu_buffer *network,
 					uint32_t network_index);
 
@@ -208,7 +222,7 @@ int ethosu_mailbox_network_info_request(struct ethosu_mailbox *mbox,
  * Return: 0 on success, else error code.
  */
 int ethosu_mailbox_cancel_inference(struct ethosu_mailbox *mbox,
-				    void *user_arg,
-				    void *inference_handle);
+				    struct ethosu_mailbox_msg *msg,
+				    int inference_handle);
 
 #endif /* ETHOSU_MAILBOX_H */
diff --git a/kernel/ethosu_network.c b/kernel/ethosu_network.c
index 86ae410..4170046 100644
--- a/kernel/ethosu_network.c
+++ b/kernel/ethosu_network.c
@@ -88,30 +88,6 @@ static int ethosu_network_release(struct inode *inode,
 	return 0;
 }
 
-static int ethosu_network_info_request(struct ethosu_network *net,
-				       struct ethosu_uapi_network_info *uapi)
-{
-	struct ethosu_network_info *info;
-	int ret;
-
-	/* Create network info request */
-	info = ethosu_network_info_create(net->edev, net, uapi);
-	if (IS_ERR(info))
-		return PTR_ERR(info);
-
-	/* Unlock the device mutex and wait for completion */
-	mutex_unlock(&net->edev->mutex);
-	ret = ethosu_network_info_wait(info, 3000);
-	mutex_lock(&net->edev->mutex);
-
-	if (ret)
-		info->msg.fail(&info->msg);
-
-	ethosu_network_info_put(info);
-
-	return ret;
-}
-
 static long ethosu_network_ioctl(struct file *file,
 				 unsigned int cmd,
 				 unsigned long arg)
diff --git a/kernel/ethosu_network_info.c b/kernel/ethosu_network_info.c
index ebb059c..da95309 100644
--- a/kernel/ethosu_network_info.c
+++ b/kernel/ethosu_network_info.c
@@ -29,37 +29,15 @@
 #include "ethosu_mailbox.h"
 #include "uapi/ethosu.h"
 
-/****************************************************************************
- * Functions
- ****************************************************************************/
+#define NETWORK_INFO_RESP_TIMEOUT_MS 3000
 
-static void ethosu_network_info_destroy(struct kref *kref)
+static inline int ethosu_network_info_send(struct ethosu_network_info *info)
 {
-	struct ethosu_network_info *info =
-		container_of(kref, struct ethosu_network_info, kref);
-
-	dev_info(info->edev->dev, "Network info destroy. handle=0x%pK\n", info);
-
-	list_del(&info->msg.list);
-
-	ethosu_network_put(info->net);
-
-	devm_kfree(info->edev->dev, info);
-}
-
-static int ethosu_network_info_send(struct ethosu_network_info *info)
-{
-	int ret;
-
 	/* Send network info request to firmware */
-	ret = ethosu_mailbox_network_info_request(&info->edev->mailbox,
-						  info,
-						  info->net->buf,
-						  info->net->index);
-	if (ret)
-		return ret;
-
-	return 0;
+	return ethosu_mailbox_network_info_request(&info->edev->mailbox,
+						   &info->msg,
+						   info->net->buf,
+						   info->net->index);
 }
 
 static void ethosu_network_info_fail(struct ethosu_mailbox_msg *msg)
@@ -92,91 +70,84 @@ static int ethosu_network_info_resend(struct ethosu_mailbox_msg *msg)
 	return 0;
 }
 
-struct ethosu_network_info *ethosu_network_info_create(
-	struct ethosu_device *edev,
-	struct ethosu_network *net,
-	struct ethosu_uapi_network_info *uapi)
+int ethosu_network_info_request(struct ethosu_network *net,
+				struct ethosu_uapi_network_info *uapi)
 {
 	struct ethosu_network_info *info;
 	int ret;
+	int timeout;
 
-	info = devm_kzalloc(edev->dev, sizeof(*info), GFP_KERNEL);
+	info = devm_kzalloc(net->edev->dev, sizeof(*info), GFP_KERNEL);
 	if (!info)
-		return ERR_PTR(-ENOMEM);
+		return -ENOMEM;
 
-	info->edev = edev;
+	info->edev = net->edev;
 	info->net = net;
 	info->uapi = uapi;
-	kref_init(&info->kref);
 	init_completion(&info->done);
 	info->msg.fail = ethosu_network_info_fail;
 	info->msg.resend = ethosu_network_info_resend;
 
-	/* Insert network info to network info list */
-	list_add(&info->msg.list, &edev->mailbox.pending_list);
+	ret = ethosu_mailbox_register(&info->edev->mailbox, &info->msg);
+	if (ret < 0)
+		goto kfree;
 
 	/* Get reference to network */
-	ethosu_network_get(net);
+	ethosu_network_get(info->net);
 
 	ret = ethosu_network_info_send(info);
 	if (ret)
-		goto put_info;
-
-	dev_info(edev->dev, "Network info create. handle=%p\n", info);
-
-	return info;
+		goto deregister;
 
-put_info:
-	ethosu_network_info_put(info);
-
-	return ERR_PTR(ret);
-}
-
-void ethosu_network_info_get(struct ethosu_network_info *info)
-{
-	kref_get(&info->kref);
-}
-
-int ethosu_network_info_put(struct ethosu_network_info *info)
-{
-	return kref_put(&info->kref, ethosu_network_info_destroy);
-}
-
-int ethosu_network_info_wait(struct ethosu_network_info *info,
-			     int timeout_ms)
-{
-	int timeout;
+	dev_info(info->edev->dev, "Network info create. Id=%d, handle=0x%p\n\n",
+		 info->msg.id, info);
 
+	/* Unlock the device mutex and wait for completion */
+	mutex_unlock(&info->edev->mutex);
 	timeout = wait_for_completion_timeout(&info->done,
-					      msecs_to_jiffies(timeout_ms));
+					      msecs_to_jiffies(
+						      NETWORK_INFO_RESP_TIMEOUT_MS));
+	mutex_lock(&info->edev->mutex);
 
-	if (!timeout) {
-		dev_warn(info->edev->dev,
-			 "Network info timed out.");
+	if (0 == timeout) {
+		dev_warn(info->edev->dev, "Network info timed out.");
 
-		return -ETIME;
+		ret = -ETIME;
+		goto deregister;
 	}
 
-	return info->errno;
+deregister:
+	ethosu_mailbox_deregister(&info->edev->mailbox, &info->msg);
+	ethosu_network_put(info->net);
+
+kfree:
+	dev_info(info->edev->dev, "Network info destroy. Id=%d, handle=0x%p\n",
+		 info->msg.id, info);
+	devm_kfree(info->edev->dev, info);
+
+	return ret;
 }
 
 void ethosu_network_info_rsp(struct ethosu_device *edev,
 			     struct ethosu_core_network_info_rsp *rsp)
 {
-	struct ethosu_network_info *info =
-		(struct ethosu_network_info *)rsp->user_arg;
-	uint32_t i;
 	int ret;
+	int id = (int)rsp->user_arg;
+	struct ethosu_mailbox_msg *msg;
+	struct ethosu_network_info *info;
+	uint32_t i;
 
-	ret = ethosu_mailbox_find(&edev->mailbox, &info->msg);
-	if (0 != ret) {
+	msg = ethosu_mailbox_find(&edev->mailbox, id);
+	if (IS_ERR(msg)) {
 		dev_warn(edev->dev,
-			 "Handle not found in network info list. handle=0x%p\n",
-			 info);
+			 "Id for network info msg not found. Id=%d\n",
+			 id);
 
 		return;
 	}
 
+	info = container_of(msg, typeof(*info), msg);
+
 	if (completion_done(&info->done))
 		return;
 
diff --git a/kernel/ethosu_network_info.h b/kernel/ethosu_network_info.h
index 680be80..facca54 100644
--- a/kernel/ethosu_network_info.h
+++ b/kernel/ethosu_network_info.h
@@ -28,7 +28,6 @@
 #include "ethosu_core_interface.h"
 #include "ethosu_mailbox.h"
 
-#include <linux/kref.h>
 #include <linux/types.h>
 #include <linux/completion.h>
 
@@ -44,7 +43,6 @@ struct ethosu_network_info {
 	struct ethosu_device            *edev;
 	struct ethosu_network           *net;
 	struct ethosu_uapi_network_info *uapi;
-	struct kref                     kref;
 	struct completion               done;
 	int                             errno;
 	struct ethosu_mailbox_msg       msg;
@@ -55,34 +53,14 @@ struct ethosu_network_info {
  ****************************************************************************/
 
 /**
- * ethosu_network_info_create() - Create network info
+ * ethosu_network_info_request() - Send a network info request
  *
  * This function must be called in the context of a user space process.
  *
- * Return: Valid pointer on success, else ERR_PTR.
+ * Return: 0 on success, .
  */
-struct ethosu_network_info *ethosu_network_info_create(
-	struct ethosu_device *edev,
-	struct ethosu_network *net,
-	struct ethosu_uapi_network_info *uapi);
-
-/**
- * ethosu_network_info_get() - Get network info
- */
-void ethosu_network_info_get(struct ethosu_network_info *info);
-
-/**
- * ethosu_network_info_put() - Put network info
- *
- * Return: 1 if object was removed, else 0.
- */
-int ethosu_network_info_put(struct ethosu_network_info *info);
-
-/**
- * ethosu_network_info_wait() - Get network info
- */
-int ethosu_network_info_wait(struct ethosu_network_info *info,
-			     int timeout);
+int ethosu_network_info_request(struct ethosu_network *net,
+				struct ethosu_uapi_network_info *uapi);
 
 /**
  * ethosu_network_info_rsp() - Handle network info response.