From bcc9c14c7e80cf39a59833b7e1c17443ef7f9788 Mon Sep 17 00:00:00 2001
From: Florent Kermarrec <florent@enjoy-digital.fr>
Date: Thu, 20 Jan 2022 17:44:35 +0100
Subject: [PATCH] cores/cpu: Fix automatic collection for
 Rocket/BlackParrot/EOS-S3.

---
 litex/soc/cores/cpu/__init__.py             | 7 +++----
 litex/soc/cores/cpu/blackparrot/__init__.py | 2 +-
 litex/soc/cores/cpu/blackparrot/core.py     | 4 ++--
 litex/soc/cores/cpu/rocket/__init__.py      | 2 +-
 litex/soc/cores/cpu/rocket/core.py          | 4 ++--
 5 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/litex/soc/cores/cpu/__init__.py b/litex/soc/cores/cpu/__init__.py
index 611ead38b..583677636 100644
--- a/litex/soc/cores/cpu/__init__.py
+++ b/litex/soc/cores/cpu/__init__.py
@@ -94,11 +94,10 @@ def collect_cpus():
                 continue
 
             # OK, it seems to be a CPU; now get the class and add it to dict.
-            cpu        = file
-            cpu_module = f"{cpu}"
+            cpu = file
             sys.path.append(path)
-            for cpu_name, cpu_cls in inspect.getmembers(importlib.import_module(cpu_module), inspect.isclass):
-                if cpu.replace("_", "") == cpu_name.lower():
+            for cpu_name, cpu_cls in inspect.getmembers(importlib.import_module(cpu), inspect.isclass):
+                if cpu_name.lower() in [cpu, cpu.replace("_", "")]:
                     cpus[cpu] = cpu_cls
 
     # Return collected CPUs.
diff --git a/litex/soc/cores/cpu/blackparrot/__init__.py b/litex/soc/cores/cpu/blackparrot/__init__.py
index 05c9fcee4..ca5482e02 100644
--- a/litex/soc/cores/cpu/blackparrot/__init__.py
+++ b/litex/soc/cores/cpu/blackparrot/__init__.py
@@ -1 +1 @@
-from litex.soc.cores.cpu.blackparrot.core import BlackParrotRV64
+from litex.soc.cores.cpu.blackparrot.core import BlackParrot
diff --git a/litex/soc/cores/cpu/blackparrot/core.py b/litex/soc/cores/cpu/blackparrot/core.py
index a7043ab3c..0f9d3df2e 100644
--- a/litex/soc/cores/cpu/blackparrot/core.py
+++ b/litex/soc/cores/cpu/blackparrot/core.py
@@ -52,9 +52,9 @@ GCC_FLAGS = {
     "sim":      "-march=rv64imafd -mabi=lp64d ",
 }
 
-# BlackParrotRV64 ----------------------------------------------------------------------------------
+# BlackParrot --------------------------------------------------------------------------------------
 
-class BlackParrotRV64(CPU):
+class BlackParrot(CPU):
     family               = "riscv"
     name                 = "blackparrot"
     human_name           = "BlackParrotRV64[imafd]"
diff --git a/litex/soc/cores/cpu/rocket/__init__.py b/litex/soc/cores/cpu/rocket/__init__.py
index b9a906106..426a2bced 100644
--- a/litex/soc/cores/cpu/rocket/__init__.py
+++ b/litex/soc/cores/cpu/rocket/__init__.py
@@ -1 +1 @@
-from litex.soc.cores.cpu.rocket.core import RocketRV64
+from litex.soc.cores.cpu.rocket.core import Rocket
diff --git a/litex/soc/cores/cpu/rocket/core.py b/litex/soc/cores/cpu/rocket/core.py
index 03b79e744..d7785cc67 100644
--- a/litex/soc/cores/cpu/rocket/core.py
+++ b/litex/soc/cores/cpu/rocket/core.py
@@ -86,9 +86,9 @@ CPU_SIZE_PARAMS = {
     "full4q":   (   256,      64,         4),
 }
 
-# Rocket RV64 --------------------------------------------------------------------------------------
+# Rocket  ------------------------------------------------------------------------------------------
 
-class RocketRV64(CPU):
+class Rocket(CPU):
     family               = "riscv"
     name                 = "rocket"
     human_name           = "RocketRV64[imac]"