Merge pull request #1530 from SpiNNakerManchester/ColouredApplication…

…Vertex

Coloured application vertex

spynnaker/pyNN/models/abstract_models/__init__.py

@@ -17,6 +17,7 @@
 from .abstract_max_spikes import AbstractMaxSpikes
 from .abstract_synapse_expandable import (
     AbstractSynapseExpandable, SYNAPSE_EXPANDER_APLX)
+from .coloured_application_vertex import ColouredApplicationVertex
 from .sends_synaptic_inputs_over_sdram import SendsSynapticInputsOverSDRAM
 from .receives_synaptic_inputs_over_sdram import (
     ReceivesSynapticInputsOverSDRAM)
@@ -28,6 +29,7 @@
 
 __all__ = ["AbstractAcceptsIncomingSynapses", "AbstractHasDelayStages",
            "AbstractMaxSpikes", "AbstractSynapseExpandable",
+           "ColouredApplicationVertex",
            "SYNAPSE_EXPANDER_APLX",
            "SendsSynapticInputsOverSDRAM", "ReceivesSynapticInputsOverSDRAM",
            "HasSynapses", "AbstractNeuronExpandable", "NEURON_EXPANDER_APLX",

spynnaker/pyNN/models/abstract_models/coloured_application_vertex.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025 The University of Manchester
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from spinn_utilities.abstract_base import abstractmethod
+from pacman.model.graphs.application import ApplicationVertex
+
+
+class ColouredApplicationVertex(ApplicationVertex):
+    """
+    Indicates an application vertex that implements n_colour_bits.
+    """
+
+    __slots__ = ()
+
+    @property
+    @abstractmethod
+    def n_colour_bits(self) -> int:
+        """
+        The number of colour bits sent by this vertex.
+
+        Assumed 0 unless overridden
+
+        :rtype: int
+        """
+        raise NotImplementedError

spynnaker/pyNN/models/common/local_only_2d_common.py

@@ -13,13 +13,24 @@
 # limitations under the License.
 from math import ceil, log2, floor
 from collections import namedtuple, defaultdict
-from pacman.model.graphs.application import ApplicationVirtualVertex
+from typing import Dict, List, Tuple, TYPE_CHECKING, Union
+
+from pacman.model.graphs.application import (
+    ApplicationVertex, ApplicationVirtualVertex)
 from pacman.model.graphs.common.slice import Slice
 from pacman.model.graphs.common.mdslice import MDSlice
+from pacman.model.routing_info import AppVertexRoutingInfo
+
 from spinn_front_end_common.utilities.constants import BYTES_PER_WORD
+
 from spynnaker.pyNN.data.spynnaker_data_view import SpynnakerDataView
+from spynnaker.pyNN.models.abstract_models import ColouredApplicationVertex
 from spynnaker.pyNN.utilities.utility_calls import get_n_bits
 
+if TYPE_CHECKING:
+    from spynnaker.pyNN.models.neuron import AbstractPopulationVertex
+    from spynnaker.pyNN.models.projection import Projection
+
 #: The number of bits in a short value
 BITS_PER_SHORT = 16
 
@@ -37,7 +48,7 @@
     "Source", ["projection", "local_delay", "delay_stage"])
 
 
-def get_div_const(value):
+def get_div_const(value: int) -> int:
     """ Get the values used to perform fast division by an integer constant
 
     :param int value: The value to be divided by
@@ -53,13 +64,13 @@ def get_div_const(value):
             + (sh1 << BITS_PER_SHORT) + m)
 
 
-def get_delay_for_source(incoming):
+def get_delay_for_source(incoming: "Projection") -> Tuple[
+        ColouredApplicationVertex, int, int, str]:
     """ Get the vertex which will send data from a given source projection,
         along with the delay stage and locally-handled delay value
 
     :param Projection incoming: The incoming projection to get the delay from
     :return: The vertex, the local delay, the delay stage, the partition id
-    :rtype: ApplicationVertex, int, int, str
     """
     # pylint: disable=protected-access
     app_edge = incoming._projection_edge
@@ -69,14 +80,18 @@ def get_delay_for_source(incoming):
     max_delay = app_edge.post_vertex.splitter.max_support_delay()
     local_delay = steps % max_delay
     delay_stage = 0
-    pre_vertex = app_edge.pre_vertex
+    pre_vertex: ColouredApplicationVertex = app_edge.pre_vertex
     if steps > max_delay:
         delay_stage = (steps // max_delay) - 1
-        pre_vertex = app_edge.delay_edge.pre_vertex
+        delay_edge = app_edge.delay_edge
+        assert delay_edge is not None
+        pre_vertex = delay_edge.pre_vertex
     return pre_vertex, local_delay, delay_stage, s_info.partition_id
 
 
-def get_rinfo_for_spike_source(pre_vertex, partition_id):
+def get_rinfo_for_spike_source(
+        pre_vertex: ApplicationVertex,
+        partition_id: str) -> Tuple[AppVertexRoutingInfo, int, int]:
     """
     Get the routing information for the source of a projection in the
     given partition.
@@ -91,6 +106,7 @@ def get_rinfo_for_spike_source(pre_vertex, partition_id):
     # Find the routing information
     r_info = routing_info.get_info_from(
             pre_vertex, partition_id)
+    assert isinstance(r_info, AppVertexRoutingInfo)
 
     n_cores = len(r_info.vertex.splitter.get_out_going_vertices(partition_id))
 
@@ -104,7 +120,8 @@ def get_rinfo_for_spike_source(pre_vertex, partition_id):
     return r_info, core_mask, mask_shift
 
 
-def get_sources_for_target(app_vertex):
+def get_sources_for_target(app_vertex: "AbstractPopulationVertex") -> Dict[
+        Tuple[ColouredApplicationVertex, str], List[Source]]:
     """
     Get all the application vertex sources that will hit the given application
     vertex.
@@ -123,7 +140,8 @@ def get_sources_for_target(app_vertex):
     return sources
 
 
-def get_first_and_last_slice(pre_vertex):
+def get_first_and_last_slice(pre_vertex: ApplicationVertex) -> \
+        Union[Tuple[Slice, Slice], Tuple[MDSlice, MDSlice]]:
     """
     Get the first and last slice of an application vertex.
 

spynnaker/pyNN/models/common/population_application_vertex.py

@@ -12,24 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import Collection, Container, Iterable, List, Optional, Tuple
+
 from spinn_utilities.overrides import overrides
 from spinn_utilities.ranged.abstract_sized import Selector
-from pacman.model.graphs.application import ApplicationVertex
+
 from pacman.model.graphs.machine import MachineVertex
 from pacman.model.graphs.common import Slice
 from pacman.model.routing_info import RoutingInfo
 from pacman.utilities.utility_calls import get_keys
+
 from spinn_front_end_common.utilities.exceptions import ConfigurationException
 from spinn_front_end_common.abstract_models import HasCustomAtomKeyMap
 from spinn_front_end_common.interface.ds import DataType
+
+from spynnaker.pyNN.models.abstract_models import ColouredApplicationVertex
 from spynnaker.pyNN.models.common.parameter_holder import ParameterHolder
 from spynnaker.pyNN.utilities.buffer_data_type import BufferDataType
 from spynnaker.pyNN.models.current_sources import AbstractCurrentSource
 from .types import Names, Values
 
 
 # pylint: disable=abstract-method
-class PopulationApplicationVertex(ApplicationVertex, HasCustomAtomKeyMap):
+class PopulationApplicationVertex(
+        ColouredApplicationVertex, HasCustomAtomKeyMap):
     """
     A vertex that can be used in a Population.
 
@@ -410,14 +415,8 @@ def inject(self, current_source: AbstractCurrentSource,
             "This Population doesn't support injection")
 
     @property
+    @overrides(ColouredApplicationVertex.n_colour_bits)
     def n_colour_bits(self) -> int:
-        """
-        The number of colour bits sent by this vertex.
-
-        Assumed 0 unless overridden
-
-        :rtype: int
-        """
         return 0
 
     @overrides(HasCustomAtomKeyMap.get_atom_key_map)

spynnaker/pyNN/models/neuron/local_only/local_only_convolution.py

@@ -30,11 +30,11 @@
     BYTES_PER_SHORT, BYTES_PER_WORD)
 
 from spynnaker.pyNN.exceptions import SynapticConfigurationException
+from spynnaker.pyNN.models.abstract_models import ColouredApplicationVertex
 from spynnaker.pyNN.models.neural_projections.connectors import (
     ConvolutionConnector, AbstractConnector)
 from spynnaker.pyNN.models.neuron.synapse_dynamics import (
     AbstractSupportsSignedWeights)
-from spynnaker.pyNN.models.common import PopulationApplicationVertex
 from spynnaker.pyNN.types import Weight_Delay_In_Types
 from spynnaker.pyNN.models.common.local_only_2d_common import (
     get_div_const, get_rinfo_for_spike_source, get_sources_for_target,
@@ -78,7 +78,7 @@ def __init__(self, delay: Weight_Delay_In_Types = None):
 
         # Store the sources to avoid recalculation
         self.__cached_sources: Dict[ApplicationVertex, Dict[
-                Tuple[PopulationApplicationVertex, str],
+                Tuple[ColouredApplicationVertex, str],
                 List[Source]]] = dict()
 
     @property
@@ -250,7 +250,7 @@ def write_parameters(
 
     def __get_sources_for_target(
             self, app_vertex: AbstractPopulationVertex) -> Dict[
-                Tuple[PopulationApplicationVertex, str], List[Source]]:
+                Tuple[ColouredApplicationVertex, str], List[Source]]:
         """
         Get all the application vertex sources that will hit the given
         application vertex.

spynnaker/pyNN/models/neuron/local_only/local_only_pool_dense.py

@@ -29,6 +29,7 @@
 from spinn_front_end_common.utilities.constants import BYTES_PER_WORD
 
 from spynnaker.pyNN.exceptions import SynapticConfigurationException
+from spynnaker.pyNN.models.abstract_models import ColouredApplicationVertex
 from spynnaker.pyNN.models.neural_projections.connectors import (
     PoolDenseConnector)
 from spynnaker.pyNN.models.neuron.synapse_dynamics import (
@@ -74,7 +75,7 @@ def __init__(self, delay: Weight_Delay_In_Types = None):
         """
         # Store the sources to avoid recalculation
         self.__cached_sources: Dict[ApplicationVertex, Dict[
-                Tuple[ApplicationVertex, str], List[Source]]] = dict()
+                Tuple[ColouredApplicationVertex, str], List[Source]]] = dict()
 
         super().__init__(delay)
         if not isinstance(self.delay, (float, int)):
@@ -158,7 +159,8 @@ def write_parameters(
             first_conn_index = len(connector_data)
             for source in source_infos:
                 # pylint: disable=protected-access
-                conn = source.projection._synapse_information.connector
+                conn = cast(PoolDenseConnector,
+                            source.projection._synapse_information.connector)
                 app_edge = source.projection._projection_edge
                 connector_data.append(conn.get_local_only_data(
                     app_edge, source.local_delay, source.delay_stage,
@@ -224,7 +226,9 @@ def write_parameters(
         spec.write_array(numpy.array(source_data, dtype=numpy.uint32))
         spec.write_array(numpy.concatenate(connector_data))
 
-    def __get_sources_for_target(self, app_vertex: AbstractPopulationVertex):
+    def __get_sources_for_target(
+            self, app_vertex: AbstractPopulationVertex) -> Dict[
+            Tuple[ColouredApplicationVertex, str], List[Source]]:
         """
         Get all the application vertex sources that will hit the given
         application vertex.

spynnaker/pyNN/models/utility_models/delays/delay_extension_vertex.py

@@ -18,6 +18,8 @@
 from pacman.model.graphs.application import (
     ApplicationEdgePartition, ApplicationVertex)
 from spinn_front_end_common.utilities.constants import BYTES_PER_WORD
+
+from spynnaker.pyNN.models.abstract_models import ColouredApplicationVertex
 from spynnaker.pyNN.exceptions import DelayExtensionException
 from spynnaker.pyNN.models.abstract_models import AbstractHasDelayStages
 from spynnaker.pyNN.utilities.constants import POP_TABLE_MAX_ROW_LENGTH
@@ -29,7 +31,7 @@
 _DELAY_PARAM_HEADER_WORDS = 9
 
 
-class DelayExtensionVertex(ApplicationVertex, AbstractHasDelayStages):
+class DelayExtensionVertex(ColouredApplicationVertex, AbstractHasDelayStages):
     """
     Provide delays to incoming spikes in multiples of the maximum delays
     of a neuron (typically 16 or 32).
@@ -86,7 +88,7 @@ def n_atoms(self) -> int:
         return self.__partition.pre_vertex.n_atoms
 
     @property
-    @overrides(ApplicationVertex.atoms_shape)
+    @overrides(ColouredApplicationVertex.atoms_shape)
     def atoms_shape(self) -> Tuple[int, ...]:
         return self.__partition.pre_vertex.atoms_shape
 
@@ -182,8 +184,6 @@ def outgoing_edges(self) -> Sequence[DelayedApplicationEdge]:
         return self.__outgoing_edges
 
     @property
+    @overrides(ColouredApplicationVertex.n_colour_bits)
     def n_colour_bits(self) -> int:
-        """
-        The number of bits for event colouring.
-        """
         return self.__n_colour_bits