GPU.js in jupyterlite.ipynb

{"metadata":{"language_info":{"codemirror_mode":{"name":"python","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8"},"kernelspec":{"name":"python","display_name":"Python (Pyodide)","language":"python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## `gpu.js` in `pyolite`\n\n`pyolite` runs inside a [WebWorker](https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API). While some features don't work there, a lot of really cool ones do.\n\n- `importScripts` can load arbitrary scripts from The Internet, or the local machine\n- the browser's GPU is often accessible there, depending on the browser and settings\n\nCombining these two, we can explore [gpu.js](https://github.com/gpujs/gpu.js), a JS library for working with GPUs, with fallback to pure JS.","metadata":{}},{"cell_type":"code","source":"import pyodide\n# basically `window` or `self`\nimport js\n# for managing types properly\nfrom pyodide.ffi import to_js\n# for the pretties\nfrom IPython.display import Markdown, JSON\nimport numpy","metadata":{"trusted":true},"execution_count":37,"outputs":[]},{"cell_type":"markdown","source":"## `importScripts`\n\n`js.importScripts` blocks and just bangs things right onto the global scope.","metadata":{}},{"cell_type":"code","source":"js.importScripts('https://cdn.jsdelivr.net/npm/gpu.js@latest/dist/gpu-browser.js')","metadata":{"trusted":true},"execution_count":38,"outputs":[]},{"cell_type":"markdown","source":"## `GPU`\n\nWe now have a `GPU` top-level object.","metadata":{}},{"cell_type":"code","source":"GPU = js.GPU","metadata":{"trusted":true},"execution_count":39,"outputs":[]},{"cell_type":"markdown","source":"but what can it do?","metadata":{}},{"cell_type":"code","source":"def gpu_capabilities():\n    checks = {\n        \"isGPUSupported\": \"GPU is in-fact supported\",\n        \"isKernelMapSupported\": \"kernel maps are supported\",\n        \"isOffscreenCanvasSupported\": \"offscreen canvas is supported\",\n        \"isWebGLSupported\": \"WebGL v1 is supported\",\n        \"isWebGL2Supported\": \"WebGL v2 is supported\",\n        \"isHeadlessGLSupported\": \"headlessgl is supported\",\n        \"isCanvasSupported\": \"canvas is supported\",\n        \"isGPUHTMLImageArraySupported\": \"the platform supports HTMLImageArray's\",\n        \"isSinglePrecisionSupported\": \"the system supports single precision float 32 values\",\n    }\n    results = [\n        \"### My Browser's `GPU` capabilities\",\n        f\"> agent string `{js.navigator.userAgent}`\" \n    ]\n    for check, message in checks.items():\n        x = \"x\" if getattr(GPU, check) else \" \"\n        results += [f\"- [{x}] {message}\"]\n    return Markdown(\"\\n\".join(results))\n\ngpu_capabilities()","metadata":{"trusted":true},"execution_count":40,"outputs":[{"execution_count":40,"output_type":"execute_result","data":{"text/plain":"<IPython.core.display.Markdown object>","text/markdown":"### My Browser's `GPU` capabilities\n> agent string `Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.5060.114 Safari/537.36`\n- [x] GPU is in-fact supported\n- [x] kernel maps are supported\n- [x] offscreen canvas is supported\n- [x] WebGL v1 is supported\n- [x] WebGL v2 is supported\n- [ ] headlessgl is supported\n- [ ] canvas is supported\n- [x] the platform supports HTMLImageArray's\n- [x] the system supports single precision float 32 values"},"metadata":{}}]},{"cell_type":"markdown","source":"## Making a `GPU` instance\n\nFirst, we need a `GPU`.","metadata":{}},{"cell_type":"code","source":"gpu = js.eval(\"new GPU()\")","metadata":{"trusted":true},"execution_count":41,"outputs":[]},{"cell_type":"markdown","source":"## Making a Kernel\n\nNext, we need a _GPU kernel_ (not to be confused with a Jupyter(Lite) kernel!). ","metadata":{}},{"cell_type":"markdown","source":"### Defining the kernel function\n\nFor starters, we'll just use good-ol `eval` to get a function, and hand that to the kernel. This will follow one of [`gpu.js`'s basic tests](https://github.com/gpujs/gpu.js/blob/develop/test/features/basic-math.js).","metadata":{}},{"cell_type":"code","source":"kernel_function = js.eval(\"\"\"function kernelFn(a, b) {\n    let sum = 0;\n    sum += a[this.thread.y][0] * b[0][this.thread.x];\n    sum += a[this.thread.y][1] * b[1][this.thread.x];\n    sum += a[this.thread.y][2] * b[2][this.thread.x];\n    return sum;\n}; kernelFn;\n\"\"\")\nkernel_function","metadata":{"trusted":true},"execution_count":42,"outputs":[{"execution_count":42,"output_type":"execute_result","data":{"text/plain":"function kernelFn(a, b) {\n    let sum = 0;\n    sum += a[this.thread.y][0] * b[0][this.thread.x];\n    sum += a[this.thread.y][1] * b[1][this.thread.x];\n    sum += a[this.thread.y][2] * b[2][this.thread.x];\n    return sum;\n}"},"metadata":{}}]},{"cell_type":"markdown","source":"### Defiing the options\n\nKernels take a number of options, but the required one is `output`, the shape of the output array. We condition the whole thing with `to_js`.","metadata":{}},{"cell_type":"code","source":"output = to_js([3, 3])","metadata":{"trusted":true},"execution_count":43,"outputs":[]},{"cell_type":"markdown","source":"### Instantiate the kernel","metadata":{}},{"cell_type":"code","source":"kernel = gpu.createKernel(kernel_function).setOutput(output)","metadata":{"trusted":true},"execution_count":44,"outputs":[]},{"cell_type":"markdown","source":"## Using the kernel","metadata":{}},{"cell_type":"markdown","source":"### Defining the inputs\n\nThese _also_ need to be properly conditioned with `to_js`.","metadata":{}},{"cell_type":"code","source":"a = to_js([\n    [1, 2, 3],\n    [4, 5, 6],\n    [7, 8, 9]\n])\nb = to_js([\n    [1, 2, 3],\n    [4, 5, 6],\n    [7, 8, 9]\n])","metadata":{"trusted":true},"execution_count":45,"outputs":[]},{"cell_type":"markdown","source":"### Calling the kernel","metadata":{}},{"cell_type":"code","source":"res = kernel(a, b)\nres","metadata":{"trusted":true},"execution_count":46,"outputs":[{"execution_count":46,"output_type":"execute_result","data":{"text/plain":"30,36,42,66,81,96,102,126,150"},"metadata":{}}]},{"cell_type":"markdown","source":"### Conditioning the output\n\nWhile it looks flat, there's actually some structure in there: it needs a little [conditioning](https://pyodide.org/en/stable/usage/type-conversions.html#using-javascript-typed-arrays-from-python) to be useful.","metadata":{}},{"cell_type":"code","source":"arr = numpy.asarray(res.to_py())\narr","metadata":{"trusted":true},"execution_count":47,"outputs":[{"execution_count":47,"output_type":"execute_result","data":{"text/plain":"array([[ 30.,  36.,  42.],\n       [ 66.,  81.,  96.],\n       [102., 126., 150.]], dtype=float32)"},"metadata":{}}]},{"cell_type":"markdown","source":"## And that's \"it\"\n\nThis is just a quick look at using `importScripts` and `GPU.js`. There's loads more to do:\n\n- avoid calling out to a CDN\n- benchmarking\n- packaging this up as something useful","metadata":{}}]}