[PATCH 2/3] firmware: qcom-scm: Support multiple waitq contexts

[Unnathi Chalicheemala <quic_uchalich@quicinc.com>]

Currently, only a single waitqueue context is supported, with waitqueue
id zero. SM8650 firmware now supports multiple waitqueue contexts, so
add support to dynamically create and support as many unique waitqueue
contexts as firmware returns to the driver.
Unique waitqueue contexts are supported using xarray to create a
hash table for associating a unique wq_ctx with a struct completion
variable for easy lookup.
The waitqueue ids can be >=0 as now we have more than one waitqueue
context.

Signed-off-by: Unnathi Chalicheemala <quic_uchalich@quicinc.com>
---
 drivers/firmware/qcom/qcom_scm-smc.c |  7 +++-
 drivers/firmware/qcom/qcom_scm.c     | 77 ++++++++++++++++++++++++++----------
 drivers/firmware/qcom/qcom_scm.h     |  3 +-
 3 files changed, 64 insertions(+), 23 deletions(-)

diff --git a/drivers/firmware/qcom/qcom_scm-smc.c b/drivers/firmware/qcom/qcom_scm-smc.c
index 16cf88acfa8e..80083e3615b1 100644
--- a/drivers/firmware/qcom/qcom_scm-smc.c
+++ b/drivers/firmware/qcom/qcom_scm-smc.c
@@ -103,7 +103,12 @@ static int __scm_smc_do_quirk_handle_waitq(struct device *dev, struct arm_smccc_
 			wq_ctx = res->a1;
 			smc_call_ctx = res->a2;
 
-			ret = qcom_scm_wait_for_wq_completion(wq_ctx);
+			if (!dev) {
+				/* Protect the dev_get_drvdata() call that follows */
+				return -EPROBE_DEFER;
+			}
+
+			ret = qcom_scm_wait_for_wq_completion(dev_get_drvdata(dev), wq_ctx);
 			if (ret)
 				return ret;
 
diff --git a/drivers/firmware/qcom/qcom_scm.c b/drivers/firmware/qcom/qcom_scm.c
index c1be8270ead1..4606c49ef155 100644
--- a/drivers/firmware/qcom/qcom_scm.c
+++ b/drivers/firmware/qcom/qcom_scm.c
@@ -21,6 +21,7 @@
 #include <linux/platform_device.h>
 #include <linux/reset-controller.h>
 #include <linux/types.h>
+#include <linux/xarray.h>
 
 #include "qcom_scm.h"
 
@@ -33,7 +34,7 @@ struct qcom_scm {
 	struct clk *iface_clk;
 	struct clk *bus_clk;
 	struct icc_path *path;
-	struct completion waitq_comp;
+	struct xarray waitq;
 	struct reset_controller_dev reset;
 
 	/* control access to the interconnect path */
@@ -1742,42 +1743,74 @@ bool qcom_scm_is_available(void)
 }
 EXPORT_SYMBOL_GPL(qcom_scm_is_available);
 
-static int qcom_scm_assert_valid_wq_ctx(u32 wq_ctx)
+static struct completion *qcom_scm_get_completion(struct qcom_scm *scm, u32 wq_ctx)
 {
-	/* FW currently only supports a single wq_ctx (zero).
-	 * TODO: Update this logic to include dynamic allocation and lookup of
-	 * completion structs when FW supports more wq_ctx values.
+	struct completion *wq;
+	struct completion *old;
+	int err;
+
+	wq = xa_load(&scm->waitq, wq_ctx);
+	if (wq) {
+		/*
+		 * Valid struct completion *wq found corresponding to
+		 * given wq_ctx. We're done here.
+		 */
+		goto out;
+	}
+
+	/*
+	 * If a struct completion *wq does not exist for wq_ctx, create it. FW
+	 * only uses a finite number of wq_ctx values, so we will be reaching
+	 * here only a few times right at the beginning of the device's uptime
+	 * and then early-exit from idr_find() above subsequently.
 	 */
-	if (wq_ctx != 0) {
-		dev_err(__scm->dev, "Firmware unexpectedly passed non-zero wq_ctx\n");
-		return -EINVAL;
+	wq = kzalloc(sizeof(*wq), GFP_ATOMIC);
+	if (!wq) {
+		wq = ERR_PTR(-ENOMEM);
+		goto out;
 	}
 
-	return 0;
+	init_completion(wq);
+
+	old = xa_store(&scm->waitq, wq_ctx, wq, GFP_ATOMIC);
+	err = xa_err(old);
+	if (err) {
+		kfree(wq);
+		wq = ERR_PTR(err);
+	}
+
+out:
+	return wq;
 }
 
-int qcom_scm_wait_for_wq_completion(u32 wq_ctx)
+int qcom_scm_wait_for_wq_completion(struct qcom_scm *scm, u32 wq_ctx)
 {
-	int ret;
+	struct completion *wq;
 
-	ret = qcom_scm_assert_valid_wq_ctx(wq_ctx);
-	if (ret)
-		return ret;
+	wq = qcom_scm_get_completion(scm, wq_ctx);
+	if (IS_ERR(wq)) {
+		pr_err("Unable to wait on invalid waitqueue for wq_ctx %d: %ld\n",
+						wq_ctx, PTR_ERR(wq));
+		return PTR_ERR(wq);
+	}
 
-	wait_for_completion(&__scm->waitq_comp);
+	wait_for_completion(wq);
 
 	return 0;
 }
 
 static int qcom_scm_waitq_wakeup(struct qcom_scm *scm, unsigned int wq_ctx)
 {
-	int ret;
+	struct completion *wq;
 
-	ret = qcom_scm_assert_valid_wq_ctx(wq_ctx);
-	if (ret)
-		return ret;
+	wq = qcom_scm_get_completion(scm, wq_ctx);
+	if (IS_ERR(wq)) {
+		pr_err("Unable to wake up invalid waitqueue for wq_ctx %d: %ld\n",
+						wq_ctx, PTR_ERR(wq));
+		return PTR_ERR(wq);
+	}
 
-	complete(&__scm->waitq_comp);
+	complete(wq);
 
 	return 0;
 }
@@ -1854,7 +1887,9 @@ static int qcom_scm_probe(struct platform_device *pdev)
 	if (ret)
 		return ret;
 
-	init_completion(&scm->waitq_comp);
+	platform_set_drvdata(pdev, scm);
+
+	xa_init(&scm->waitq);
 
 	__scm = scm;
 	__scm->dev = &pdev->dev;
diff --git a/drivers/firmware/qcom/qcom_scm.h b/drivers/firmware/qcom/qcom_scm.h
index 4532907e8489..d54df5a2b690 100644
--- a/drivers/firmware/qcom/qcom_scm.h
+++ b/drivers/firmware/qcom/qcom_scm.h
@@ -62,7 +62,8 @@ struct qcom_scm_res {
 	u64 result[MAX_QCOM_SCM_RETS];
 };
 
-int qcom_scm_wait_for_wq_completion(u32 wq_ctx);
+struct qcom_scm;
+int qcom_scm_wait_for_wq_completion(struct qcom_scm *scm, u32 wq_ctx);
 int scm_get_wq_ctx(u32 *wq_ctx, u32 *flags, u32 *more_pending);
 
 #define SCM_SMC_FNID(s, c)	((((s) & 0xFF) << 8) | ((c) & 0xFF))

-- 
2.25.1