SnoopCompileBasics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Presenter: restart kernel and clear output!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg\n",
    "Pkg.activate(\".\")\n",
    "Pkg.precompile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# SnoopCompile basics\n",
    "\n",
    "## Part 3 of \"Package development: improving engineering quality & latency,\" JuliaCon 2021\n",
    "\n",
    "Tim Holy\n",
    "\n",
    "- Collecting \"snoop\" data\n",
    "- Graphical tools\n",
    "  + flamegraphs (e.g., profiling)\n",
    "  + profile-guided despecialization\n",
    "- `precompile` statement generation\n",
    "- More reasons to want inferrability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Comparing JET and SnoopCompile\n",
    "\n",
    "JET (see previous presentation) has both similarities and differences with [SnoopCompile](https://github.com/timholy/SnoopCompile.jl)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Most obviously, SnoopCompile focuses on *latency* (shortening \"time to first plot\") whereas JET focuses on *correctness*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "However, there are deep commonalities. A key reason is that both packages \"spy\" on type-inference to gather data about code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "But even here they have differences."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "```julia\n",
    "sum2(list) = [sum(list[1]), sum(list[2])]\n",
    "\n",
    "data = Any[[1,2,3], Any[1,2,3]]\n",
    "\n",
    "sum2(data)\n",
    "\n",
    "# Compare to one that JET does detect:\n",
    "# sum(Any[1,2,3])\n",
    "```\n",
    "\n",
    "The script above, when passed into JET, yields \"No errors\" despite the fact that the commented-out line would generate errors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "```julia\n",
    "sum2(::Vector{Any})\n",
    "├─ sum(::Vector{Int})\n",
    "└─ sum(::Vector{Any})\n",
    "```\n",
    "\n",
    "Both of the latter calls are made by *runtime dispatch*. But JET is a static analyzer---it uses types, not values---and it stops when the chain of inferrability breaks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "In contrast, SnoopCompile is a dynamic analyzer that acts more like a profiler:\n",
    "1. turn on snooping\n",
    "2. run some code\n",
    "3. turn off snooping\n",
    "4. return the snooping data\n",
    "\n",
    "Demo (run in a fresh session):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "using SnoopCompile\n",
    "sum2(list) = [sum(list[1]), sum(list[2])]\n",
    "tinf = tinfdemo2 = @snoopi_deep begin    \n",
    "    sum2(Any[[1,2,3], Any[1,2,3]])\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "using AbstractTrees\n",
    "print_tree(tinf, maxdepth=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Here you can see the call to `sum(::Vector{Any})` as a fresh entry into inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Graphic tools: the flamegraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "using ProfileSVG   # good for notebooks; use ProfileView from the REPL\n",
    "ProfileSVG.view(flamegraph(tinf))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "This profiles *inference*, not runtime performance. Width = inference time, height = inference depth. Empty spaces are when something else is happening (LLVM codegen, native codegen, or computation).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If you use ProfileView instead of ProfileSVG, you also get:\n",
    "- left-click: display the complete `MethodInstance` at the REPL\n",
    "- right-click (two-finger tap on a laptop): open the corresponding source file & line in your editor\n",
    "\n",
    "Just browsing the flamegraph can give you a lot of insight about what takes time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# A real-world example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "using Flux\n",
    "\n",
    "# From Flux's introductory documentation\n",
    "actual(x) = 4x + 2\n",
    "loss(predict, x, y) = Flux.Losses.mse(predict(x), y)\n",
    "\n",
    "x_train, x_test = hcat(0:5...), hcat(6:10...)\n",
    "y_train, y_test = actual.(x_train), actual.(x_test)\n",
    "\n",
    "tinf = tinfflux = @snoopi_deep begin\n",
    "    predict = Dense(1, 1)\n",
    "    parameters = params(predict)\n",
    "    Flux.train!((x, y) -> loss(predict, x, y), parameters, [(x_train, y_train)], Descent())\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "using ProfileSVG\n",
    "ProfileSVG.view(flamegraph(tinf); maxframes=50000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Observations:\n",
    "- approximately half the time was spent on inference\n",
    "- a lot is red: non-precompilable (maybe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Spoiler: just precompiling `Zygote._generate_pullback_via_decomposition(::Type)` shaves ~3s from the execution time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# When is specialization worthwhile? Profile-guided despecialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Profile\n",
    "@profile begin\n",
    "    predict = Dense(1, 1)\n",
    "    parameters = params(predict)\n",
    "    Flux.train!((x, y) -> loss(predict, x, y), parameters, [(x_train, y_train)], Descent())\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "using PyPlot: PyPlot, plt\n",
    "pgdsgui(tinf; by=exclusive)   # compare self runtime vs self inference time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Such plots can be useful in deciding when compiler specialization is worthwhile (if you're testing realistic workloads)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "(switch to live REPL demo, `pgds_demo.jl`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Despecialization is typically your most impactful way to reduce latency (see [details](https://pretalx.com/juliacon2021/talk/LE38LV/)):\n",
    "- where possible, reduce type-diversity\n",
    "- where possible, standardize argument types for slow-to-compile methods (handle input diversity through \"`convert` stub functions\" similarly to [this style recommendation](https://docs.julialang.org/en/v1/manual/style-guide/#Handle-excess-argument-diversity-in-the-caller))\n",
    "- annotate arguments with `@nospecialize` and possibly use `Base.inferencebarrier`\n",
    "- reduce optimization of entire module: `Base.Experimental.@compiler_options compile=min optimize=0 infer=false\n",
    "`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Precompile statement generation\n",
    "\n",
    "SnoopCompile can prepare lists of `precompile(f, types)` statements for incorporation into your package.\n",
    "\n",
    "Alternatively, you can just execute code during build time:\n",
    "\n",
    "```julia\n",
    "# Put this in your package code\n",
    "if ccall(:jl_generating_output, Cint, ()) == 1\n",
    "    # this runs only when we are precompiling the package\n",
    "    foo([1,2,3])   # precompiles `foo` for `Vector{Int}`\n",
    "end\n",
    "```\n",
    "\n",
    "But this isn't ideal if execution has side effects (e.g., plotting in a new window)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# Parcel precompile directives by package\n",
    "# For short output, we'll limit ourselves to just MethodInstances taking > 100ms\n",
    "# In practice, you might often choose a smaller threshold (~10ms)\n",
    "ttot, pcs = SnoopCompile.parcel(tinfflux; tmin=0.1);\n",
    "ttot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "pcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# Write all precompiles to package files\n",
    "SnoopCompile.write(\"/tmp/FluxPCs\", pcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "readlines(\"/tmp/FluxPCs/precompile_Zygote.jl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "tmod, pcmod = pcs[end].second;\n",
    "mi = pcmod[end][2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Uh-oh. This is hard to make consistent. Can we find something it called?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "ProfileSVG.view(flamegraph(tinfflux); maxframes=50000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "So this one call accounts for 3s of inference time. Precompile it and profit!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Inference triggers\n",
    "\n",
    "Flames that go all the way down to the bottom are fresh entry points into inference due to runtime dispatch: recall that\n",
    "```julia\n",
    "sum2(list) = [sum(list[1]), sum(list[2])]\n",
    "tinf = tinfdemo2 = @snoopi_deep begin    \n",
    "    sum2(Any[[1,2,3], Any[1,2,3]])\n",
    "end\n",
    "```\n",
    "gave us (in a fresh session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tinfdemo2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "using ProfileSVG\n",
    "ProfileSVG.view(flamegraph(tinfdemo2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Runtime inference -> no backedges -> need another `precompile` statement."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "\n",
    "The more separate nontrivial flames you have, the more `precompile` statements you'll need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Problem 1: what if you don't \"own\" the method? `sum(::AbstractVector)` belongs to `Base`, not your package"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Problem 2: you might own the method but not the types..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "tinfflux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ProfileSVG.view(flamegraph(tinfflux); maxframes=50000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The red bars depend on a user-specific function:\n",
    "```julia\n",
    "loss(predict, x, y) = Flux.Losses.mse(predict(x), y)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "There is no way for packages to anticipate every possible user function. Unsolvable?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Solve it in the *user's* application: it knows about both Flux and the specific functions of interest, so if calls are inferrable then *the whole stack becomes precompilable*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Our main remaining topic: improving inferrability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Julia 1.7.0-pre3-ns 1.7.0-beta3",
   "language": "julia",
   "name": "julia-1.7.0-pre3-ns-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}