static DEFINE_MUTEX(mbox_lock);
 
-static int send_mbox_cmd(struct pci_dev *pdev, u16 cmd_id, u32 cmd_data, u64 *cmd_resp)
+static int wait_for_mbox_ready(struct proc_thermal_device *proc_priv)
 {
-       struct proc_thermal_device *proc_priv;
        u32 retries, data;
        int ret;
 
-       mutex_lock(&mbox_lock);
-       proc_priv = pci_get_drvdata(pdev);
-
        /* Poll for rb bit == 0 */
        retries = MBOX_RETRY_COUNT;
        do {
-               data = readl((void __iomem *) (proc_priv->mmio_base + MBOX_OFFSET_INTERFACE));
+               data = readl(proc_priv->mmio_base + MBOX_OFFSET_INTERFACE);
                if (data & BIT_ULL(MBOX_BUSY_BIT)) {
                        ret = -EBUSY;
                        continue;
                break;
        } while (--retries);
 
+       return ret;
+}
+
+static int send_mbox_write_cmd(struct pci_dev *pdev, u16 id, u32 data)
+{
+       struct proc_thermal_device *proc_priv;
+       u32 reg_data;
+       int ret;
+
+       proc_priv = pci_get_drvdata(pdev);
+
+       mutex_lock(&mbox_lock);
+
+       ret = wait_for_mbox_ready(proc_priv);
        if (ret)
                goto unlock_mbox;
 
-       if (cmd_id == MBOX_CMD_WORKLOAD_TYPE_WRITE)
-               writel(cmd_data, (void __iomem *) ((proc_priv->mmio_base + MBOX_OFFSET_DATA)));
-
+       writel(data, (proc_priv->mmio_base + MBOX_OFFSET_DATA));
        /* Write command register */
-       data = BIT_ULL(MBOX_BUSY_BIT) | cmd_id;
-       writel(data, (void __iomem *) ((proc_priv->mmio_base + MBOX_OFFSET_INTERFACE)));
+       reg_data = BIT_ULL(MBOX_BUSY_BIT) | id;
+       writel(reg_data, (proc_priv->mmio_base + MBOX_OFFSET_INTERFACE));
 
-       /* Poll for rb bit == 0 */
-       retries = MBOX_RETRY_COUNT;
-       do {
-               data = readl((void __iomem *) (proc_priv->mmio_base + MBOX_OFFSET_INTERFACE));
-               if (data & BIT_ULL(MBOX_BUSY_BIT)) {
-                       ret = -EBUSY;
-                       continue;
-               }
+       ret = wait_for_mbox_ready(proc_priv);
 
-               if (data) {
-                       ret = -ENXIO;
-                       goto unlock_mbox;
-               }
+unlock_mbox:
+       mutex_unlock(&mbox_lock);
+       return ret;
+}
 
-               ret = 0;
+static int send_mbox_read_cmd(struct pci_dev *pdev, u16 id, u64 *resp)
+{
+       struct proc_thermal_device *proc_priv;
+       u32 reg_data;
+       int ret;
 
-               if (!cmd_resp)
-                       break;
+       proc_priv = pci_get_drvdata(pdev);
 
-               if (cmd_id == MBOX_CMD_WORKLOAD_TYPE_READ)
-                       *cmd_resp = readl((void __iomem *) (proc_priv->mmio_base + MBOX_OFFSET_DATA));
-               else
-                       *cmd_resp = readq((void __iomem *) (proc_priv->mmio_base + MBOX_OFFSET_DATA));
+       mutex_lock(&mbox_lock);
 
-               break;
-       } while (--retries);
+       ret = wait_for_mbox_ready(proc_priv);
+       if (ret)
+               goto unlock_mbox;
+
+       /* Write command register */
+       reg_data = BIT_ULL(MBOX_BUSY_BIT) | id;
+       writel(reg_data, (proc_priv->mmio_base + MBOX_OFFSET_INTERFACE));
+
+       ret = wait_for_mbox_ready(proc_priv);
+       if (ret)
+               goto unlock_mbox;
+
+       if (id == MBOX_CMD_WORKLOAD_TYPE_READ)
+               *resp = readl(proc_priv->mmio_base + MBOX_OFFSET_DATA);
+       else
+               *resp = readq(proc_priv->mmio_base + MBOX_OFFSET_DATA);
 
 unlock_mbox:
        mutex_unlock(&mbox_lock);
        return ret;
 }
 
-int processor_thermal_send_mbox_cmd(struct pci_dev *pdev, u16 cmd_id, u32 cmd_data, u64 *cmd_resp)
+int processor_thermal_send_mbox_read_cmd(struct pci_dev *pdev, u16 id, u64 *resp)
 {
-       return send_mbox_cmd(pdev, cmd_id, cmd_data, cmd_resp);
+       return send_mbox_read_cmd(pdev, id, resp);
 }
-EXPORT_SYMBOL_GPL(processor_thermal_send_mbox_cmd);
+EXPORT_SYMBOL_NS_GPL(processor_thermal_send_mbox_read_cmd, INT340X_THERMAL);
+
+int processor_thermal_send_mbox_write_cmd(struct pci_dev *pdev, u16 id, u32 data)
+{
+       return send_mbox_write_cmd(pdev, id, data);
+}
+EXPORT_SYMBOL_NS_GPL(processor_thermal_send_mbox_write_cmd, INT340X_THERMAL);
 
 /* List of workload types */
 static const char * const workload_types[] = {
        NULL
 };
 
-
 static ssize_t workload_available_types_show(struct device *dev,
                                               struct device_attribute *attr,
                                               char *buf)
 
        data |= ret;
 
-       ret = send_mbox_cmd(pdev, MBOX_CMD_WORKLOAD_TYPE_WRITE, data, NULL);
+       ret = send_mbox_write_cmd(pdev, MBOX_CMD_WORKLOAD_TYPE_WRITE, data);
        if (ret)
                return false;
 
        u64 cmd_resp;
        int ret;
 
-       ret = send_mbox_cmd(pdev, MBOX_CMD_WORKLOAD_TYPE_READ, 0, &cmd_resp);
+       ret = send_mbox_read_cmd(pdev, MBOX_CMD_WORKLOAD_TYPE_READ, &cmd_resp);
        if (ret)
                return false;
 
        .name = "workload_request"
 };
 
-
-
 static bool workload_req_created;
 
 int proc_thermal_mbox_add(struct pci_dev *pdev, struct proc_thermal_device *proc_priv)
        int ret;
 
        /* Check if there is a mailbox support, if fails return success */
-       ret = send_mbox_cmd(pdev, MBOX_CMD_WORKLOAD_TYPE_READ, 0, &cmd_resp);
+       ret = send_mbox_read_cmd(pdev, MBOX_CMD_WORKLOAD_TYPE_READ, &cmd_resp);
        if (ret)
                return 0;
 
 
 #include <linux/pci.h>
 #include "processor_thermal_device.h"
 
+MODULE_IMPORT_NS(INT340X_THERMAL);
+
 struct mmio_reg {
        int read_only;
        u32 offset;
                                     struct device_attribute *attr,
                                     const char *buf, size_t count)
 {
-       u16 cmd_id = 0x0008;
-       u64 cmd_resp;
+       u16 id = 0x0008;
        u32 input;
        int ret;
 
        if (ret)
                return ret;
 
-       ret = processor_thermal_send_mbox_cmd(to_pci_dev(dev), cmd_id, input, &cmd_resp);
+       ret = processor_thermal_send_mbox_write_cmd(to_pci_dev(dev), id, input);
        if (ret)
                return ret;
 
                                    struct device_attribute *attr,
                                    char *buf)
 {
-       u16 cmd_id = 0x0007;
-       u64 cmd_resp;
+       u16 id = 0x0007;
+       u64 resp;
        int ret;
 
-       ret = processor_thermal_send_mbox_cmd(to_pci_dev(dev), cmd_id, 0, &cmd_resp);
+       ret = processor_thermal_send_mbox_read_cmd(to_pci_dev(dev), id, &resp);
        if (ret)
                return ret;
 
-       return sprintf(buf, "%llu\n", cmd_resp);
+       return sprintf(buf, "%llu\n", resp);
 }
 
 static ssize_t ddr_data_rate_show(struct device *dev,
                                  struct device_attribute *attr,
                                  char *buf)
 {
-       u16 cmd_id = 0x0107;
-       u64 cmd_resp;
+       u16 id = 0x0107;
+       u64 resp;
        int ret;
 
-       ret = processor_thermal_send_mbox_cmd(to_pci_dev(dev), cmd_id, 0, &cmd_resp);
+       ret = processor_thermal_send_mbox_read_cmd(to_pci_dev(dev), id, &resp);
        if (ret)
                return ret;
 
-       return sprintf(buf, "%llu\n", cmd_resp);
+       return sprintf(buf, "%llu\n", resp);
 }
 
 static DEVICE_ATTR_RW(rfi_restriction);