From 9b6683e61fb724dee51eeea5456f18c61ec363c4 Mon Sep 17 00:00:00 2001
From: lijuan <juan.li@spacemit.com>
Date: Tue, 24 Dec 2024 20:48:28 +0800
Subject: [PATCH] ai-trap: scan device-tree when booting time

Change-Id: I3075960b7d3a20262b4979816e9521f4c4f169c2
Signed-off-by: lijuan <juan.li@spacemit.com>
---
 arch/riscv/include/asm/processor.h | 3 ++-
 arch/riscv/kernel/cpufeature.c     | 9 ++++-----
 arch/riscv/kernel/traps.c          | 4 +---
 kernel/smp.c                       | 4 ++++
 4 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/arch/riscv/include/asm/processor.h b/arch/riscv/include/asm/processor.h
index af6962ccfc2f..d19c24adac97 100644
--- a/arch/riscv/include/asm/processor.h
+++ b/arch/riscv/include/asm/processor.h
@@ -165,7 +165,8 @@ int riscv_of_parent_hartid(struct device_node *node, unsigned long *hartid);
 
 extern void riscv_fill_hwcap(void);
 extern int arch_dup_task_struct(struct task_struct *dst, struct task_struct *src);
-extern struct cpumask ai_core_mask_get(void);
+extern void ai_core_mask_get(void);
+extern struct cpumask	ai_cpu_mask;
 
 extern unsigned long signal_minsigstksz __ro_after_init;
 
diff --git a/arch/riscv/kernel/cpufeature.c b/arch/riscv/kernel/cpufeature.c
index 0721e6c767e5..be6401842a3a 100644
--- a/arch/riscv/kernel/cpufeature.c
+++ b/arch/riscv/kernel/cpufeature.c
@@ -72,15 +72,15 @@ bool __riscv_isa_extension_available(const unsigned long *isa_bitmap, unsigned i
 }
 EXPORT_SYMBOL_GPL(__riscv_isa_extension_available);
 
-struct cpumask ai_core_mask_get(void)
+struct cpumask	ai_cpu_mask;
+void ai_core_mask_get(void)
 {
 	struct device_node *node;
 	const char *cpu_ai;
-	struct cpumask	cpu_mask;
 	unsigned long hartid;
 	int rc;
 
-	cpumask_clear(&cpu_mask);
+	cpumask_clear(&ai_cpu_mask);
 
 	for_each_of_cpu_node(node) {
 		rc = riscv_of_processor_hartid(node, &hartid);
@@ -92,11 +92,10 @@ struct cpumask ai_core_mask_get(void)
 		}
 
 		if(!strcmp(cpu_ai, "true")) {
-			cpumask_set_cpu(hartid, &cpu_mask);
+			cpumask_set_cpu(hartid, &ai_cpu_mask);
 		}
 	}
 
-	return cpu_mask;
 }
 
 static int riscv_ext_zicbom_validate(const struct riscv_isa_ext_data *data,
diff --git a/arch/riscv/kernel/traps.c b/arch/riscv/kernel/traps.c
index 67f2f81fd6c4..865f191ab428 100644
--- a/arch/riscv/kernel/traps.c
+++ b/arch/riscv/kernel/traps.c
@@ -174,10 +174,8 @@ asmlinkage __visible __trap_section void do_trap_insn_illegal(struct pt_regs *re
 		__get_user(epc, (u32 __user *)regs->epc);
 		if ((epc & AI_OPCODE_MASK0) == AI_OPCODE_MATCH0 ||
 			(epc & AI_OPCODE_MASK1) == AI_OPCODE_MATCH1) {
-			struct cpumask mask;
-			mask = ai_core_mask_get();
 			local_irq_enable();
-			sched_setaffinity(current->pid, &mask);
+			sched_setaffinity(current->pid, &ai_cpu_mask);
 			local_irq_disable();
 			irqentry_exit_to_user_mode(regs);
 			return;
diff --git a/kernel/smp.c b/kernel/smp.c
index 3eeffeaf5450..77cb94c8c879 100644
--- a/kernel/smp.c
+++ b/kernel/smp.c
@@ -987,6 +987,10 @@ void __init smp_init(void)
 
 	/* Any cleanup work */
 	smp_cpus_done(setup_max_cpus);
+
+#ifdef CONFIG_BIND_THREAD_TO_AICORES
+	ai_core_mask_get();
+#endif
 }
 
 /*