Idea: scope- (rather than type-) centric composable optimizations

I was inspired by this paper and the following observations to sketch out a model of compositional optimizations that I thought I would share with the community.

  • Simple interpretation can be faster than Julia’s per-function JIT for non-hot code paths.
  • Julia has done a lot of work to make compilation transparent, but it’s still hard to control/compose directly how optimizations are applied.
  • Optimizations are very much type-centric, i.e. dictated by types, rather than a combination of types and algorith/control flow.

I think the idea of using a lazy computation graph combined with scopes to declare and compose, which stops short of using a GUI to interactively apply high-level optimizations as in the paper, might be a viable long-term evolution path for Julia that:

  • Addresses unnecessary compilation and the time to first plot issue, since unadorned code can simply be interpreted with Python-like performance.
  • Prevents Julia turning into yet another opaque JIT like V8 that seems to be the current direction with interpretation + heuristic JIT application
  • Allows domain knowledge to dictate what optimizations are applied at the top.
  • Allows new optimizations to be packaged as libraries.
  • Optimizations can be written in Julia itself as pattern matches applied to the computation graph.
  • Does not intermix the mathematical/algorithmic description of programs versus their implementation using lower-level control flow primitives (for loops, etc.).
  • Is more deliberate and avoids pathological cases of having to dig through @code_warntype to satisfy the compiler about type stability.

In a nutshell, I think types are a great first step towards giving developers a nice way to negotiate optimization with the compiler but they are not the whole solution.

A tiny code example showing the concept:

@with UnrollLoop(; n=16), ParallelLoop()
  z = Conv2D(x, y)
  # nested optimization scope
  @with OtherOptimization()
     ...
  end
end

I have written a short blog post on the subject for more context if anyone is interested and would love to hear your thoughts and alternative ideas for extending Julia’s performance and expressiveness.

This seems relevant to the compiler hooks @Keno and @staticfloat have been working on.

Also you might want to look at https://github.com/MikeInnes/Mjolnir.jl which seems to do some of that as well.

2 Likes

Related future julia features:

  • The ability to set the optimization level on a per-module basis (coming on Julia 1.5) Pull request link
  • An infrastructure to add custom compiler passes (coming in Julia, 1.6) Pull request link
1 Like