# Copyright (c) 2022 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 typing import Iterator, MutableMapping, Optional, Tuple, cast
from typing_extensions import Never
import numpy
from numpy import uint32
from numpy.typing import NDArray
from spinn_utilities.helpful_functions import is_singleton
from spinn_utilities.ranged import RangeDictionary, RangedList
from pacman.model.graphs.common import Slice
from pacman.model.placements import Placement
from spinn_front_end_common.interface.ds import DataSpecificationBase
from spinn_front_end_common.utilities.constants import BYTES_PER_WORD
from spinn_front_end_common.utilities.helpful_functions import (
locate_memory_region_for_placement)
from spynnaker.pyNN.utilities.struct import Struct, StructRepeat
from spynnaker.pyNN.data import SpynnakerDataView
from spynnaker.pyNN.models.neuron.population_vertex import (
PopulationVertex)
from spynnaker.pyNN.models.neuron.population_machine_neurons import (
NeuronRegions)
def _all_one_val_gen(rd: RangeDictionary[float]) -> bool:
"""
Determine if all the values of a dictionary are the same, assuming we
already know that they are generatable
.. note::
A random distribution is considered the same if the same distribution
is used for all neurons.
"""
for key in rd.keys():
if is_singleton(rd[key]):
return True
if not rd[key].range_based():
return False
for i, (_start, _stop, _val) in enumerate(rd[key].iter_ranges()):
if i > 0:
return False
return True
[docs]
class NeuronData(object):
"""
Holds and creates the data for a group of neurons.
"""
__slots__ = (
# The application vertex
"__app_vertex",
# The data to be written for all vertices, if applicable
"__neuron_data",
# The number of structs in neuron_data
"__neuron_data_n_structs",
# The neuron recording data for all vertices, if applicable
"__neuron_recording_data",
# Whether an attempt has been made to generate neuron data
"__generation_done",
# Whether to generate things on the machine
"__gen_on_machine")
def __init__(self, app_vertex: PopulationVertex):
"""
:param app_vertex: Vertex to get data from.
"""
self.__app_vertex = app_vertex
self.__neuron_data: Optional[NDArray[uint32]] = None
self.__neuron_recording_data: Optional[NDArray[uint32]] = None
self.__generation_done = False
self.__gen_on_machine = False
self.__neuron_data_n_structs = 0
@property
def gen_on_machine(self) -> bool:
"""
Whether the neuron data can be generated on the machine or not.
"""
return self.__gen_on_machine
[docs]
def generate_data(self) -> None:
"""
Do the data generation internally.
"""
if self.__generation_done:
return
self.__generation_done = True
if not self.__app_vertex.can_generate_on_machine():
self.__gen_on_machine = False
return
# Check that all parameters and state variables have a single range
params = self.__app_vertex.parameters
state_vars = self.__app_vertex.state_variables
if not _all_one_val_gen(params) or not _all_one_val_gen(state_vars):
# Note at this point, we can still generate ranges on machine,
# just that it has to be different per core
self.__gen_on_machine = True
self.__neuron_data = None
return
# Go through all the structs and make all the data
structs = self.__app_vertex.neuron_impl.structs
values = _MergedDict(params, state_vars)
all_data = [struct.get_generator_data(values) for struct in structs]
self.__neuron_data = numpy.concatenate(all_data)
self.__neuron_data_n_structs = len(structs)
# Get the neuron recording data
neuron_recorder = self.__app_vertex.neuron_recorder
if neuron_recorder.is_global_generatable:
self.__neuron_recording_data = neuron_recorder.get_generator_data()
# If we get here, we know everything is generated on machine
self.__gen_on_machine = True
[docs]
def write_data(self, spec: DataSpecificationBase, vertex_slice: Slice,
neuron_regions: NeuronRegions,
gen_on_machine: bool = True) -> None:
"""
Write the generated data.
:param spec: The data specification to write to
:param vertex_slice: The vertex slice to generate for
:param neuron_regions: The regions to write to
:param gen_on_machine: Whether to allow generation on machine
"""
if gen_on_machine:
self.generate_data()
spec.reserve_memory_region(
region=neuron_regions.neuron_params,
size=self.__app_vertex.get_sdram_usage_for_neuron_params(
vertex_slice.n_atoms),
label="neuron_params")
neuron_recorder = self.__app_vertex.neuron_recorder
spec.reserve_memory_region(
region=neuron_regions.neuron_recording,
size=neuron_recorder.get_metadata_sdram_usage_in_bytes(
vertex_slice.n_atoms),
label="neuron recording")
if self.gen_on_machine and gen_on_machine:
if self.__neuron_data is not None:
data = self.__neuron_data
n_structs = self.__neuron_data_n_structs
else:
n_structs, data = self.__get_neuron_builder_data(vertex_slice)
header = self.__get_neuron_builder_header(
vertex_slice, n_structs, neuron_regions)
if self.__neuron_recording_data is not None:
rec_data = self.__neuron_recording_data
else:
rec_data = neuron_recorder.get_generator_data(vertex_slice)
n_words = len(data) + len(header) + len(rec_data)
spec.reserve_memory_region(
region=neuron_regions.neuron_builder,
size=n_words * BYTES_PER_WORD,
label="neuron_builder")
spec.switch_write_focus(neuron_regions.neuron_builder)
spec.write_array(header)
spec.write_array(data)
spec.write_array(rec_data)
else:
spec.switch_write_focus(neuron_regions.neuron_params)
neuron_data = self.__get_neuron_param_data(vertex_slice)
spec.write_array(neuron_data)
neuron_recorder.write_neuron_recording_region(
spec, neuron_regions.neuron_recording, vertex_slice)
spec.reserve_memory_region(
region=neuron_regions.initial_values,
size=self.__app_vertex.get_sdram_usage_for_neuron_params(
vertex_slice.n_atoms),
label="initial_values")
def __get_neuron_param_data(self, vertex_slice: Slice) -> NDArray:
"""
Get neuron parameter data for a slice.
:param vertex_slice: The slice to get the data for
"""
structs = self.__app_vertex.neuron_impl.structs
values = _MergedDict(self.__app_vertex.parameters,
self.__app_vertex.state_variables)
return numpy.concatenate([
self.__get_struct_data(struct, values, vertex_slice)
for struct in structs])
def __get_struct_data(self, struct: Struct, values: '_MergedDict',
vertex_slice: Slice) -> NDArray[uint32]:
"""
Get the data for a structure.
:param struct: The structure to get the data for
:param values: The values to fill in the structure with
:param vertex_slice: The slice to get the values for
"""
if struct.repeat_type == StructRepeat.GLOBAL:
return struct.get_data(values)
return struct.get_data(values, vertex_slice)
def __get_neuron_builder_data(
self, vertex_slice: Slice) -> Tuple[int, NDArray[uint32]]:
"""
Get the data to build neuron parameters with.
:param vertex_slice: The slice to get the parameters for
:return: The number of structures and the data
"""
structs = self.__app_vertex.neuron_impl.structs
values = _MergedDict(self.__app_vertex.parameters,
self.__app_vertex.state_variables)
return len(structs), numpy.concatenate([
self.__get_builder_data(struct, values, vertex_slice)
for struct in structs])
def __get_builder_data(
self, struct: Struct, values: '_MergedDict',
vertex_slice: Slice) -> NDArray[uint32]:
"""
Get the builder data for a structure.
:param struct: The structure to get the data for
:param values: The values to fill in the structure with
:param vertex_slice: The slice to get the values for
"""
if struct.repeat_type == StructRepeat.GLOBAL:
return struct.get_generator_data(values)
return struct.get_generator_data(values, vertex_slice)
def __get_neuron_builder_header(
self, vertex_slice: Slice, n_structs: int,
neuron_regions: NeuronRegions) -> NDArray[uint32]:
"""
Get the header of the neuron builder region.
:param vertex_slice: The slice to put in the header
:param n_structs: The number of structures to generate
:param neuron_regions: The regions to point to
"""
return numpy.array([
neuron_regions.neuron_params,
neuron_regions.neuron_recording,
*self.__app_vertex.pop_seed,
*self.__app_vertex.core_seed(vertex_slice),
n_structs, vertex_slice.n_atoms], dtype=uint32)
[docs]
def read_data(self, placement: Placement,
neuron_regions: NeuronRegions) -> None:
"""
Read the current state of the data from the machine into the
application vertex.
:param placement:
The placement of the vertex to read
:param neuron_regions: The regions to read from
"""
merged_dict = _MergedDict(self.__app_vertex.parameters,
self.__app_vertex.state_variables)
self.__do_read_data(
placement, neuron_regions.neuron_params, merged_dict)
[docs]
def read_initial_data(self, placement: Placement,
neuron_regions: NeuronRegions) -> None:
"""
Read the initial state of the data from the machine into the
application vertex.
:param placement:
The placement of the vertex to read
:param neuron_regions: The regions to read from
"""
merged_dict = _MergedDict(self.__app_vertex.parameters,
self.__app_vertex.initial_state_variables)
self.__do_read_data(
placement, neuron_regions.initial_values, merged_dict)
def __do_read_data(self, placement: Placement, region: int,
results: '_MergedDict') -> None:
"""
Perform the reading of data.
:param placement: Where the vertex is on the machine
:param region: The region to read from
:param results: Where to write the results to
"""
address = locate_memory_region_for_placement(placement, region)
vertex_slice = placement.vertex.vertex_slice
data_size = self.__app_vertex.get_sdram_usage_for_neuron_params(
vertex_slice.n_atoms)
block = SpynnakerDataView.read_memory(
placement.x, placement.y, address, data_size)
offset = 0
for struct in self.__app_vertex.neuron_impl.structs:
if struct.repeat_type == StructRepeat.GLOBAL:
struct.read_data(block, cast(RangeDictionary, results), offset)
offset += struct.get_size_in_whole_words() * BYTES_PER_WORD
else:
struct.read_data(block, cast(RangeDictionary, results), offset,
vertex_slice)
offset += (
struct.get_size_in_whole_words(vertex_slice.n_atoms) *
BYTES_PER_WORD)
[docs]
def reset_generation(self) -> None:
"""
Reset generation so it is done again.
"""
self.__neuron_data = None
self.__neuron_recording_data = None
self.__generation_done = False
self.__gen_on_machine = False
self.__neuron_data_n_structs = 0
class _MergedDict(MutableMapping[str, RangedList[float]]):
__slots__ = (
"__params",
"__state_vars")
def __init__(self, params: RangeDictionary[float],
state_vars: RangeDictionary[float]):
"""
:param params: The parameter values
:param state_vars: The State values
"""
self.__params = params
self.__state_vars = state_vars
def __contains__(self, key: object) -> bool:
if not isinstance(key, str):
return False
return key in self.__params or key in self.__state_vars
def __getitem__(self, key: str) -> RangedList[float]:
if key in self.__params:
return self.__params[key]
return self.__state_vars[key]
def __setitem__(self, key: str, value: RangedList[float]) -> None:
if key in self.__params:
self.__params[key] = value
elif key in self.__state_vars:
self.__state_vars[key] = value
else:
raise KeyError(f"No such key {key}")
def __delitem__(self, __v: str) -> Never:
raise NotImplementedError("items may not be deleted")
def __iter__(self) -> Iterator[str]:
yield from self.__params.keys()
yield from self.__state_vars.keys()
def __len__(self) -> int:
return len(self.__params) + len(self.__state_vars)