Use ids for identifying messages sent to EthosU firmware

Avoid to use pointers to kernel memory as identify messages, prior this
change, if memory is reused that can lead to unexpected conflicts.

Remove ref counting from capabilities and network_info since memory is
freed in only one place.

Finally, extract the capabilities code in its own files.

Change-Id: I9ab7257f3ab85fa3347be65b3739e0daf9b5a441

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.