init

lib_audio_dsp/python/audio_dsp/design/thread.py

@@ -0,0 +1,88 @@
+# Copyright 2024-2025 XMOS LIMITED.
+# This Software is subject to the terms of the XMOS Public Licence: Version 1.
+"""Contains classes for adding a thread to the DSP pipeline."""
+
+from typing import Type
+
+from audio_dsp.design.composite_stage import CompositeStage, _StageOrComposite
+from audio_dsp.design.stage import Stage, StageOutputList, find_config
+
+
+class DSPThreadStage(Stage):
+    """
+    Stage for the DSP thread. Does not support processing of data through it. Only
+    used for DSP thread level control commands, for example, querying the max cycles
+    consumed by the thread.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(config=find_config("dsp_thread"), **kwargs)
+        self.create_outputs(0)
+
+    def add_to_dot(self, dot):
+        """
+        Exclude this stage from the dot diagram.
+
+        Parameters
+        ----------
+        dot : graphviz.Diagraph
+            dot instance to add edges to.
+        """
+
+
+class Thread(CompositeStage):
+    """
+    A composite stage used to represent a thread in the pipeline. Create
+    using Pipeline.thread rather than instantiating directly.
+
+    Parameters
+    ----------
+    id : int
+        Thread index
+    kwargs : dict
+        forwarded to __init__ of CompositeStage
+
+    Attributes
+    ----------
+    id : int
+        Thread index
+    thread_stage : Stage
+        DSPThreadStage stage
+    """
+
+    def __init__(self, id: int, **kwargs):
+        super().__init__(name=f"Thread {id}", **kwargs)
+        self.id = id
+        self.thread_stage = None
+
+    def __enter__(self):
+        """Support for context manager."""
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Support for context manager."""
+        ...
+
+    def add_thread_stage(self):
+        """Add to this thread the stage which manages thread level commands."""
+        self.thread_stage = self.stage(DSPThreadStage, StageOutputList(), label=f"thread{self.id}")
+
+    def stage(
+        self,
+        stage_type: Type[_StageOrComposite],
+        inputs: StageOutputList,
+        label: str | None = None,
+        **kwargs,
+    ) -> _StageOrComposite:
+        """
+        Create a new stage or composite stage and
+        register it with this composite stage.
+
+        This is a wrapper around :func:`audio_dsp.design.composite_stage.stage`
+        but adds a thread crossing counter.
+        """
+        for i in inputs.edges:
+            if i:
+                i.crossings.append(self.id)
+        stage = super().stage(stage_type, inputs, label, **kwargs)
+        return stage