Prod neuron (#783)

* prod neuron

* trying to get prod neuron to work...

* trying to get prod neuron cpu to work...

* prod neuron process cpu backend working with unit test

* remove init file from prod_neuron

* fixed prod neuron license headers

* renamed prod_neuron to prodneuron. Docstrings.

src/lava/proc/prodneuron/models.py

@@ -0,0 +1,38 @@
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: BSD-3-Clause
+# See: https://spdx.org/licenses/
+
+import numpy as np
+from lava.magma.core.sync.protocols.loihi_protocol import LoihiProtocol
+from lava.magma.core.model.py.ports import PyInPort, PyOutPort
+from lava.magma.core.model.py.type import LavaPyType
+from lava.magma.core.resources import CPU
+from lava.magma.core.decorator import implements, requires, tag
+from lava.magma.core.model.py.model import PyLoihiProcessModel
+
+from lava.proc.prodneuron.process import ProdNeuron
+
+
+@implements(proc=ProdNeuron, protocol=LoihiProtocol)
+@requires(CPU)
+@tag('fixed_pt')
+class PyProdNeuronModelFixed(PyLoihiProcessModel):
+    """Fixed point implementation of ProdNeuron"""
+    a_in1 = LavaPyType(PyInPort.VEC_DENSE, np.int32, precision=24)
+    a_in2 = LavaPyType(PyInPort.VEC_DENSE, np.int32, precision=24)
+    s_out = LavaPyType(PyOutPort.VEC_DENSE, np.int32, precision=24)
+    v: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=24)
+    vth: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=24)
+    exp: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=24)
+
+    def run_spk(self) -> None:
+        a_in_data1 = self.a_in1.recv()
+        a_in_data2 = self.a_in2.recv()
+
+        v = a_in_data1 * a_in_data2
+        v >>= self.exp
+
+        is_spike = np.abs(v) > self.vth
+        sp_out = v * is_spike
+
+        self.s_out.send(sp_out)

src/lava/proc/prodneuron/process.py

@@ -0,0 +1,50 @@
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: BSD-3-Clause
+# See: https://spdx.org/licenses/
+
+import numpy as np
+import typing as ty
+from lava.magma.core.process.process import AbstractProcess
+from lava.magma.core.process.variable import Var
+from lava.magma.core.process.ports.ports import InPort, OutPort
+
+
+class ProdNeuron(AbstractProcess):
+    def __init__(
+            self,
+            shape: ty.Tuple[int, ...],
+            vth: ty.Optional[int] = 1,
+            exp: ty.Optional[int] = 0) -> None:
+        """ProdNeuron
+
+        Multiplies two graded inputs and outputs result as graded spike.
+
+        v[t] = (a_in1 * a_in2) >> exp
+        s_out = v[t] * (v[t] > vth)
+
+        Parameters
+        ----------
+
+        shape : tuple(int)
+            Number and topology of ProdNeuron neurons.
+        vth : int
+            Threshold
+        exp : int
+            Fixed-point base
+        """
+        super().__init__(shape=shape)
+
+        self.a_in1 = InPort(shape=shape)
+        self.a_in2 = InPort(shape=shape)
+
+        self.s_out = OutPort(shape=shape)
+
+        self.vth = Var(shape=(1,), init=vth)
+        self.exp = Var(shape=(1,), init=exp)
+
+        self.v = Var(shape=shape, init=np.zeros(shape, 'int32'))
+
+    @property
+    def shape(self) -> ty.Tuple[int, ...]:
+        """Return shape of the Process."""
+        return self.proc_params['shape']

tests/lava/proc/prodneuron/test_prod_neuron.py

@@ -0,0 +1,76 @@
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: BSD-3-Clause
+# See: https://spdx.org/licenses/
+
+import unittest
+import numpy as np
+
+from lava.proc.prodneuron.process import ProdNeuron
+from lava.proc.dense.process import Dense
+from lava.proc import io
+
+from lava.magma.core.run_conditions import RunSteps
+from lava.magma.core.run_configs import Loihi2SimCfg
+
+
+class TestProdNeuronProc(unittest.TestCase):
+    def test_prod_neuron_out(self):
+        weight_exp = 7
+        num_steps = 10
+
+        weights1 = np.zeros((1, 1))
+        weights1[0, 0] = 1
+
+        weights1 *= 2**weight_exp
+        weights1 = weights1.astype('int')
+
+        weights2 = np.zeros((1, 1))
+        weights2[0, 0] = 0.5
+
+        weights2 *= 2**weight_exp
+        weights2 = weights2.astype('int')
+
+        inp_data1 = np.zeros((weights1.shape[1], num_steps))
+        inp_data2 = np.zeros((weights2.shape[1], num_steps))
+
+        inp_data1[:, 2] = 10
+        inp_data1[:, 6] = 30
+        inp_data2[:, :] = 20
+
+        dense1 = Dense(weights=weights1,
+                       num_message_bits=24,
+                       weight_exp=-weight_exp)
+        dense2 = Dense(weights=weights2,
+                       num_message_bits=24,
+                       weight_exp=-weight_exp)
+
+        vec1 = ProdNeuron(shape=(weights1.shape[0],))
+
+        generator1 = io.source.RingBuffer(data=inp_data1)
+        generator2 = io.source.RingBuffer(data=inp_data2)
+        logger = io.sink.RingBuffer(shape=(weights1.shape[0],),
+                                    buffer=num_steps)
+
+        generator1.s_out.connect(dense1.s_in)
+        dense1.a_out.connect(vec1.a_in1)
+
+        generator2.s_out.connect(dense2.s_in)
+        dense2.a_out.connect(vec1.a_in2)
+
+        vec1.s_out.connect(logger.a_in)
+
+        vec1.run(condition=RunSteps(num_steps=num_steps),
+                 run_cfg=Loihi2SimCfg(select_tag='fixed_pt'))
+        out_data = logger.data.get().astype('int')
+        vec1.stop()
+
+        ch1 = (weights1 @ inp_data1) / 2**weight_exp
+        ch2 = (weights2 @ inp_data2) / 2**weight_exp
+
+        expected_out = ch1 * ch2
+        # then there is one extra timestep from hardware
+        self.assertTrue(np.all(expected_out[:, :-1] == out_data[:, 1:]))
+
+
+if __name__ == '__main__':
+    unittest.main()