test_modules.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test modules of recognition pipeline\n",
    "Written by Yujun Lin\n",
    "\n",
    "\n",
    "## Prepare\n",
    "import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import lattice ice40\n",
      "import lattice mantle40\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "#os.environ['MANTLE_TARGET'] = 'ice40'\n",
    "from magma import *\n",
    "set_mantle_target(\"ice40\")\n",
    "import mantle\n",
    "import math\n",
    "from mantle.lattice.ice40 import ROMB, SB_LUT4\n",
    "\n",
    "from magma.simulator import PythonSimulator\n",
    "from magma.scope import Scope\n",
    "from magma.bit_vector import BitVector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Global Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of bits for num_cycles: 4\n",
      "number of bits for num_classes: 3\n",
      "number of bits for bit counter output: 5\n",
      "number of bits for bit counter output accumulator: 9\n"
     ]
    }
   ],
   "source": [
    "num_cycles = 16\n",
    "num_classes = 8\n",
    "# operand width\n",
    "N = 16\n",
    "# number of bits for num_cycles\n",
    "n = int(math.ceil(math.log2(num_cycles)))\n",
    "# number of bits for num_classes\n",
    "b = int(math.ceil(math.log2(num_classes)))\n",
    "# number of bits for bit counter output\n",
    "n_bc = int(math.floor(math.log2(N))) + 1\n",
    "# number of bits for bit counter output accumulator\n",
    "n_bc_adder = int(math.floor(math.log2(N*num_cycles))) + 1\n",
    "print('number of bits for num_cycles: %d' % n)\n",
    "print('number of bits for num_classes: %d' % b)\n",
    "print('number of bits for bit counter output: %d' % n_bc)\n",
    "print('number of bits for bit counter output accumulator: %d' % n_bc_adder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Control module\n",
    "\n",
    "generate address for weight and image block\n",
    "\n",
    "`IDX` means the idx-th row of weight matrix\n",
    "\n",
    "`CYCLE` means the cycle-th block of idx-th row of weight matrix\n",
    "\n",
    "`CYCLE` also means the cycle-th block of image vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Controller(Circuit):\n",
    "    name = \"Controller\"\n",
    "    IO = ['CLK', In(Clock), 'IDX', Out(Bits(b)),\n",
    "          'CYCLE', Out(Bits(n))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        adder_cycle = mantle.Add(n, cin=False, cout=False)\n",
    "        reg_cycle = mantle.Register(n, has_reset=True)\n",
    "        adder_idx = mantle.Add(b, cin=False, cout=False)\n",
    "        reg_idx = mantle.Register(b, has_ce=True)\n",
    "        wire(io.CLK, reg_cycle.CLK)\n",
    "        wire(io.CLK, reg_idx.CLK)\n",
    "        wire(reg_cycle.O, adder_cycle.I0)\n",
    "        wire(bits(1, n), adder_cycle.I1)\n",
    "        wire(adder_cycle.O, reg_cycle.I)\n",
    "        comparison_cycle = mantle.EQ(n)\n",
    "        wire(reg_cycle.O, comparison_cycle.I0)\n",
    "        wire(bits(num_cycles-1, n), comparison_cycle.I1)\n",
    "        # if cycle-th is the last, then switch to next idx (accumulate idx) and clear cycle\n",
    "        wire(comparison_cycle.O, reg_cycle.RESET)\n",
    "        wire(comparison_cycle.O, reg_idx.CE)\n",
    "        comparison_idx = mantle.EQ(b)\n",
    "        wire(reg_idx.O, comparison_idx.I0)\n",
    "        wire(bits(num_classes-1, b), comparison_idx.I1)\n",
    "        wire(reg_idx.O, adder_idx.I0)\n",
    "        wire(bits(0, b-1), adder_idx.I1[1:])\n",
    "        nand_gate = mantle.NAnd()\n",
    "        wire(comparison_cycle.O, nand_gate.I0)\n",
    "        wire(comparison_idx.O, nand_gate.I1)\n",
    "        # after all idx rows, we stop accumulating idx\n",
    "        wire(nand_gate.O, adder_idx.I1[0])\n",
    "        wire(adder_idx.O, reg_idx.I)\n",
    "        wire(reg_idx.O, io.IDX)\n",
    "        wire(adder_cycle.O, io.CYCLE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'PythonSimulator' object has no attribute 'step'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-51b380be7d75>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m96\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m     \u001b[0msimulator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m     \u001b[0msimulator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mevaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m     \u001b[0mclk\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msimulator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_value\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mTest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCLK\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'PythonSimulator' object has no attribute 'step'"
     ]
    }
   ],
   "source": [
    "class Test(Circuit):\n",
    "    name = \"Test\"\n",
    "    IO = ['CLK', In(Clock), 'IDX', Out(Bits(b)), 'CYCLE', Out(Bits(n)),\n",
    "          'CONTROL', Out(Bit)]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        # IF\n",
    "        controller = Controller()\n",
    "        reg_1_cycle = mantle.DefineRegister(n)()\n",
    "        reg_1_control = mantle.DFF(init=1)\n",
    "        wire(io.CLK, controller.CLK)\n",
    "        wire(io.CLK, reg_1_cycle.CLK)\n",
    "        wire(io.CLK, reg_1_control.CLK)\n",
    "        reg_1_idx = controller.IDX\n",
    "        wire(controller.CYCLE, reg_1_cycle.I)\n",
    "        wire(1, reg_1_control.I)\n",
    "        \n",
    "        wire(reg_1_idx, io.IDX)\n",
    "        wire(reg_1_cycle.O, io.CYCLE)\n",
    "        wire(reg_1_control.O, io.CONTROL)\n",
    "\n",
    "simulator = PythonSimulator(Test, clock=Test.CLK)\n",
    "waveforms = []\n",
    "\n",
    "for i in range(96):\n",
    "    simulator.step()\n",
    "    simulator.evaluate()\n",
    "    clk = simulator.get_value(Test.CLK)\n",
    "    o = simulator.get_value(Test.IDX)\n",
    "    c = simulator.get_value(Test.CYCLE)\n",
    "    ctl = simulator.get_value(Test.CONTROL)\n",
    "    waveforms.append([clk, ctl] + o + c)\n",
    "names = [\"CLK\", \"CTL\"]\n",
    "for i in range(n):\n",
    "    names.append(\"IDX[{}]\".format(i))\n",
    "for i in range(b):\n",
    "    names.append(\"CYC[{}]\".format(i))\n",
    "from magma.waveform import waveform\n",
    "waveform(waveforms, names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ROM module\n",
    "\n",
    "Test Unit for Rom Reading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestReadRom(Circuit):\n",
    "    name = \"TestReadRom2\"\n",
    "    IO = ['IDX', In(Bits(b)), 'CYCLE', In(Bits(n)), 'CLK', In(Clock),\n",
    "          'WEIGHT', Out(Bits(N)), 'IMAGE', Out(Bits(N))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        weights_list = [1] + [2**16-1]*15 + [3] + [2**16-1]*15 + ([0] + [2**16-1]*15)*((256-32)//16)\n",
    "        weigths_rom = ROMB(weights_list)\n",
    "        lut_list = []\n",
    "        for i in range(N):\n",
    "            lut_list.append(SB_LUT4(LUT_INIT=1))\n",
    "        wire(io.CYCLE, weigths_rom.RADDR[:n])\n",
    "        wire(io.IDX, weigths_rom.RADDR[n:n+b])\n",
    "        if n + b < 8:\n",
    "            wire(bits(0, 8-n-b), weigths_rom.RADDR[n+b:])\n",
    "        wire(1, weigths_rom.RE)\n",
    "        wire(weigths_rom.RDATA, io.WEIGHT)\n",
    "        wire(io.CLK, weigths_rom.RCLK)\n",
    "        for i in range(N):\n",
    "            wire(io.CYCLE, bits([lut_list[i].I0, lut_list[i].I1, lut_list[i].I2, lut_list[i].I3]))\n",
    "            wire(lut_list[i].O, io.IMAGE[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Test(Circuit):\n",
    "    name = \"Test\"\n",
    "    IO = ['CLK', In(Clock), 'WEIGHT', Out(Bits(N)), 'IMAGE', Out(Bits(N)),\n",
    "          'IDX', Out(Bits(b)), 'CYCLE', Out(Bits(n)), 'CONTROL', Out(Bit)]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        # IF - get cycle_id, label_index_id\n",
    "        controller = Controller()\n",
    "        reg_1_cycle = mantle.DefineRegister(n)()\n",
    "        reg_1_control = mantle.DFF(init=1)\n",
    "        wire(io.CLK, controller.CLK)\n",
    "        wire(io.CLK, reg_1_cycle.CLK)\n",
    "        wire(io.CLK, reg_1_control.CLK)\n",
    "        reg_1_idx = controller.IDX\n",
    "        wire(controller.CYCLE, reg_1_cycle.I)\n",
    "        wire(1, reg_1_control.I)\n",
    "        # RR - get weight block, image block of N bits\n",
    "        readROM = TestReadRom()\n",
    "        wire(reg_1_idx, readROM.IDX)\n",
    "        wire(reg_1_cycle.O, readROM.CYCLE)\n",
    "        reg_2 = mantle.DefineRegister(N + b + n)()\n",
    "        reg_2_control = mantle.DFF()\n",
    "        reg_2_weight = readROM.WEIGHT\n",
    "        wire(io.CLK, reg_2.CLK)\n",
    "        wire(io.CLK, readROM.CLK)\n",
    "        wire(io.CLK, reg_2_control.CLK)\n",
    "        wire(readROM.IMAGE, reg_2.I[:N])\n",
    "        wire(reg_1_idx, reg_2.I[N:N + b])\n",
    "        wire(reg_1_cycle.O, reg_2.I[N + b:])\n",
    "        wire(reg_1_control.O, reg_2_control.I)\n",
    "        \n",
    "        wire(reg_2_weight, io.WEIGHT)\n",
    "        wire(reg_2.O[:N], io.IMAGE)\n",
    "        wire(reg_2.O[N:N+b], io.IDX)\n",
    "        wire(reg_2.O[N+b:], io.CYCLE)\n",
    "        wire(reg_2_control.O, io.CONTROL)\n",
    "\n",
    "simulator = PythonSimulator(Test, clock=Test.CLK)\n",
    "waveforms = []\n",
    "\n",
    "for i in range(96):\n",
    "    simulator.step()\n",
    "    simulator.evaluate()\n",
    "    clk = simulator.get_value(Test.CLK)\n",
    "    w = simulator.get_value(Test.WEIGHT)\n",
    "    i = simulator.get_value(Test.IMAGE)\n",
    "    d = simulator.get_value(Test.IDX)\n",
    "    c = simulator.get_value(Test.CYCLE)\n",
    "    ctl = simulator.get_value(Test.CONTROL)\n",
    "    waveforms.append([clk, ctl] + w + i + d + c)\n",
    "names = [\"CLK\", \"CTL\"]\n",
    "for i in range(N):\n",
    "    names.append(\"WGT[{}]\".format(i))\n",
    "for i in range(N):\n",
    "    names.append(\"IMG[{}]\".format(i))\n",
    "for i in range(n):\n",
    "    names.append(\"IDX[{}]\".format(i))\n",
    "for i in range(b):\n",
    "    names.append(\"CYC[{}]\".format(i))\n",
    "from magma.waveform import waveform\n",
    "waveform(waveforms, names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pop Count Unit\n",
    "4/8/16 bit pop count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4-bit pop count\n",
    "class BitCounter4(Circuit):\n",
    "    name = \"BitCounter4\"\n",
    "    IO = ['I', In(Bits(4)), 'O', Out(Bits(3))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        lut_list = []\n",
    "        lut_list.append(SB_LUT4(LUT_INIT=int('0110100110010110', 2)))\n",
    "        lut_list.append(SB_LUT4(LUT_INIT=int('0111111011101000', 2)))\n",
    "        lut_list.append(SB_LUT4(LUT_INIT=int('1000000000000000', 2)))\n",
    "        for i in range(3):\n",
    "            wire(io.I, bits([lut_list[i].I0, lut_list[i].I1, lut_list[i].I2, lut_list[i].I3]))\n",
    "            wire(lut_list[i].O, io.O[i])\n",
    "\n",
    "\n",
    "# 8-bit pop count\n",
    "class BitCounter8(Circuit):\n",
    "    name = \"BitCounter8\"\n",
    "    IO = ['I', In(Bits(8)), 'O', Out(Bits(4))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        counter_1 = BitCounter4()\n",
    "        counter_2 = BitCounter4()\n",
    "        wire(io.I[:4], counter_1.I)\n",
    "        wire(io.I[4:], counter_2.I)\n",
    "        adders = [mantle.HalfAdder()] + [mantle.FullAdder() for _ in range(2)]\n",
    "        for i in range(3):\n",
    "            wire(counter_1.O[i], adders[i].I0)\n",
    "            wire(counter_2.O[i], adders[i].I1)\n",
    "            if i > 0:\n",
    "                wire(adders[i-1].COUT, adders[i].CIN)\n",
    "            wire(adders[i].O, io.O[i])\n",
    "        wire(adders[-1].COUT, io.O[-1])\n",
    "\n",
    "\n",
    "# 16-bit pop count\n",
    "class BitCounter16(Circuit):\n",
    "    name = 'BitCounter16'\n",
    "    IO = ['I', In(Bits(16)), 'O', Out(Bits(5))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        counter_1 = BitCounter8()\n",
    "        counter_2 = BitCounter8()\n",
    "        wire(io.I[:8], counter_1.I)\n",
    "        wire(io.I[8:], counter_2.I)\n",
    "        adders = [mantle.HalfAdder()] + [mantle.FullAdder() for _ in range(3)]\n",
    "        for i in range(4):\n",
    "            wire(counter_1.O[i], adders[i].I0)\n",
    "            wire(counter_2.O[i], adders[i].I1)\n",
    "            if i > 0:\n",
    "                wire(adders[i-1].COUT, adders[i].CIN)\n",
    "            wire(adders[i].O, io.O[i])\n",
    "        wire(adders[-1].COUT, io.O[-1])\n",
    "\n",
    "\n",
    "# pop count\n",
    "def DefineBitCounter(n):\n",
    "    if n <= 4:\n",
    "        return BitCounter4\n",
    "    elif n <= 8:\n",
    "        return BitCounter8\n",
    "    elif n <= 16:\n",
    "        return BitCounter16\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Test(Circuit):\n",
    "    name = \"Test\"\n",
    "    IO = ['CLK', In(Clock), 'COUNT', Out(Bits(n_bc_adder)), 'CONTROL', Out(Bit)]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        # IF - get cycle_id, label_index_id\n",
    "        controller = Controller()\n",
    "        reg_1_cycle = mantle.Register(n)\n",
    "        reg_1_control = mantle.DFF(init=1)\n",
    "        wire(io.CLK, controller.CLK)\n",
    "        wire(io.CLK, reg_1_cycle.CLK)\n",
    "        wire(io.CLK, reg_1_control.CLK)\n",
    "        reg_1_idx = controller.IDX\n",
    "        wire(controller.CYCLE, reg_1_cycle.I)\n",
    "        wire(1, reg_1_control.I)\n",
    "        # RR - get weight block, image block of N bits\n",
    "        readROM = TestReadRom()\n",
    "        wire(reg_1_idx, readROM.IDX)\n",
    "        wire(reg_1_cycle.O, readROM.CYCLE)\n",
    "        reg_2 = mantle.Register(N + b + n)\n",
    "        reg_2_control = mantle.DFF()\n",
    "        reg_2_weight = readROM.WEIGHT\n",
    "        wire(io.CLK, reg_2.CLK)\n",
    "        wire(io.CLK, readROM.CLK)\n",
    "        wire(io.CLK, reg_2_control.CLK)\n",
    "        wire(readROM.IMAGE, reg_2.I[:N])\n",
    "        wire(reg_1_idx, reg_2.I[N:N + b])\n",
    "        wire(reg_1_cycle.O, reg_2.I[N + b:])\n",
    "        wire(reg_1_control.O, reg_2_control.I)\n",
    "        # EX - NXOr for multiplication, pop count and accumulate the result for activation\n",
    "        multiplier = mantle.NXOr(height=2, width=N)\n",
    "        bit_counter = DefineBitCounter(N)()\n",
    "        adder = mantle.Add(n_bc_adder, cin=False, cout=False)\n",
    "        mux_for_adder_0 = mantle.Mux(height=2, width=n_bc_adder)\n",
    "        mux_for_adder_1 = mantle.Mux(height=2, width=n_bc_adder)\n",
    "        reg_3_1 = mantle.Register(n_bc_adder)\n",
    "        reg_3_2 = mantle.Register(b + n)\n",
    "        wire(io.CLK, reg_3_1.CLK)\n",
    "        wire(io.CLK, reg_3_2.CLK)\n",
    "        wire(reg_2_weight, multiplier.I0)\n",
    "        wire(reg_2.O[:N], multiplier.I1)\n",
    "        wire(multiplier.O, bit_counter.I)\n",
    "        wire(bits(0, n_bc_adder), mux_for_adder_0.I0)\n",
    "        wire(bit_counter.O, mux_for_adder_0.I1[:n_bc])\n",
    "        if n_bc_adder > n_bc:\n",
    "            wire(bits(0, n_bc_adder - n_bc), mux_for_adder_0.I1[n_bc:])\n",
    "        # only when data read is ready (i.e. control signal is high), accumulate the pop count result\n",
    "        wire(reg_2_control.O, mux_for_adder_0.S)\n",
    "        wire(reg_3_1.O, mux_for_adder_1.I0)\n",
    "        wire(bits(0, n_bc_adder), mux_for_adder_1.I1)\n",
    "        if n == 4:\n",
    "            comparison_3 = SB_LUT4(LUT_INIT=int('0'*15+'1', 2))\n",
    "            wire(reg_2.O[N+b:], bits([comparison_3.I0, comparison_3.I1, comparison_3.I2, comparison_3.I3]))\n",
    "        else:\n",
    "            comparison_3 = mantle.EQ(n)\n",
    "            wire(reg_2.O[N+b:], comparison_3.I0)\n",
    "            wire(bits(0, n), comparison_3.I1)\n",
    "        wire(comparison_3.O, mux_for_adder_1.S)\n",
    "        wire(mux_for_adder_0.O, adder.I0)\n",
    "        wire(mux_for_adder_1.O, adder.I1)\n",
    "        wire(adder.O, reg_3_1.I)\n",
    "        wire(reg_2.O[N:], reg_3_2.I)\n",
    "\n",
    "        wire(reg_3_1.O, io.COUNT)\n",
    "        wire(reg_2_control.O, io.CONTROL)\n",
    "\n",
    "simulator = PythonSimulator(Test, clock=Test.CLK)\n",
    "waveforms = []\n",
    "\n",
    "for i in range(128):\n",
    "    simulator.step()\n",
    "    simulator.evaluate()\n",
    "    clk = simulator.get_value(Test.CLK)\n",
    "    o = simulator.get_value(Test.COUNT)\n",
    "    ctl = simulator.get_value(Test.CONTROL)\n",
    "    waveforms.append([clk, ctl] + o)\n",
    "names = [\"CLK\", \"CTL\"]\n",
    "for i in range(n_bc_adder):\n",
    "    names.append(\"COUNT[{}]\".format(i))\n",
    "from magma.waveform import waveform\n",
    "waveform(waveforms, names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classifier Module\n",
    "\n",
    "using compare operation to decide the final prediction label of image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Classifier(Circuit):\n",
    "    name = \"Classifier\"\n",
    "    IO = ['I', In(Bits(n_bc_adder)), 'IDX', In(Bits(b)), 'CLK', In(Clock),\n",
    "          'O', Out(Bits(b)), 'M', Out(Bits(n_bc_adder))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        comparison = mantle.UGT(n_bc_adder)\n",
    "        reg_count = mantle.Register(n_bc_adder, has_ce=True)\n",
    "        reg_idx = mantle.Register(b, has_ce=True)\n",
    "        wire(io.I, comparison.I0)\n",
    "        wire(reg_count.O, comparison.I1)\n",
    "        wire(comparison.O, reg_count.CE)\n",
    "        wire(comparison.O, reg_idx.CE)\n",
    "        wire(io.CLK, reg_count.CLK)\n",
    "        wire(io.CLK, reg_idx.CLK)\n",
    "        wire(io.I, reg_count.I)\n",
    "        wire(io.IDX, reg_idx.I)\n",
    "        wire(reg_idx.O, io.O)\n",
    "        wire(reg_count.O, io.M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Test(Circuit):\n",
    "    name = \"Test\"\n",
    "    IO = ['CLK', In(Clock), 'MAX', Out(Bits(n_bc_adder)),\n",
    "          'IDX', Out(Bits(b)), 'COUNT', Out(Bits(n_bc_adder))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        # IF - get cycle_id, label_index_id\n",
    "        controller = Controller()\n",
    "        reg_1_cycle = mantle.Register(n)\n",
    "        reg_1_control = mantle.DFF(init=1)\n",
    "        wire(io.CLK, controller.CLK)\n",
    "        wire(io.CLK, reg_1_cycle.CLK)\n",
    "        wire(io.CLK, reg_1_control.CLK)\n",
    "        reg_1_idx = controller.IDX\n",
    "        wire(controller.CYCLE, reg_1_cycle.I)\n",
    "        wire(1, reg_1_control.I)\n",
    "        # RR - get weight block, image block of N bits\n",
    "        readROM = TestReadRom()\n",
    "        wire(reg_1_idx, readROM.IDX)\n",
    "        wire(reg_1_cycle.O, readROM.CYCLE)\n",
    "        reg_2 = mantle.Register(N + b + n)\n",
    "        reg_2_control = mantle.DFF()\n",
    "        reg_2_weight = readROM.WEIGHT\n",
    "        wire(io.CLK, reg_2.CLK)\n",
    "        wire(io.CLK, readROM.CLK)\n",
    "        wire(io.CLK, reg_2_control.CLK)\n",
    "        wire(readROM.IMAGE, reg_2.I[:N])\n",
    "        wire(reg_1_idx, reg_2.I[N:N + b])\n",
    "        wire(reg_1_cycle.O, reg_2.I[N + b:])\n",
    "        wire(reg_1_control.O, reg_2_control.I)\n",
    "        # EX - NXOr for multiplication, pop count and accumulate the result for activation\n",
    "        multiplier = mantle.NXOr(height=2, width=N)\n",
    "        bit_counter = DefineBitCounter(N)()\n",
    "        adder = mantle.Add(n_bc_adder, cin=False, cout=False)\n",
    "        mux_for_adder_0 = mantle.Mux(height=2, width=n_bc_adder)\n",
    "        mux_for_adder_1 = mantle.Mux(height=2, width=n_bc_adder)\n",
    "        reg_3_1 = mantle.Register(n_bc_adder)\n",
    "        reg_3_2 = mantle.Register(b + n)\n",
    "        wire(io.CLK, reg_3_1.CLK)\n",
    "        wire(io.CLK, reg_3_2.CLK)\n",
    "        wire(reg_2_weight, multiplier.I0)\n",
    "        wire(reg_2.O[:N], multiplier.I1)\n",
    "        wire(multiplier.O, bit_counter.I)\n",
    "        wire(bits(0, n_bc_adder), mux_for_adder_0.I0)\n",
    "        wire(bit_counter.O, mux_for_adder_0.I1[:n_bc])\n",
    "        if n_bc_adder > n_bc:\n",
    "            wire(bits(0, n_bc_adder - n_bc), mux_for_adder_0.I1[n_bc:])\n",
    "        # only when data read is ready (i.e. control signal is high), accumulate the pop count result\n",
    "        wire(reg_2_control.O, mux_for_adder_0.S)\n",
    "        wire(reg_3_1.O, mux_for_adder_1.I0)\n",
    "        wire(bits(0, n_bc_adder), mux_for_adder_1.I1)\n",
    "        if n == 4:\n",
    "            comparison_3 = SB_LUT4(LUT_INIT=int('0'*15+'1', 2))\n",
    "            wire(reg_2.O[N+b:], bits([comparison_3.I0, comparison_3.I1, comparison_3.I2, comparison_3.I3]))\n",
    "        else:\n",
    "            comparison_3 = mantle.EQ(n)\n",
    "            wire(reg_2.O[N+b:], comparison_3.I0)\n",
    "            wire(bits(0, n), comparison_3.I1)\n",
    "        wire(comparison_3.O, mux_for_adder_1.S)\n",
    "        wire(mux_for_adder_0.O, adder.I0)\n",
    "        wire(mux_for_adder_1.O, adder.I1)\n",
    "        wire(adder.O, reg_3_1.I)\n",
    "        wire(reg_2.O[N:], reg_3_2.I)\n",
    "        # CF - classify the image\n",
    "        classifier = Classifier()\n",
    "        reg_4 = mantle.Register(n + b)\n",
    "        reg_4_idx = classifier.O\n",
    "        wire(io.CLK, classifier.CLK)\n",
    "        wire(io.CLK, reg_4.CLK)\n",
    "        wire(reg_3_1.O, classifier.I)\n",
    "        wire(reg_3_2.O[:b], classifier.IDX)\n",
    "        wire(reg_3_2.O, reg_4.I)\n",
    "        \n",
    "        wire(reg_3_1.O, io.COUNT)\n",
    "        wire(classifier.O, io.IDX)\n",
    "        wire(classifier.M, io.MAX)\n",
    "\n",
    "simulator = PythonSimulator(Test, clock=Test.CLK)\n",
    "waveforms = []\n",
    "\n",
    "for i in range(128):\n",
    "    simulator.step()\n",
    "    simulator.evaluate()\n",
    "    clk = simulator.get_value(Test.CLK)\n",
    "    o = simulator.get_value(Test.IDX)\n",
    "    m = simulator.get_value(Test.MAX)\n",
    "    c = simulator.get_value(Test.COUNT)\n",
    "    waveforms.append([clk] + o + m + c)\n",
    "names = [\"CLK\"]\n",
    "for i in range(b):\n",
    "    names.append(\"IDX[{}]\".format(i))\n",
    "for i in range(n_bc_adder):\n",
    "    names.append(\"MAX[{}]\".format(i))\n",
    "for i in range(n_bc_adder):\n",
    "    names.append(\"CNT[{}]\".format(i))\n",
    "from magma.waveform import waveform\n",
    "waveform(waveforms, names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Classifier(Circuit):\n",
    "    name = \"Classifier\"\n",
    "    IO = ['I', In(Bits(n_bc_adder)), 'IDX', In(Bits(b)), 'CLK', In(Clock), 'O', Out(Bits(b))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        comparison = mantle.UGT(n_bc_adder)\n",
    "        reg_count = mantle.Register(n_bc_adder, has_ce=True)\n",
    "        reg_idx = mantle.Register(b, has_ce=True)\n",
    "        wire(io.I, comparison.I0)\n",
    "        wire(reg_count.O, comparison.I1)\n",
    "        wire(comparison.O, reg_count.CE)\n",
    "        wire(comparison.O, reg_idx.CE)\n",
    "        wire(io.CLK, reg_count.CLK)\n",
    "        wire(io.CLK, reg_idx.CLK)\n",
    "        wire(io.I, reg_count.I)\n",
    "        wire(io.IDX, reg_idx.I)\n",
    "        wire(reg_idx.O, io.O)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipeline Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Test(Circuit):\n",
    "    name = \"Test\"\n",
    "    IO = ['CLK', In(Clock), 'O', Out(Bits(b)), 'IDX', Out(Bits(b))]\n",
    "    @classmethod\n",
    "    def definition(io):\n",
    "        # IF - get cycle_id, label_index_id\n",
    "        controller = Controller()\n",
    "        reg_1_cycle = mantle.Register(n)\n",
    "        reg_1_control = mantle.DFF(init=1)\n",
    "        wire(io.CLK, controller.CLK)\n",
    "        wire(io.CLK, reg_1_cycle.CLK)\n",
    "        wire(io.CLK, reg_1_control.CLK)\n",
    "        reg_1_idx = controller.IDX\n",
    "        wire(controller.CYCLE, reg_1_cycle.I)\n",
    "        wire(1, reg_1_control.I)\n",
    "        # RR - get weight block, image block of N bits\n",
    "        readROM = TestReadRom()\n",
    "        wire(reg_1_idx, readROM.IDX)\n",
    "        wire(reg_1_cycle.O, readROM.CYCLE)\n",
    "        reg_2 = mantle.Register(N + b + n)\n",
    "        reg_2_control = mantle.DFF()\n",
    "        reg_2_weight = readROM.WEIGHT\n",
    "        wire(io.CLK, reg_2.CLK)\n",
    "        wire(io.CLK, readROM.CLK)\n",
    "        wire(io.CLK, reg_2_control.CLK)\n",
    "        wire(readROM.IMAGE, reg_2.I[:N])\n",
    "        wire(reg_1_idx, reg_2.I[N:N + b])\n",
    "        wire(reg_1_cycle.O, reg_2.I[N + b:])\n",
    "        wire(reg_1_control.O, reg_2_control.I)\n",
    "        # EX - NXOr for multiplication, pop count and accumulate the result for activation\n",
    "        multiplier = mantle.NXOr(height=2, width=N)\n",
    "        bit_counter = DefineBitCounter(N)()\n",
    "        adder = mantle.Add(n_bc_adder, cin=False, cout=False)\n",
    "        mux_for_adder_0 = mantle.Mux(height=2, width=n_bc_adder)\n",
    "        mux_for_adder_1 = mantle.Mux(height=2, width=n_bc_adder)\n",
    "        reg_3_1 = mantle.Register(n_bc_adder)\n",
    "        reg_3_2 = mantle.Register(b + n)\n",
    "        wire(io.CLK, reg_3_1.CLK)\n",
    "        wire(io.CLK, reg_3_2.CLK)\n",
    "        wire(reg_2_weight, multiplier.I0)\n",
    "        wire(reg_2.O[:N], multiplier.I1)\n",
    "        wire(multiplier.O, bit_counter.I)\n",
    "        wire(bits(0, n_bc_adder), mux_for_adder_0.I0)\n",
    "        wire(bit_counter.O, mux_for_adder_0.I1[:n_bc])\n",
    "        if n_bc_adder > n_bc:\n",
    "            wire(bits(0, n_bc_adder - n_bc), mux_for_adder_0.I1[n_bc:])\n",
    "        # only when data read is ready (i.e. control signal is high), accumulate the pop count result\n",
    "        wire(reg_2_control.O, mux_for_adder_0.S)\n",
    "        wire(reg_3_1.O, mux_for_adder_1.I0)\n",
    "        wire(bits(0, n_bc_adder), mux_for_adder_1.I1)\n",
    "        if n == 4:\n",
    "            comparison_3 = SB_LUT4(LUT_INIT=int('0'*15+'1', 2))\n",
    "            wire(reg_2.O[N+b:], bits([comparison_3.I0, comparison_3.I1, comparison_3.I2, comparison_3.I3]))\n",
    "        else:\n",
    "            comparison_3 = mantle.EQ(n)\n",
    "            wire(reg_2.O[N+b:], comparison_3.I0)\n",
    "            wire(bits(0, n), comparison_3.I1)\n",
    "        wire(comparison_3.O, mux_for_adder_1.S)\n",
    "        wire(mux_for_adder_0.O, adder.I0)\n",
    "        wire(mux_for_adder_1.O, adder.I1)\n",
    "        wire(adder.O, reg_3_1.I)\n",
    "        wire(reg_2.O[N:], reg_3_2.I)\n",
    "        # CF - classify the image\n",
    "        classifier = Classifier()\n",
    "        reg_4 = mantle.Register(n + b)\n",
    "        reg_4_idx = classifier.O\n",
    "        wire(io.CLK, classifier.CLK)\n",
    "        wire(io.CLK, reg_4.CLK)\n",
    "        wire(reg_3_1.O, classifier.I)\n",
    "        wire(reg_3_2.O[:b], classifier.IDX)\n",
    "        wire(reg_3_2.O, reg_4.I)\n",
    "        # WB - wait to show the result until the end\n",
    "        reg_5 = mantle.Register(b, has_ce=True)\n",
    "        comparison_5_1 = mantle.EQ(b)\n",
    "        comparison_5_2 = mantle.EQ(n)\n",
    "        and_gate = mantle.And()\n",
    "        wire(io.CLK, reg_5.CLK)\n",
    "        wire(reg_4_idx, reg_5.I)\n",
    "        wire(reg_4.O[:b], comparison_5_1.I0)\n",
    "        wire(bits(num_classes - 1, b), comparison_5_1.I1)\n",
    "        wire(reg_4.O[b:], comparison_5_2.I0)\n",
    "        wire(bits(num_cycles - 1, n), comparison_5_2.I1)\n",
    "        wire(comparison_5_1.O, and_gate.I0)\n",
    "        wire(comparison_5_2.O, and_gate.I1)\n",
    "        wire(and_gate.O, reg_5.CE)\n",
    "        wire(reg_5.O, io.O)\n",
    "        \n",
    "        wire(classifier.O, io.IDX)\n",
    "\n",
    "\n",
    "simulator = PythonSimulator(Test, clock=Test.CLK)\n",
    "waveforms = []\n",
    "\n",
    "for i in range(300):\n",
    "    simulator.step()\n",
    "    simulator.evaluate()\n",
    "    clk = simulator.get_value(Test.CLK)\n",
    "    o = simulator.get_value(Test.O)\n",
    "    i = simulator.get_value(Test.IDX)\n",
    "    waveforms.append([clk] + o + i)\n",
    "names = [\"CLK\"]\n",
    "for i in range(b):\n",
    "    names.append(\"O[{}]\".format(i))\n",
    "for i in range(b):\n",
    "    names.append(\"I[{}]\".format(i))\n",
    "from magma.waveform import waveform\n",
    "waveform(waveforms, names)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}