From 363afb5fc9a8b60f521a5fafbaba4b9958dee05c Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Wed, 18 Oct 2023 15:12:56 +0800
Subject: [PATCH] artiq_ddb_template: add support for user LEDs

Add support for additional user LEDs.
---
 artiq/frontend/artiq_ddb_template.py | 40 +++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 7 deletions(-)

diff --git a/artiq/frontend/artiq_ddb_template.py b/artiq/frontend/artiq_ddb_template.py
index 5038563d6..eeca02c8e 100755
--- a/artiq/frontend/artiq_ddb_template.py
+++ b/artiq/frontend/artiq_ddb_template.py
@@ -24,6 +24,27 @@ def get_cpu_target(description):
     else:
         raise NotImplementedError
 
+
+def get_num_leds(description):
+    drtio_role = description["drtio_role"]
+    target = description["target"]
+    hw_rev = description["hw_rev"]
+    kasli_board_leds = {
+        "v1.0": 4,
+        "v1.1": 6,
+        "v2.0": 3
+    }
+    if target == "kasli":
+        if hw_rev in ("v1.0", "v1.1") and drtio_role != "standalone":
+            # LEDs are used for DRTIO status on v1.0 and v1.1
+            return kasli_board_leds[hw_rev] - 3
+        return kasli_board_leds[hw_rev]
+    elif target == "kasli_soc":
+        return 2
+    else:
+        raise ValueError
+     
+
 def process_header(output, description):
     print(textwrap.dedent("""
         # Autogenerated for the {variant} variant
@@ -701,8 +722,8 @@ class PeripheralManager:
         processor = getattr(self, "process_"+str(peripheral["type"]))
         return processor(rtio_offset, peripheral)
 
-    def add_board_leds(self, rtio_offset, board_name=None):
-        for i in range(2):
+    def add_board_leds(self, rtio_offset, board_name=None, num_leds=2):
+        for i in range(num_leds):
             if board_name is None:
                 led_name = self.get_name("led")
             else:
@@ -716,7 +737,7 @@ class PeripheralManager:
                 }}""",
                 name=led_name,
                 channel=rtio_offset+i)
-        return 2
+        return num_leds
 
 
 def split_drtio_eem(peripherals):
@@ -745,9 +766,10 @@ def process(output, primary_description, satellites):
     for peripheral in local_peripherals:
         n_channels = pm.process(rtio_offset, peripheral)
         rtio_offset += n_channels
-    if drtio_role == "standalone":
-        n_channels = pm.add_board_leds(rtio_offset)
-        rtio_offset += n_channels
+
+    num_leds = get_num_leds(primary_description)
+    pm.add_board_leds(rtio_offset, num_leds=num_leds)
+    rtio_offset += num_leds
 
     for destination, description in satellites:
         if description["drtio_role"] != "satellite":
@@ -766,7 +788,11 @@ def process(output, primary_description, satellites):
         for peripheral in peripherals:
             n_channels = pm.process(rtio_offset, peripheral)
             rtio_offset += n_channels
-    
+
+        num_leds = get_num_leds(description)
+        pm.add_board_leds(rtio_offset, num_leds=num_leds)
+        rtio_offset += num_leds
+
     for i, peripheral in enumerate(drtio_peripherals):
         if not("drtio_destination" in peripheral):
             if primary_description["target"] == "kasli":