Peformance Optimization

Project.toml

@@ -10,6 +10,7 @@ CodeInfoTools = "bc773b8a-8374-437a-b9f2-0e9785855863"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]

src/Libtask.jl

@@ -3,6 +3,7 @@ module Libtask
 using CodeInfoTools
 using FunctionWrappers: FunctionWrapper
 using LRUCache
+using MacroTools
 
 export TapedTask, consume, produce
 export TArray, tzeros, tfill, TRef

src/tapedfunction.jl

@@ -62,7 +62,7 @@ struct TypedFunction{OT, IT<:Tuple}
     TypedFunction{OT, IT}(f::Function) where {OT, IT<:Tuple} = new{OT, IT}(f, Ref{OT}())
 end
 
-function (f::TypedFunction{OT, IT})(args...) where {OT, IT<:Tuple}
+@inline function (f::TypedFunction{OT, IT})(args...) where {OT, IT<:Tuple}
     output = f.func(args...)
     f.retval[] = OT === Nothing ? nothing : output
     return f.retval[]
@@ -71,21 +71,17 @@ end
 struct Box{T}
     id::Symbol
     get::TypedFunction{T, Tuple{TapedFunction, Symbol}}
-    set::TypedFunction{Nothing, Tuple{TapedFunction, Symbol, T}}
 
     function Box{T}(id::Symbol) where T
-        return new(id,
-                   TypedFunction{T, Tuple{TapedFunction, Symbol}}(_inner_getter),
-                   TypedFunction{Nothing, Tuple{TapedFunction, Symbol, T}}(_inner_setter))
+        return new(id, TypedFunction{T, Tuple{TapedFunction, Symbol}}(_inner_getter))
     end
 end
 
 @inline _inner_getter(tf::TapedFunction, v::Symbol) = tf.bindings[v]
-@inline _inner_setter(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
 @inline _lookup(tf::TapedFunction, v) = v
 @inline _lookup(tf::TapedFunction, v::Box{T}) where T = v.get(tf, v.id)
-@inline _update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
-@inline _update_var!(tf::TapedFunction, v::Box{T}, c::T) where T = v.set(tf, v.id, c)
+@inline _update_var!(tf::TapedFunction, v::Symbol, c) = (tf.bindings[v] = c; nothing)
+@inline _update_var!(tf::TapedFunction, v::Box{T}, c::T) where T = (tf.bindings[v.id] = c; nothing)
 
 """
     Instruction
@@ -113,9 +109,10 @@ end
 @inline val(x) = x
 @inline val(x::GlobalRef) = getproperty(x.mod, x.name)
 @inline val(x::QuoteNode) = eval(x)
-@inline val(x::TapedFunction) = x.func
 @inline result(t::TapedFunction) = t.bindings[t.retval]
 
+const SLOTS = [Symbol("_", i) for i in 1:100]
+
 function (tf::TapedFunction)(args...; callback=nothing, continuation=false)
     if !continuation # reset counter and retval to run from the start
         tf.counter = 1;
@@ -126,7 +123,7 @@ function (tf::TapedFunction)(args...; callback=nothing, continuation=false)
     if tf.counter <= 1
         haskey(tf.bindings, :_1) && _update_var!(tf, :_1, tf.func)
         for i in 1:length(args)
-            slot = Symbol("_", i + 1)
+            slot = i < length(SLOTS) ? SLOTS[i + 1] : Symbol("_", i + 1)
             haskey(tf.bindings, slot) && _update_var!(tf, slot, args[i])
         end
     end
@@ -179,20 +176,29 @@ function Base.show(io::IO, instr::GotoInstruction)
     println(io, "GotoInstruction(", instr.condition, ", dest=", instr.dest, ")")
 end
 
-function (instr::Instruction{F})(tf::TapedFunction) where F
-    # catch run-time exceptions / errors.
-    try
-        func = val(_lookup(tf, instr.func))
-        inputs = map(x -> val(_lookup(tf, x)), instr.input)
-        output = func(inputs...)
-        _update_var!(tf, instr.output, output)
-        tf.counter += 1
-    catch e
-        println("counter=", tf.counter)
-        println("tf=", tf)
-        println(e, catch_backtrace());
-        rethrow(e);
+@generated function (instr::Instruction{F, N})(tf::TapedFunction) where {F, N}
+    arity = instr.parameters[2]
+    body = quote
+        try
+            func = val(_lookup(tf, instr.func))
+            output = func() # will inject arguments later
+            _update_var!(tf, instr.output, output)
+            tf.counter += 1
+        catch e
+            # catch run-time exceptions / errors.
+            println("counter=", tf.counter)
+            println("tf=", tf)
+            println(e, catch_backtrace());
+            rethrow(e);
+        end
     end
+    body = MacroTools.striplines(body)
+    # inject arguments
+    call_args = body.args[1].args[1].args[2].args[2].args
+    for i in 1:arity
+        push!(call_args, :(val(_lookup(tf, instr.input[$i]))))
+    end
+    body
 end
 
 function (instr::GotoInstruction)(tf::TapedFunction)
@@ -257,39 +263,57 @@ function translate!(tape::RawTape, ir::Core.CodeInfo)
 
     for (idx, line) in enumerate(ir.code)
         isa(line, Core.Const) && (line = line.val) # unbox Core.Const
-        ins = translate!!(Core.SSAValue(idx), line, bindings, ir)
+        isconst = isa(ir.ssavaluetypes[idx], Core.Const)
+        ins = translate!!(Core.SSAValue(idx), line, bindings, isconst, ir)
         push!(tape, ins)
     end
     return (bindings, tape)
 end
 
 const IRVar = Union{Core.SSAValue, Core.SlotNumber}
 
+function _const_instruction(var::IRVar, v, bindings::Dict{Symbol, Any}, ir)
+    if isa(var, Core.SSAValue)
+        box = bind_var!(var, bindings, ir)
+        bindings[box.id] = v
+        return GotoInstruction(Box{Bool}(:_true), 0) # NOOP
+    end
+    return Instruction(identity, (bind_var!(v, bindings, ir),), bind_var!(var, bindings, ir))
+end
+
 function translate!!(var::IRVar, line::Core.NewvarNode,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+                     bindings::Dict{Symbol, Any}, isconst::Bool, @nospecialize(ir))
     # use a noop to ensure the 1-to-1 mapping from ir.code to instructions
     # on tape. see GotoInstruction.dest.
     return GotoInstruction(Box{Bool}(:_true), 0)
 end
 
 function translate!!(var::IRVar, line::GlobalRef,
-                     bindings::Dict{Symbol, Any}, ir)
+                     bindings::Dict{Symbol, Any}, isconst::Bool, ir)
+    if isconst
+        v = ir.ssavaluetypes[var.id].val
+        return _const_instruction(var, v, bindings, ir)
+    end
     return Instruction(() -> val(line), (), bind_var!(var, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Core.SlotNumber,
-                     bindings::Dict{Symbol, Any}, ir)
+                     bindings::Dict{Symbol, Any}, isconst::Bool, ir)
+    if isconst
+        v = ir.ssavaluetypes[var.id].val
+        return _const_instruction(var, v, bindings, ir)
+    end
     return Instruction(identity, (bind_var!(line, bindings, ir),), bind_var!(var, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Core.TypedSlot,
-                     bindings::Dict{Symbol, Any}, ir)
+                     bindings::Dict{Symbol, Any}, isconst::Bool, ir)
     input_box = bind_var!(Core.SlotNumber(line.id), bindings, ir)
     return Instruction(identity, (input_box,), bind_var!(var, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Core.GotoIfNot,
-                     bindings::Dict{Symbol, Any}, ir)
+                     bindings::Dict{Symbol, Any}, isconst::Bool, ir)
     _cond = bind_var!(line.cond, bindings, ir)
     cond = if isa(_cond, Bool)
         Box{Bool}(_cond ? :_true : :_false)
@@ -300,29 +324,35 @@ function translate!!(var::IRVar, line::Core.GotoIfNot,
 end
 
 function translate!!(var::IRVar, line::Core.GotoNode,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+                     bindings::Dict{Symbol, Any}, isconst::Bool, @nospecialize(ir))
     return GotoInstruction(Box{Bool}(:_false), line.label)
 end
 
 function translate!!(var::IRVar, line::Core.ReturnNode,
-                     bindings::Dict{Symbol, Any}, ir)
+                     bindings::Dict{Symbol, Any}, isconst::Bool, ir)
     return ReturnInstruction(bind_var!(line.val, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Expr,
-                     bindings::Dict{Symbol, Any}, ir::Core.CodeInfo)
+                     bindings::Dict{Symbol, Any}, isconst::Bool, ir::Core.CodeInfo)
     head = line.head
     _bind_fn = (x) -> bind_var!(x, bindings, ir)
     if head === :new
         args = map(_bind_fn, line.args)
         return Instruction(__new__, args |> Tuple, _bind_fn(var))
     elseif head === :call
+        if isconst
+            v = ir.ssavaluetypes[var.id].val
+            return _const_instruction(var, v, bindings, ir)
+        end
         args = map(_bind_fn, line.args)
         # args[1] is the function
         func = line.args[1]
         if Meta.isexpr(func, :static_parameter) # func is a type parameter
             func = ir.parent.sparam_vals[func.args[1]]
-        else # isa(func, GlobalRef) or a var?
+        elseif isa(func, GlobalRef)
+            func = val(func)
+        else # a var?
             func = args[1] # a var(box)
         end
         return Instruction(func, args[2:end] |> Tuple, _bind_fn(var))
@@ -331,8 +361,12 @@ function translate!!(var::IRVar, line::Expr,
         lhs = line.args[1]
         rhs = line.args[2] # the right hand side, maybe a Expr, or a var, or ...
         if Meta.isexpr(rhs, (:new, :call))
-            return translate!!(lhs, rhs, bindings, ir)
+            return translate!!(lhs, rhs, bindings, false, ir)
         else # rhs is a single value
+            if isconst
+                v = ir.ssavaluetypes[var.id].val
+                return Instruction(identity, (_bind_fn(v),), _bind_fn(lhs))
+            end
             return Instruction(identity, (_bind_fn(rhs),), _bind_fn(lhs))
         end
     else
@@ -366,7 +400,7 @@ tape_copy(x::Core.Box) = Core.Box(tape_copy(x.contents))
 # tape_copy(x::Dict) = deepcopy(x)
 
 function copy_bindings(old::Dict{Symbol, Any})
-    newb = Dict{Symbol, Any}()
+    newb = copy(old)
     for (k, v) in old
         newb[k] = tape_copy(v)
     end