test_8308.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2dac2c6c-7219-44d1-9442-ccd083cc6f27",
   "metadata": {},
   "source": [
    "# Compiler-Target Switch\n",
    "\n",
    "Configure the target:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "813946b6-ca4c-4ede-a80a-7047ea6c9fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['POLIASTRO_TARGET'] = 'cuda'  # cpu, parallel, cuda\n",
    "os.environ['POLIASTRO_INLINE'] = 'no'  # yes, no"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "447c2a1e-f842-47ab-bf7f-0e19682ce65f",
   "metadata": {},
   "source": [
    "Variation of code from my current `poliastro` dev branch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9cb94749-124b-473b-9789-13a98c1686d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import warnings\n",
    "\n",
    "import numba as nb\n",
    "\n",
    "\n",
    "class JitWarning(UserWarning):\n",
    "    pass\n",
    "\n",
    "\n",
    "TARGET = os.environ.get('POLIASTRO_TARGET', 'cpu')\n",
    "\n",
    "if TARGET not in ('cpu', 'parallel', 'cuda'):  # numba 0.54.0, 19 August 2021, removed AMD ROCm target\n",
    "    raise ValueError(f'unknown target \"{TARGET:s}\"')\n",
    "if TARGET == 'parallel' and sys.maxsize <= 2**31:  # paying respect to poliastro#1399\n",
    "    raise ValueError('target \"parallel\" not supported on 32bit systems')\n",
    "if TARGET == 'cuda':\n",
    "    from numba import cuda  # explicit import required and only performed if target is switched to cuda\n",
    "\n",
    "INLINE = os.environ.get('POLIASTRO_INLINE', 'no')\n",
    "if INLINE not in ('yes', 'no'):\n",
    "    raise ValueError(f'unknown value for inline \"{INLINE:s}\"')\n",
    "INLINE = INLINE == 'yes'\n",
    "\n",
    "CACHE = os.environ.get('POLIASTRO_CACHE', 'no')\n",
    "if CACHE not in ('yes', 'no'):\n",
    "    raise ValueError(f'unknown value for cache \"{TARGET:s}\"')\n",
    "CACHE = CACHE == 'yes'\n",
    "if TARGET == 'cuda' and CACHE:\n",
    "    warnings.warn(\n",
    "        'caching is not supported for target \"cuda\"',\n",
    "        JitWarning,\n",
    "        stacklevel=2,\n",
    "    )\n",
    "\n",
    "PRECISIONS = ('f4', 'f8')  # TODO allow f2, i.e. half, for CUDA at least?\n",
    "\n",
    "NOPYTHON = True  # only for debugging, True by default\n",
    "\n",
    "_VECTOR = 'Tuple([f,f,f])'  # TODO hope for support of \"f[:]\" return values in cuda target\n",
    "_MATRIX = f'Tuple([{_VECTOR:s},{_VECTOR:s},{_VECTOR:s}])'  # TODO see above\n",
    "\n",
    "\n",
    "def _parse_signatures(signature):\n",
    "    \"\"\"\n",
    "    Automatically generate signatures for floats, vectors and matrices\n",
    "    \"\"\"\n",
    "    if '->' in signature:  # this is likely a layout for guvectorize\n",
    "        return signature\n",
    "    if not any(notation in signature for notation in ('f', 'V', 'M')):  # leave this signature as it is\n",
    "        return signature\n",
    "    if any(level in signature for level in PRECISIONS):  # leave this signature as it is\n",
    "        return signature\n",
    "    signature = signature.replace('V', _VECTOR)\n",
    "    signature = signature.replace('M', _MATRIX)\n",
    "    signature = [signature.replace('f', dtype) for dtype in PRECISIONS]\n",
    "    print(signature)\n",
    "    return signature\n",
    "\n",
    "\n",
    "def hjit(*args, **kwargs):\n",
    "    \"\"\"\n",
    "    Scalar helper, pre-configured, internal, switches compiler targets.\n",
    "    Functions decorated by it can only be called directly if TARGET is cpu or parallel.\n",
    "    \"\"\"\n",
    "\n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    if len(args) > 0 and isinstance(args[0], str):\n",
    "        args = _parse_signatures(args[0]), *args[1:]\n",
    "\n",
    "    cfg = {}\n",
    "    if TARGET in ('cpu', 'parallel'):\n",
    "        cfg.update({'nopython': NOPYTHON, 'inline': 'always' if INLINE else 'never', 'cache': CACHE})\n",
    "    if TARGET == 'cuda':\n",
    "        cfg.update({'device': True, 'inline': INLINE})\n",
    "    cfg.update(kwargs)\n",
    "\n",
    "    wjit = cuda.jit if TARGET == 'cuda' else nb.jit\n",
    "\n",
    "    def wrapper(func):\n",
    "        return wjit(\n",
    "            *args,\n",
    "            **cfg,\n",
    "        )(func)\n",
    "\n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "\n",
    "    return wrapper\n",
    "\n",
    "\n",
    "def vjit(*args, **kwargs):\n",
    "    \"\"\"\n",
    "    Vectorize on array, pre-configured, user-facing, switches compiler targets.\n",
    "    Functions decorated by it can always be called directly if needed.\n",
    "    \"\"\"\n",
    "\n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    if len(args) > 0 and isinstance(args[0], str):\n",
    "        args = _parse_signatures(args[0]), *args[1:]\n",
    "\n",
    "    cfg = {'target': TARGET}\n",
    "    if TARGET in ('cpu', 'parallel'):\n",
    "        cfg.update({'nopython': NOPYTHON, 'cache': CACHE})\n",
    "    cfg.update(kwargs)\n",
    "\n",
    "    def wrapper(func):\n",
    "        return nb.vectorize(\n",
    "            *args,\n",
    "            **cfg,\n",
    "        )(func)\n",
    "\n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "\n",
    "    return wrapper\n",
    "\n",
    "\n",
    "def gjit(*args, **kwargs):\n",
    "    \"\"\"\n",
    "    Generalized vectorize on array, pre-configured, user-facing, switches compiler targets.\n",
    "    Functions decorated by it can always be called directly if needed.\n",
    "    \"\"\"\n",
    "\n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    if len(args) > 0 and isinstance(args[0], str):\n",
    "        args = _parse_signatures(args[0]), *args[1:]\n",
    "\n",
    "    cfg = {'target': TARGET}\n",
    "    if TARGET in ('cpu', 'parallel'):\n",
    "        cfg.update({'nopython': NOPYTHON, 'cache': CACHE})\n",
    "    cfg.update(kwargs)\n",
    "\n",
    "    def wrapper(func):\n",
    "        return nb.guvectorize(\n",
    "            *args,\n",
    "            **cfg,\n",
    "        )(func)\n",
    "\n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "\n",
    "    return wrapper\n",
    "\n",
    "\n",
    "def jit(*args, **kwargs):\n",
    "    \"\"\"\n",
    "    Regular (n)jit, pre-configured, potentially user-facing, always CPU compiler target.\n",
    "    Functions decorated by it can be called directly.\n",
    "    \"\"\"\n",
    "\n",
    "    if len(args) == 1 and callable(args[0]):\n",
    "        func = args[0]\n",
    "        args = tuple()\n",
    "    else:\n",
    "        func = None\n",
    "\n",
    "    cfg = {'nopython': NOPYTHON, 'inline': 'never'}  # DOES NOT SWITCH INLINE TO PRESERVE OLD TESTED BEHAVIOR\n",
    "    cfg.update(kwargs)\n",
    "\n",
    "    def wrapper(func):\n",
    "\n",
    "        return nb.jit(\n",
    "            *args,\n",
    "            **cfg,\n",
    "        )(func)\n",
    "\n",
    "    if func is not None:\n",
    "        return wrapper(func)\n",
    "\n",
    "    return wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b17c6e1-58e9-44fa-adad-6f2ebff84396",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0d81c9c2-085d-44ab-ab7b-2c82a2ccf9b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['f4(f4)', 'f8(f8)']\n",
      "['f4(f4)', 'f8(f8)']\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from math import cos, sin\n",
    "\n",
    "COMPLEXITY = 2 ** 11\n",
    "SIZE = 2 ** 16\n",
    "\n",
    "@hjit('f(f)')\n",
    "def helper(scalar: float) -> float:\n",
    "    res: float = 0.0\n",
    "    for idx in range(COMPLEXITY):\n",
    "        if idx % 2 == round(scalar) % 2:\n",
    "            res += sin(idx)\n",
    "        else:\n",
    "            res -= cos(idx)\n",
    "    return res\n",
    "\n",
    "@vjit('f(f)')\n",
    "def test(d: float) -> float:\n",
    "    return helper(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5690a85-6260-464c-9b7c-639a5b83797c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ernst/Desktop/PROJEKTE/prj.TST2/github.numba/numba/cuda/dispatcher.py:495: NumbaPerformanceWarning: Grid size 103 will likely result in GPU under-utilization due to low occupancy.\n",
      "  warn(NumbaPerformanceWarning(msg))\n"
     ]
    },
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [4]\u001b[0m, in \u001b[0;36m<cell line: 5>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m data_f8 \u001b[38;5;241m=\u001b[39m (np\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mrandom(SIZE) \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m128\u001b[39m)\u001b[38;5;241m.\u001b[39mastype(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mf8\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      4\u001b[0m result_f4 \u001b[38;5;241m=\u001b[39m test(data_f4)\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m result_f4\u001b[38;5;241m.\u001b[39mdtype \u001b[38;5;241m==\u001b[39m np\u001b[38;5;241m.\u001b[39mfloat32\n\u001b[1;32m      7\u001b[0m result_f8 \u001b[38;5;241m=\u001b[39m test(data_f8)\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m result_f8\u001b[38;5;241m.\u001b[39mdtype \u001b[38;5;241m==\u001b[39m np\u001b[38;5;241m.\u001b[39mfloat64\n",
      "\u001b[0;31mAssertionError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "data_f4 = (np.random.random(SIZE) * 128).astype('f4')\n",
    "data_f8 = (np.random.random(SIZE) * 128).astype('f8')\n",
    "\n",
    "result_f4 = test(data_f4)\n",
    "assert result_f4.dtype == np.float32\n",
    "\n",
    "result_f8 = test(data_f8)\n",
    "assert result_f8.dtype == np.float64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46f44669-9b05-48ce-846c-b0b68c081b8a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bd3579a-5d5f-49ee-9107-504f5b735bb1",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}