OpenMP code generation for synaptic propagation in Brian 2

brian_openmp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brian2 import *\n",
    "set_device('cpp_standalone')  # Use C++ standalone mode\n",
    "prefs.devices.cpp_standalone.openmp_threads = 2  # Switch on OpenMP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING    OpenMP code is not yet well tested, and may be inaccurate. [brian2.devices.cpp_standalone.device.openmp]\n"
     ]
    }
   ],
   "source": [
    "# Neurons that don't do anything, we are only interested in the synapses\n",
    "group1 = NeuronGroup(10, 'x : 1', threshold='x > 1', reset='x = 0')\n",
    "group2 = NeuronGroup(10, 'y : 1')\n",
    "syn = Synapses(group1, group2, 'w : 1',  # synaptic weight variable\n",
    "               on_pre='w += 1'  # for each spike, increase synaptic weight\n",
    "               )\n",
    "syn.connect()  # Connect all-to-all\n",
    "\n",
    "# \"Run\" the simulation for no time at all... we are only interested in the code\n",
    "run(0*ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can print the code that is generated for the synaptic event handling (the `on_pre` above) by using:\n",
    "```Python\n",
    "print(syn.pre.codeobj.code)\n",
    "```\n",
    "This code is very long, though, and contains a lot of boring \"boilerplate\" code. Here's only the interesting bit:\n",
    "```C++\n",
    "#pragma omp parallel\n",
    "{\n",
    "std::vector<int> *_spiking_synapses = synapses_pre.peek();\n",
    "const int _num_spiking_synapses = _spiking_synapses->size();\n",
    "#pragma omp for schedule(static)\n",
    "for(int _spiking_synapse_idx=0;\n",
    "    _spiking_synapse_idx<_num_spiking_synapses;\n",
    "    _spiking_synapse_idx++)\n",
    "{\n",
    "    const size_t _idx = (*_spiking_synapses)[_spiking_synapse_idx];\n",
    "    double w = _ptr_array_synapses_w[_idx];\n",
    "    w += 1;\n",
    "    _ptr_array_synapses_w[_idx] = w;\n",
    "}\n",
    "}\n",
    "```\n",
    "(Note the OpenMP pragmas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2\n",
    "Almost identical to the previous one, the only difference is the `on_pre` statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A few ugly lines that are only necessary if we want to run more than\n",
    "# one standalone simulation in the same script\n",
    "start_scope()\n",
    "device.reinit()\n",
    "device.activate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Neurons that don't do anything, we are only interested in the synapses\n",
    "group1 = NeuronGroup(10, 'x : 1', threshold='x > 1', reset='x = 0')\n",
    "group2 = NeuronGroup(10, 'y : 1')\n",
    "syn = Synapses(group1, group2, 'w : 1',  # synaptic weight variable\n",
    "               on_pre='y += 1'  # for each spike, increase post-synaptic y\n",
    "               )\n",
    "syn.connect()  # Connect all-to-all\n",
    "\n",
    "# \"Run\" the simulation for no time at all... we are only interested in the code\n",
    "run(0*ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the generated code for this example:\n",
    "```C++\n",
    "#pragma omp parallel\n",
    "{\n",
    "std::vector<int> *_spiking_synapses = synapses_1_pre.peek();\n",
    "const int _num_spiking_synapses = _spiking_synapses->size();\n",
    "#pragma omp master\n",
    "{\n",
    "    for(int _spiking_synapse_idx=0;\n",
    "        _spiking_synapse_idx<_num_spiking_synapses;\n",
    "        _spiking_synapse_idx++)\n",
    "    {\n",
    "        const size_t _idx = (*_spiking_synapses)[_spiking_synapse_idx];\n",
    "        const int32_t _postsynaptic_idx = _ptr_array_synapses_1__synaptic_post[_idx];\n",
    "        double y = _ptr_array_neurongroup_4_y[_postsynaptic_idx];\n",
    "        y += 1;\n",
    "        _ptr_array_neurongroup_4_y[_postsynaptic_idx] = y;\n",
    "    }\n",
    "}\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_scope()\n",
    "device.reinit()\n",
    "device.activate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Neurons that don't do anything, we are only interested in the synapses\n",
    "group1 = NeuronGroup(10, 'x : 1', threshold='x > 1', reset='x = 0')\n",
    "group2 = NeuronGroup(10, 'y : 1')\n",
    "syn = Synapses(group1, group2, 'w : 1',  # synaptic weight variable\n",
    "               on_pre='y += 1'  # for each spike, increase post-synaptic y\n",
    "               )\n",
    "syn.connect(j='i')  # Connect one-to-one\n",
    "\n",
    "# \"Run\" the simulation for no time at all... we are only interested in the code\n",
    "run(0*ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generated code is identical to the previous example:\n",
    "```C++\n",
    "#pragma omp parallel\n",
    "{\n",
    "std::vector<int> *_spiking_synapses = synapses_pre.peek();\n",
    "const int _num_spiking_synapses = _spiking_synapses->size();\n",
    "#pragma omp master\n",
    "{\n",
    "    for(int _spiking_synapse_idx=0;\n",
    "        _spiking_synapse_idx<_num_spiking_synapses;\n",
    "        _spiking_synapse_idx++)\n",
    "    {\n",
    "        const size_t _idx = (*_spiking_synapses)[_spiking_synapse_idx];\n",
    "        const int32_t _postsynaptic_idx = _ptr_array_synapses__synaptic_post[_idx];\n",
    "        double y = _ptr_array_neurongroup_1_y[_postsynaptic_idx];\n",
    "        y += 1;\n",
    "        _ptr_array_neurongroup_1_y[_postsynaptic_idx] = y;\n",
    "    }\n",
    "}\n",
    "}\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:brian2_py37] *",
   "language": "python",
   "name": "conda-env-brian2_py37-py"
  },
  "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}