Mjolnir

Mjolnir is a hybrid approach to partial evaluation / abstract interpretation, with an implementation in Julia. It can be thought of as a blend of operator-overloading based tracing (as in JAX, PyTorch Script, staged programming systems etc.) and dataflow-based abstract interpretation (as in the type inference systems of Julia, TypeScript and Crystal). It is aimed at package developers rather than Julia end-users.

Mjolnir can reproduce the compact, linear traces (aka computation graphs or Wengert lists) of tracing systems.

julia> function pow(x, n)
         r = 1
         while n > 0
           n -= 1
           r *= x
         end
         return r
       end
pow (generic function with 1 method)

julia> using Mjolnir

julia> @trace pow(::Int, 3)
1: (%1 :: const(pow), %2 :: Int64, %3 :: const(3))
  %4 = (*)(1, %2) :: Int64
  %5 = (*)(%4, %2) :: Int64
  %6 = (*)(%5, %2) :: Int64
  return %6

However, it avoids several of the downsides of those systems. It supports arbitrary Julia types (not just 'tensors' but also strings and structs). It supports value-dependent control flow (as it can encode branches in the trace). It supports side effects and mutating operators. Functions like println don't have to be evaluated at compile time. It can enforce its assumptions (i.e. referential transparency) rather than making the user responsible for them, and can generate diagnostics when there are issues. Mjolnir can thus compile a much wider range of Julia programs than OO approaches.

julia> @trace pow(::Int, ::Int)
1: (%1 :: const(pow), %2 :: Int64, %3 :: Int64)
  %4 = (>)(%3, 0) :: Bool
  br 3 (1) unless %4
  br 2 (%3, 1)
2: (%5 :: Int64, %6 :: Int64)
  %7 = (-)(%5, 1) :: Int64
  %8 = (*)(%6, %2) :: Int64
  %9 = (>)(%7, 0) :: Bool
  br 3 (%8) unless %9
  br 2 (%7, %8)
3: (%10 :: Int64)
  return %10

julia> function pow(x, n)
         r = 1
         while n > 0
           n -= 1
           r *= x
           @show r
         end
         return r
       end
pow (generic function with 1 method)

julia> @trace pow(2, 3)
1: (%1 :: const(pow), %2 :: const(2), %3 :: const(3))
  %4 = (println)("r = ", "2") :: Nothing
  %5 = (println)("r = ", "4") :: Nothing
  %6 = (println)("r = ", "8") :: Nothing
  return 8

Mjolnir is designed to be highly customisable, and to give as much control as possible to packages that build on it.