trusty: benchmark: Support CPU pinning in the kernel.

Bug: 315443297
Test: build.py qemu-generic-arm64-test-debug  --test com.android.kernel.pacbench
Change-Id: Ib46baa4fc2cd84566b68e4b665dc66c6002a90d1

1 files changed, 27 insertions, 2 deletions

diff --git a/include/shared/lk/trusty_benchmark.h b/include/shared/lk/trusty_benchmark.h
index 8e13fcf..62dada7 100644
--- a/include/shared/lk/trusty_benchmark.h
+++ b/include/shared/lk/trusty_benchmark.h
@@ -132,9 +132,13 @@
 #include "trusty_bench_option_cb.h"
 #include "trusty_bench_print_tables.h"
 #include "trusty_unittest.h"
+#ifdef TRUSTY_USERSPACE
 #ifdef WITH_PTHREAD
 #include <lib/thread/pthread.h>
 #endif
+#elif WITH_SMP
+#include <kernel/mp.h>
+#endif
 #include <uapi/err.h>
 
 __BEGIN_CDECLS
@@ -181,8 +185,7 @@ static inline size_t bench_get_cpu_idx(void) {
 /*
  * Helper macros to run on tests on all CPUs
  */
-
-#ifdef WITH_PTHREAD
+#if defined(TRUSTY_USERSPACE) && defined(WITH_PTHREAD)
 static int trusty_bench_multi_cpus_setup(void) {
     if (trusty_bench_nb_cpu > 1) {
         cpu_set_t cpu_set;
@@ -211,10 +214,32 @@ static int trusty_bench_multi_cpus_teardown(void) {
     }
     return NO_ERROR;
 }
+#elif !defined(TRUSTY_USERSPACE) && WITH_SMP
+static int trusty_bench_multi_cpus_setup(void) {
+    if (trusty_bench_nb_cpu > 1) {
+        const int cpu = bench_state.cur_param_idx / trusty_cur_bench_nb_params;
+
+        if (cpu < SMP_MAX_CPUS && mp_is_cpu_active(cpu)) {
+            thread_set_pinned_cpu(get_current_thread(), cpu);
+        } else {
+            return EINVAL;
+        }
+    }
+
+    return NO_ERROR;
+}
+
+static int trusty_bench_multi_cpus_teardown(void) {
+    if (trusty_bench_nb_cpu > 1) {
+        thread_set_pinned_cpu(get_current_thread(), -1);
+    }
+    return NO_ERROR;
+}
 #else
 static int trusty_bench_multi_cpus_setup(void) {
     return NO_ERROR;
 }
+
 static int trusty_bench_multi_cpus_teardown(void) {
     return NO_ERROR;
 }