{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy\n",
    "import numpy\n",
    "import math\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wait until conda gets 5.0 if you want RawKernel. Not fun to setup CUDA with libraries and build. All credit for original code codes to Jim Pivarski here: https://github.com/jpivarski/python-numpy-mini-course/blob/evaluated/7-gpu.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cupy.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "\n",
    "\n",
    "def prepare(height, width, numpy=numpy):\n",
    "    y, x = numpy.ogrid[-1 : 0 : height * 1j, -1.5 : 0 : width * 1j]\n",
    "    c = x + y * 1j\n",
    "    fractal = numpy.zeros(c.shape, dtype=numpy.int32)\n",
    "    return c, fractal\n",
    "\n",
    "\n",
    "def run(c, fractal, maxiterations=20):\n",
    "    fractal *= 0  # set fractal to maxiterations without replacing it\n",
    "    fractal += maxiterations\n",
    "    z = c\n",
    "    for i in range(maxiterations):\n",
    "        z **= 2\n",
    "        z += c\n",
    "        diverge = z.real ** 2 + z.imag ** 2 > 2 ** 2\n",
    "        z[diverge] = 2\n",
    "        diverge &= fractal == maxiterations\n",
    "        fractal[diverge] = i\n",
    "\n",
    "    return fractal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, fractal = prepare(8000, 12000, numpy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "_ = run(c, fractal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, fractal = prepare(8000, 12000, cupy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "_ = run(c, fractal)\n",
    "cupy.cuda.Stream.null.synchronize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cupy_single = cupy.ElementwiseKernel(\n",
    "    \"complex128 cpx, int32 maxiterations\",\n",
    "    \"int32 res\",\n",
    "    \"\"\"\n",
    "    res = maxiterations;\n",
    "    complex<double> z = cpx;\n",
    "\n",
    "    for (int i=0; i<maxiterations; i++) {\n",
    "        z = z*z + cpx;\n",
    "        \n",
    "        if(z.real()*z.real() + z.imag()*z.imag() > 4) {\n",
    "            res = i;\n",
    "            break;\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    \"\"\",\n",
    "    \"fract_el\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, _ = prepare(8000, 12000, cupy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "fractal = cupy_single(c, 20)\n",
    "cupy.cuda.Stream.null.synchronize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cupy_kernel = cupy.RawKernel(\n",
    "    \"\"\"\n",
    "extern \"C\" \n",
    "__global__ void fractal(double* c, int* fractal, int height, int width, int maxiterations) {\n",
    "    const int x = threadIdx.x + blockIdx.x*blockDim.x;\n",
    "    const int y = threadIdx.y + blockIdx.y*blockDim.y;\n",
    "    double creal = c[2 * (x + height*y)];\n",
    "    double cimag = c[2 * (x + height*y) + 1];\n",
    "    double zreal = creal;\n",
    "    double zimag = cimag;\n",
    "    fractal[x + height*y] = maxiterations;\n",
    "    for (int i = 0;  i < maxiterations;  i++) {\n",
    "        double zreal2 = zreal*zreal - zimag*zimag + creal;\n",
    "        double zimag2 = 2*zreal*zimag + cimag;\n",
    "        zreal = zreal2;\n",
    "        zimag = zimag2;\n",
    "        if (zreal*zreal + zimag*zimag > 4) {\n",
    "            fractal[x + height*y] = i;\n",
    "            break;\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\"\"\",\n",
    "    \"fractal\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_pycuda(height, width, maxiterations=20):\n",
    "    y, x = cupy.ogrid[-1 : 0 : height * 1j, -1.5 : 0 : width * 1j]\n",
    "    grid = (int(math.ceil(height / 32)), int(math.ceil(width / 32)))\n",
    "    c = x + y * 1j\n",
    "    fractal = cupy.empty(c.shape, dtype=cupy.int32) + maxiterations\n",
    "    cupy_kernel(\n",
    "        grid,\n",
    "        (32, 32, 1),\n",
    "        [\n",
    "            c.view(cupy.double),\n",
    "            fractal,\n",
    "            cupy.int32(height),\n",
    "            cupy.int32(width),\n",
    "            cupy.int32(maxiterations),\n",
    "        ],\n",
    "    )\n",
    "    return c, fractal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "_, fractal = run_pycuda(8000, 12000)\n",
    "cupy.cuda.Stream.null.synchronize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, fractal = run_pycuda(8000, 12000)\n",
    "plt.imshow(fractal.get())"
   ]
  }
 ],
 "metadata": {
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}