From dd0f21c560f7139fc0c30bea1c264d8cb09144fe Mon Sep 17 00:00:00 2001
From: Anup Patel <anup.patel@wdc.com>
Date: Thu, 26 Mar 2020 18:35:38 +0530
Subject: [PATCH] lib: sbi_scratch: Introduce sbi_scratch_last_hartid() API

The patch adds sbi_scratch_last_hartid() API which returns
last HART id having a scratch space. We can use this new API
to optimize places where we iterate over HART id from 0 to
SBI_HARTMASK_MAX_BITS.

Signed-off-by: Anup Patel <anup.patel@wdc.com>
Reviewed-by: Atish Patra <atish.patra@wdc.com>
---
 include/sbi/sbi_scratch.h | 6 ++++++
 lib/sbi/sbi_hsm.c         | 4 ++--
 lib/sbi/sbi_init.c        | 2 +-
 lib/sbi/sbi_scratch.c     | 5 ++++-
 4 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/include/sbi/sbi_scratch.h b/include/sbi/sbi_scratch.h
index e848900..8086207 100644
--- a/include/sbi/sbi_scratch.h
+++ b/include/sbi/sbi_scratch.h
@@ -113,6 +113,12 @@ extern struct sbi_scratch *hartid_to_scratch_table[];
 #define sbi_hartid_to_scratch(__hartid) \
 	hartid_to_scratch_table[__hartid]
 
+/** Last HART id having a sbi_scratch pointer */
+extern u32 last_hartid_having_scratch;
+
+/** Get last HART id having a sbi_scratch pointer */
+#define sbi_scratch_last_hartid()	last_hartid_having_scratch
+
 #endif
 
 #endif
diff --git a/lib/sbi/sbi_hsm.c b/lib/sbi/sbi_hsm.c
index e7b2f04..4330a22 100644
--- a/lib/sbi/sbi_hsm.c
+++ b/lib/sbi/sbi_hsm.c
@@ -88,7 +88,7 @@ bool sbi_hsm_hart_started(u32 hartid)
 int sbi_hsm_hart_started_mask(ulong hbase, ulong *out_hmask)
 {
 	ulong i;
-	ulong hcount = SBI_HARTMASK_MAX_BITS;
+	ulong hcount = sbi_scratch_last_hartid() + 1;
 
 	*out_hmask = 0;
 	if (hcount <= hbase)
@@ -153,7 +153,7 @@ int sbi_hsm_init(struct sbi_scratch *scratch, u32 hartid, bool cold_boot)
 			return SBI_ENOMEM;
 
 		/* Initialize hart state data for every hart */
-		for (i = 0; i < SBI_HARTMASK_MAX_BITS; i++) {
+		for (i = 0; i <= sbi_scratch_last_hartid(); i++) {
 			rscratch = sbi_hartid_to_scratch(i);
 			if (!rscratch)
 				continue;
diff --git a/lib/sbi/sbi_init.c b/lib/sbi/sbi_init.c
index f0ce209..f448b79 100644
--- a/lib/sbi/sbi_init.c
+++ b/lib/sbi/sbi_init.c
@@ -128,7 +128,7 @@ static void wake_coldboot_harts(struct sbi_scratch *scratch, u32 hartid)
 	coldboot_done = 1;
 
 	/* Send an IPI to all HARTs waiting for coldboot */
-	for (int i = 0; i < SBI_HARTMASK_MAX_BITS; i++) {
+	for (int i = 0; i <= sbi_scratch_last_hartid(); i++) {
 		if ((i != hartid) &&
 		    sbi_hartmask_test_hart(i, &coldboot_wait_hmask))
 			sbi_platform_ipi_send(plat, i);
diff --git a/lib/sbi/sbi_scratch.c b/lib/sbi/sbi_scratch.c
index 8b63570..96bae5b 100644
--- a/lib/sbi/sbi_scratch.c
+++ b/lib/sbi/sbi_scratch.c
@@ -14,6 +14,7 @@
 #include <sbi/sbi_scratch.h>
 #include <sbi/sbi_string.h>
 
+u32 last_hartid_having_scratch = SBI_HARTMASK_MAX_BITS;
 struct sbi_scratch *hartid_to_scratch_table[SBI_HARTMASK_MAX_BITS] = { 0 };
 
 static spinlock_t extra_lock = SPIN_LOCK_INITIALIZER;
@@ -32,6 +33,8 @@ int sbi_scratch_init(struct sbi_scratch *scratch)
 		hartid_to_scratch_table[i] =
 			((hartid2scratch)scratch->hartid_to_scratch)(i,
 					sbi_platform_hart_index(plat, i));
+		if (hartid_to_scratch_table[i])
+			last_hartid_having_scratch = i;
 	}
 
 	return 0;
@@ -71,7 +74,7 @@ done:
 	spin_unlock(&extra_lock);
 
 	if (ret) {
-		for (i = 0; i < SBI_HARTMASK_MAX_BITS; i++) {
+		for (i = 0; i < sbi_scratch_last_hartid(); i++) {
 			rscratch = sbi_hartid_to_scratch(i);
 			if (!rscratch)
 				continue;