1 files changed, 21 insertions, 13 deletions

diff --git a/drivers/soc/qcom/rpmh.c b/drivers/soc/qcom/rpmh.c
index 9fb25e627698..035091fd44b8 100644
--- a/drivers/soc/qcom/rpmh.c
+++ b/drivers/soc/qcom/rpmh.c
@@ -80,6 +80,7 @@ void rpmh_tx_done(const struct tcs_request *msg, int r)
         struct rpmh_request *rpm_msg = container_of(msg, struct rpmh_request,
                                                     msg);
         struct completion *compl = rpm_msg->completion;
+        bool free = rpm_msg->needs_free;
 
         rpm_msg->err = r;
 
@@ -94,7 +95,7 @@ void rpmh_tx_done(const struct tcs_request *msg, int r)
         complete(compl);
 
 exit:
-        if (rpm_msg->needs_free)
+        if (free)
                 kfree(rpm_msg);
 }
 
@@ -347,11 +348,12 @@ int rpmh_write_batch(const struct device *dev, enum rpmh_state state,
 {
         struct batch_cache_req *req;
         struct rpmh_request *rpm_msgs;
-        DECLARE_COMPLETION_ONSTACK(compl);
+        struct completion *compls;
         struct rpmh_ctrlr *ctrlr = get_rpmh_ctrlr(dev);
         unsigned long time_left;
         int count = 0;
-        int ret, i, j;
+        int ret, i;
+        void *ptr;
 
         if (!cmd || !n)
                 return -EINVAL;
@@ -361,10 +363,15 @@ int rpmh_write_batch(const struct device *dev, enum rpmh_state state,
         if (!count)
                 return -EINVAL;
 
-        req = kzalloc(sizeof(*req) + count * sizeof(req->rpm_msgs[0]),
+        ptr = kzalloc(sizeof(*req) +
+                      count * (sizeof(req->rpm_msgs[0]) + sizeof(*compls)),
                       GFP_ATOMIC);
-        if (!req)
+        if (!ptr)
                 return -ENOMEM;
+
+        req = ptr;
+        compls = ptr + sizeof(*req) + count * sizeof(*rpm_msgs);
+
         req->count = count;
         rpm_msgs = req->rpm_msgs;
 
@@ -379,25 +386,26 @@ int rpmh_write_batch(const struct device *dev, enum rpmh_state state,
         }
 
         for (i = 0; i < count; i++) {
-                rpm_msgs[i].completion = &compl;
+                struct completion *compl = &compls[i];
+
+                init_completion(compl);
+                rpm_msgs[i].completion = compl;
                 ret = rpmh_rsc_send_data(ctrlr_to_drv(ctrlr), &rpm_msgs[i].msg);
                 if (ret) {
                         pr_err("Error(%d) sending RPMH message addr=%#x\n",
                                ret, rpm_msgs[i].msg.cmds[0].addr);
-                        for (j = i; j < count; j++)
-                                rpmh_tx_done(&rpm_msgs[j].msg, ret);
                         break;
                 }
         }
 
         time_left = RPMH_TIMEOUT_MS;
-        for (i = 0; i < count; i++) {
-                time_left = wait_for_completion_timeout(&compl, time_left);
+        while (i--) {
+                time_left = wait_for_completion_timeout(&compls[i], time_left);
                 if (!time_left) {
                         /*
                          * Better hope they never finish because they'll signal
-                         * the completion on our stack and that's bad once
-                         * we've returned from the function.
+                         * the completion that we're going to free once
+                         * we've returned from this function.
                          */
                         WARN_ON(1);
                         ret = -ETIMEDOUT;
@@ -406,7 +414,7 @@ int rpmh_write_batch(const struct device *dev, enum rpmh_state state,
         }
 
 exit:
-        kfree(req);
+        kfree(ptr);
 
         return ret;
 }