{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83b016f7-695a-43b7-8baf-4bd0d435a0c8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import subprocess\n",
    "import itertools\n",
    "from collections import defaultdict\n",
    "import importlib.machinery\n",
    "import os\n",
    "import time\n",
    "import cython\n",
    "from matplotlib.patches import Rectangle\n",
    "    \n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['hatch.linewidth'] = 5.0 \n",
    "cycle = [x['color'] for x in mpl.rcParams['axes.prop_cycle']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "705868b4-c013-4fc2-a4c3-bdaca6c90d05",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "experiment_labels = []\n",
    "sizes = defaultdict(lambda: [])\n",
    "times = defaultdict(lambda: [])\n",
    "\n",
    "from sysconfig import get_paths as gp\n",
    "suffix = importlib.machinery.EXTENSION_SUFFIXES[0]\n",
    "\n",
    "# Path to pybind11 git repository\n",
    "pybind11_path = '/home/wjakob/pybind11/include'\n",
    "\n",
    "# Path to pybind11 git repository (smartholder branch)\n",
    "pybind11_sh_path = '/home/wjakob/pybind11_sh/include'\n",
    "\n",
    "# Path to boost (in this case, assumed to be installed by the OS)\n",
    "boost_path = '/usr/include/boost'\n",
    "\n",
    "cmd_base = ['clang++', '-march=native', '-shared', '-rpath', '..', '-std=c++17', '-I', '../include', '-I', gp()['include'],\n",
    "            '-Wno-deprecated-declarations', '-fPIC', f'-L{boost_path}/stage/lib', '-L..', '-fno-stack-protector',\n",
    "            '-DPYBIND11_USE_SMART_HOLDER_AS_DEFAULT']\n",
    "\n",
    "def gen_file(name, func, libs=('cython', 'boost', 'pybind11', 'pybind11_sh', 'nanobind')):\n",
    "    for i, lib in enumerate(libs):    \n",
    "        for opt_mode, opt_flags in {'debug' : ['-O0', '-g3'], 'opt' : ['-Os', '-g0']}.items():\n",
    "            if lib != 'cython':\n",
    "                fname = name + '_' + lib + '.cpp'\n",
    "            else:\n",
    "                fname = name + '_' + lib + '_' + opt_mode + '.pyx'\n",
    "            \n",
    "            with open(fname, 'w') as f:\n",
    "                if lib == 'boost':\n",
    "                    f.write(f'#include <boost/python.hpp>\\n')\n",
    "                    f.write(f'namespace py = boost::python;\\n\\n')\n",
    "                    f.write(f'BOOST_PYTHON_MODULE({name}_{lib}_{opt_mode}) {{\\n')\n",
    "                elif lib == 'nanobind':\n",
    "                    f.write(f'#include <nanobind/nanobind.h>\\n\\n')\n",
    "                    f.write(f'namespace py = nanobind;\\n\\n')\n",
    "                    f.write(f'NB_MODULE({name}_{lib}_{opt_mode}, m) {{\\n')\n",
    "                elif lib.startswith('pybind11'):\n",
    "                    f.write(f'#include <pybind11/pybind11.h>\\n\\n')\n",
    "                    f.write(f'namespace py = pybind11;\\n\\n')\n",
    "                    f.write(f'PYBIND11_MODULE({name}_{lib}_{opt_mode}, m) {{\\n')\n",
    "                elif lib == 'cython':\n",
    "                    f.write(f'from libc.stdint cimport uint16_t, int32_t, uint32_t, int64_t, uint64_t\\n')\n",
    "\n",
    "                func(f, lib)\n",
    "                if lib != 'cython':\n",
    "                    f.write(f'}}\\n')\n",
    "\n",
    "            fname_out = name + '_' + lib + '_' + opt_mode  + suffix\n",
    "            cmd = cmd_base + opt_flags + [name + '_' + lib + '.cpp', '-o', fname_out]\n",
    "            if lib == 'nanobind':\n",
    "                cmd += ['-lnanobind']\n",
    "            elif lib == 'boost':\n",
    "                cmd += ['-I', boost_path, '-lboost_python310']\n",
    "            elif lib == 'pybind11':\n",
    "                cmd += ['-I', pybind11_path]\n",
    "            elif lib == 'pybind11_sh':\n",
    "                cmd += ['-I', pybind11_sh_path]\n",
    "                \n",
    "            print(' '.join(cmd))\n",
    "            time_list = []\n",
    "            for l in range(5):\n",
    "                time_before = time.perf_counter()\n",
    "                if lib == 'cython':\n",
    "                    subprocess.check_call(['cython3', '-3',  '--cplus', fname, '-o', name + '_' + lib + '.cpp'])\n",
    "                subprocess.check_call(cmd)\n",
    "                time_after = time.perf_counter()\n",
    "                time_list.append(time_after-time_before)\n",
    "            time_list.sort()\n",
    "                \n",
    "            if opt_mode != 'debug':\n",
    "                subprocess.check_call(['strip', fname_out])\n",
    "            if i == 0:\n",
    "                experiment_labels.append(name + ' [' + opt_mode + ']')\n",
    "            sizes[lib].append(os.path.getsize(fname_out) / (1024 * 1024))\n",
    "            times[lib].append(time_list[len(time_list)//2])\n",
    "\n",
    "\n",
    "            \n",
    "def gen_func(f, lib):\n",
    "    types = [ 'uint16_t', 'int32_t', 'uint32_t', 'int64_t', 'uint64_t', 'float' ]\n",
    "    if lib == 'boost':\n",
    "        prefix = 'py::'\n",
    "    else:\n",
    "        prefix = 'm.'\n",
    "    for i, t in enumerate(itertools.permutations(types)):\n",
    "        args = f'{t[0]} a, {t[1]} b, {t[2]} c, {t[3]} d, {t[4]} e, {t[5]} f'\n",
    "        if lib != 'cython':\n",
    "            f.write('    %sdef(\"test_%04i\", +[](%s) { return a+b+c+d+e+f; });\\n' % (prefix, i, args))\n",
    "        else:\n",
    "            f.write('cpdef float test_%04i(%s):\\n    return a+b+c+d+e+f\\n\\n' % (i, args))\n",
    "\n",
    "\n",
    "def gen_class(f, lib):\n",
    "    types = [ 'uint16_t', 'int32_t', 'uint32_t', 'int64_t', 'uint64_t', 'float' ]\n",
    "\n",
    "    for i, t in enumerate(itertools.permutations(types)):\n",
    "        if lib == 'boost':\n",
    "            prefix = ''\n",
    "            postfix = f', py::init<{t[0]}, {t[1]}, {t[2]}, {t[3]}, {t[4]}, {t[4]}>()'\n",
    "            func_prefix = 'py::def'\n",
    "\n",
    "        else:\n",
    "            prefix = 'm, '\n",
    "            postfix = ''\n",
    "            func_prefix = 'm.def'\n",
    "\n",
    "        if lib != 'cython':\n",
    "            f.write(f'    struct Struct{i} {{\\n')\n",
    "            f.write(f'        {t[0]} a; {t[1]} b; {t[2]} c; {t[3]} d; {t[4]} e; {t[5]} f;\\n')\n",
    "            f.write(f'        Struct{i}({t[0]} a, {t[1]} b, {t[2]} c, {t[3]} d, {t[4]} e, {t[5]} f) : a(a), b(b), c(c), d(d), e(e), f(f) {{ }}\\n')\n",
    "            f.write(f'        float sum() const {{ return a+b+c+d+e+f; }}\\n')\n",
    "            f.write(f'    }};\\n')\n",
    "        else:\n",
    "            f.write(f'cdef class Struct{i}:\\n')\n",
    "            f.write(f'    cdef {t[0]} a\\n')\n",
    "            f.write(f'    cdef {t[1]} b\\n')\n",
    "            f.write(f'    cdef {t[2]} c\\n')\n",
    "            f.write(f'    cdef {t[3]} d\\n')\n",
    "            f.write(f'    cdef {t[4]} e\\n')\n",
    "            f.write(f'    cdef {t[5]} f\\n\\n')\n",
    "            f.write(f'    def __cinit__(self, {t[0]} a, {t[1]} b, {t[2]} c, {t[3]} d, {t[4]} e, {t[5]} f):\\n')\n",
    "            f.write(f'        self.a = a\\n')\n",
    "            f.write(f'        self.b = b\\n')\n",
    "            f.write(f'        self.c = c\\n')\n",
    "            f.write(f'        self.d = d\\n')\n",
    "            f.write(f'        self.e = e\\n')\n",
    "            f.write(f'        self.f = f\\n\\n')\n",
    "            f.write(f'    cpdef float sum(self):\\n')\n",
    "            f.write(f'        return self.a+self.b+self.c+self.d+self.e+self.f\\n\\n')\n",
    "            continue\n",
    "\n",
    "        f.write(f'    py::class_<Struct{i}>({prefix}\\\"Struct{i}\\\"{postfix})\\n')\n",
    "        \n",
    "        if lib != 'boost':\n",
    "                f.write(f'        .def(py::init<{t[0]}, {t[1]}, {t[2]}, {t[3]}, {t[4]}, {t[5]}>())\\n')\n",
    "        f.write(f'        .def(\"sum\", &Struct{i}::sum);\\n\\n')\n",
    "        \n",
    "        if i > 250:\n",
    "            break;\n",
    "        \n",
    "        \n",
    "gen_file('func', gen_func)\n",
    "gen_file('class', gen_class)\n",
    "experiment_labels = ['func [debug]', 'func [opt]', 'class [debug]', 'class [opt]']\n",
    "\n",
    "print(experiment_labels)\n",
    "print(dict(sizes))\n",
    "print(dict(times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67eac20-ffb5-4c58-ba08-64e7404a54a9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "plot_colors = {\n",
    "    'boost': cycle[1],\n",
    "    'pybind11': cycle[3],\n",
    "    'pybind11_sh': cycle[5],\n",
    "    'cython' : cycle[4],\n",
    "    'nanobind': cycle[0]\n",
    "}\n",
    "plot_labels = {\n",
    "    'boost' : 'Boost.Python',\n",
    "    'pybind11' : 'pybind11',\n",
    "    'pybind11_sh' : 'pybind11 + smart_holder',\n",
    "    'cython' : 'Cython',\n",
    "    'nanobind' : 'nanobind'\n",
    "}\n",
    "\n",
    "def bars(data, ylim_scale = 1, figsize_scale = 1, width_scale=1.0, debug_shift=0.1):\n",
    "    ylim = 0\n",
    "    for n, d in data.items():\n",
    "        if len(d) == 0:\n",
    "            continue\n",
    "        ylim = max(max(d), ylim)\n",
    "    ylim *= ylim_scale * 1.3\n",
    "\n",
    "    def adj(ann):\n",
    "        for i, a in enumerate(ann):\n",
    "            if a.xy[1] > ylim*.9:\n",
    "                a.xy = (a.xy[0], ylim * 0.8)\n",
    "                if i%2 == 1:\n",
    "                    a.set_color('white')\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=[11.25*figsize_scale, 3*figsize_scale])\n",
    "    width = 1.0/(len(data) + 1)*width_scale\n",
    "    x = np.arange(4)\n",
    "\n",
    "    result = []\n",
    "    for i, n in enumerate(plot_labels):\n",
    "        d = data[n]\n",
    "        if len(d) == 0:\n",
    "            continue\n",
    "\n",
    "        col = plot_colors[n]\n",
    "        if col != 'None':\n",
    "            kwargs = { 'edgecolor': 'black', 'color': col }\n",
    "        else:\n",
    "            kwargs =  {'edgecolor': 'white', 'hatch' : '/', 'color':cycle[7]}\n",
    "            \n",
    "        bar = ax.bar(x+width*(i -(len(data)-1)/2), d, width, label=plot_labels[n], align='center', **kwargs)\n",
    "        result.append(bar)\n",
    "        \n",
    "    ax.add_patch(Rectangle((-0.65+debug_shift, -1), 1, 28, facecolor='white', alpha=.8, edgecolor='None'))\n",
    "    ax.add_patch(Rectangle((1.4+debug_shift, -1), 1, 25, facecolor='white', alpha=.8, edgecolor='None'))\n",
    "\n",
    "    for i, n in enumerate(plot_labels):\n",
    "        d = data[n]\n",
    "        if len(d) == 0:\n",
    "            continue\n",
    "\n",
    "        bar = result[i]\n",
    "        if n == 'nanobind':\n",
    "            adj(ax.bar_label(bar, fmt='%.2f'))\n",
    "        else:\n",
    "            improvement = np.array(d) / np.array(data['nanobind'])\n",
    "            improvement = ['%.2f\\n(x%.1f)' % (d[i], v) for i, v in enumerate(improvement)]\n",
    "            adj(ax.bar_label(bar, labels=improvement, padding=3))\n",
    "        \n",
    "    ax.set_ylim(0, ylim)\n",
    "    ax.set_xticks(x, experiment_labels)\n",
    "    return fig, ax\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60cb243a-e94d-4dd0-af67-91b7c5007571",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = bars(times, ylim_scale=0.93, figsize_scale=1.1, width_scale=1)\n",
    "ax.set_ylabel('Time (seconds)')\n",
    "ax.set_title('Compilation time')\n",
    "ax.set_xlim(-0.45,3.45)\n",
    "\n",
    "ax.legend(loc='upper left')\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.savefig('times.png', facecolor='white', dpi=200, bbox_inches='tight', pad_inches = 0)\n",
    "plt.savefig('times.svg', facecolor='white', bbox_inches='tight', pad_inches = 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b23e02b4-6e98-49a3-a60f-14142757af71",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = bars(sizes, ylim_scale=.085, figsize_scale=1.1)\n",
    "ax.set_ylabel('Size (MiB)')\n",
    "ax.set_title('Binary size')\n",
    "ax.set_xlim(-0.45,3.45)\n",
    "ax.legend(loc='lower left')\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.savefig('sizes.png', facecolor='white', dpi=200, bbox_inches='tight', pad_inches = 0)\n",
    "plt.savefig('sizes.svg', facecolor='white', bbox_inches='tight', pad_inches = 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f3b5d2-74d6-415c-ad3d-b8fb0f37eddb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import cppyy\n",
    "if not hasattr(cppyy.gbl, 'test_0000'):\n",
    "    cppyy.include('cppyy.h')\n",
    "\n",
    "plot_colors = {\n",
    "    'boost': cycle[1],\n",
    "    'cython': cycle[4],\n",
    "    'pybind11': cycle[3],\n",
    "    'cppyy' : cycle[8],\n",
    "    'python': 'None',\n",
    "    'nanobind': cycle[0]\n",
    "}\n",
    "\n",
    "plot_labels = {\n",
    "    'boost' : 'Boost.Python',\n",
    "    'pybind11' : 'pybind11',    \n",
    "    'cppyy' : 'cppyy',\n",
    "    'cython' : 'Cython',\n",
    "    'nanobind' : 'nanobind',\n",
    "    'python' : 'Python'\n",
    "}\n",
    "\n",
    "class native_module:\n",
    "    @staticmethod\n",
    "    def test_0000(a, b, c, d, e, f):\n",
    "        return a + b + c + d +e + f\n",
    "\n",
    "    \n",
    "    class Struct0:\n",
    "        def __init__(self, a, b, c, d, e, f):\n",
    "            self.a = a\n",
    "            self.b = b\n",
    "            self.c = c\n",
    "            self.d = d\n",
    "            self.e = e\n",
    "            self.f = f\n",
    "\n",
    "        def sum(self):\n",
    "            return self.a + self.b + self.c + self.e + self.f\n",
    "    \n",
    "\n",
    "rtimes = defaultdict(lambda: [])\n",
    "for name in ['func', 'class']:\n",
    "    its = 10000000 if name == 'func' else 2500000\n",
    "    for lib in plot_labels:\n",
    "        for mode in ['debug', 'opt']:\n",
    "            if lib == 'cppyy':\n",
    "                m = cppyy.gbl\n",
    "            elif lib == 'nanobind_sh':\n",
    "                continue # Performance identical, not an interesting data point\n",
    "            elif lib == 'python':\n",
    "                m = native_module\n",
    "            else:\n",
    "                m = importlib.import_module(f'{name}_{lib}_{mode}')\n",
    "         \n",
    "            time_list = []\n",
    "            for i in range(5):\n",
    "                time_before = time.perf_counter()\n",
    "                if name == 'func':\n",
    "                    func = m.test_0000\n",
    "                    for i in range(its):\n",
    "                        func(1,2,3,4,5,6)\n",
    "                elif name == 'class':\n",
    "                    cls = m.Struct0\n",
    "                    sum_member = cls.sum\n",
    "                    for i in range(its):\n",
    "                        sum_member(cls(1,2,3,4,5,6))\n",
    "\n",
    "                time_after = time.perf_counter()\n",
    "                time_list.append(time_after-time_before)\n",
    "            time_list.sort()\n",
    "\n",
    "            rtimes[lib].append(time_list[len(time_list)//2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ad743ce-33c4-4dd1-8a0f-2ccbcf02aa1c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = bars(rtimes, ylim_scale=.188, figsize_scale=1.25, width_scale=1, debug_shift=.1)\n",
    "ax.set_ylabel('Time (seconds)')\n",
    "ax.set_title('Runtime performance')\n",
    "ax.set_xlim(-0.45,3.45)\n",
    "ax.legend()\n",
    "fig.tight_layout()\n",
    "plt.savefig('perf.png', facecolor='white', dpi=200, bbox_inches='tight', pad_inches = 0)\n",
    "plt.savefig('perf.svg', facecolor='white', bbox_inches='tight', pad_inches = 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f54fb73-d150-48c0-862f-57abdbce9875",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
