Prototype implementation of recursion

Project.toml

@@ -3,7 +3,7 @@ uuid = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 license = "MIT"
 desc = "Tape based task copying in Turing"
 repo = "https://github.com/TuringLang/Libtask.jl.git"
-version = "0.8.6"
+version = "0.8.7"
 
 [deps]
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"

perf/benchmark.jl

@@ -109,6 +109,22 @@ benchmark_driver!(neural_net, xs...)
 
 ####################################################################
 
+layer(w, x) = relu(w * x)
+Libtask.is_primitive(::typeof(layer), w, x) = false
+
+function neural_net(w1, w2, w3, x1, callback=nothing)
+ x2 = layer(w1, x1)
+ x3 = layer(w2, x2)
+ ret = sigmoid(LinearAlgebra.dot(w3, x3))
+ callback !== nothing && callback(ret)
+ return ret
+end
+
+xs = (randn(10,10), randn(10,10), randn(10), rand(10))
+benchmark_driver!(neural_net, xs...)
+
+####################################################################
+
 println("======= breakdown benchmark =======")
 
 x = rand(100000)

perf/p0.jl

@@ -1,5 +1,5 @@
 using Libtask
-using Turing, DynamicPPL, AdvancedPS
+using Turing, DynamicPPL, AdvancedPS, Random
 using BenchmarkTools
 
 @model gdemo(x, y) = begin
@@ -13,7 +13,9 @@ end
 
 
 # Case 1: Sample from the prior.
-m = Turing.Core.TracedModel(gdemo(1.5, 2.), SampleFromPrior(), VarInfo())
+m = Turing.Core.TracedModel(
+ gdemo(1.5, 2.), SampleFromPrior(), VarInfo(), MersenneTwister(123456)
+);
 f = m.evaluator[1];
 args = m.evaluator[2:end];
 
@@ -26,7 +28,9 @@ println("Run a tape...")
 @btime t.tf(args...)
 
 # Case 2: SMC sampler
-m = Turing.Core.TracedModel(gdemo(1.5, 2.), Sampler(SMC(50)), VarInfo());
+m = Turing.Core.TracedModel(
+ gdemo(1.5, 2.), Sampler(SMC(50)), VarInfo(), MersenneTwister(123456)
+);
 f = m.evaluator[1];
 args = m.evaluator[2:end];
 

perf/p2.jl

@@ -52,7 +52,8 @@ Random.seed!(2)
 iterations = 500
 model_fun = infiniteGMM(data)
 
-m = Turing.Core.TracedModel(model_fun, Sampler(SMC(50)), VarInfo())
+rng = MersenneTwister(123456)
+m = Turing.Core.TracedModel(model_fun, Sampler(SMC(50)), VarInfo(), rng)
 f = m.evaluator[1]
 args = m.evaluator[2:end]
 

src/tapedfunction.jl

@@ -162,7 +162,7 @@ function (tf::TapedFunction)(args...; callback=nothing, continuation=false)
  # run the raw tape
  while true
  ins = tf.tape[tf.counter]
- ins(tf)
+ ins(tf, callback)
  callback !== nothing && callback()
  tf.retval_binding_slot != 0 && break
  end
@@ -211,12 +211,17 @@ function Base.show(io::IO, instr::CondGotoInstruction)
  println(io, "CondGotoInstruction(", instr.condition, ", dest=", instr.dest, ")")
 end
 
-function (instr::Instruction{F})(tf::TapedFunction) where F
+function (instr::Instruction{F})(tf::TapedFunction, callback=nothing) where F
  # catch run-time exceptions / errors.
  try
  func = F === Int ? _lookup(tf, instr.func) : instr.func
  inputs = map(x -> _lookup(tf, x), instr.input)
- output = func(inputs...)
+ output = if is_primitive(func, inputs...)
+ func(inputs...)
+ else
+ tf_inner = TapedFunction(func, inputs..., cache=true)
+ tf_inner(inputs...; callback=callback)
+ end
  _update_var!(tf, instr.output, output)
  tf.counter += 1
  catch e
@@ -227,11 +232,11 @@ function (instr::Instruction{F})(tf::TapedFunction) where F
  end
 end
 
-function (instr::GotoInstruction)(tf::TapedFunction)
+function (instr::GotoInstruction)(tf::TapedFunction, callback=nothing)
  tf.counter = instr.dest
 end
 
-function (instr::CondGotoInstruction)(tf::TapedFunction)
+function (instr::CondGotoInstruction)(tf::TapedFunction, callback=nothing)
  cond = _lookup(tf, instr.condition)
  if cond
  tf.counter += 1
@@ -240,11 +245,11 @@ function (instr::CondGotoInstruction)(tf::TapedFunction)
  end
 end
 
-function (instr::ReturnInstruction)(tf::TapedFunction)
+function (instr::ReturnInstruction)(tf::TapedFunction, callback=nothing)
  tf.retval_binding_slot = instr.arg
 end
 
-function (instr::NOOPInstruction)(tf::TapedFunction)
+function (instr::NOOPInstruction)(tf::TapedFunction, callback=nothing)
  tf.counter += 1
 end
 
@@ -450,6 +455,18 @@ function translate!!(var, line, bindings, isconst, ir)
  throw(ErrorException("Unknown IR code"))
 end
 
+## primitives.
+
+"""
+ is_primitive(f, args...)
+
+Should a function be recursed into, or should it be treated as a single instruction, when
+encountered inside of a `TapedFunction`. If `is_primitive(f, args...)` is `true`, then
+the instruction will not be traced into. Conversely, if `is_primitive(f, args...)` is
+`false`, a `TapedFunction` is constructed.
+"""
+is_primitive(f, args...) = true
+
 ## copy Bindings, TapedFunction
 
 """

test/tf.jl

@@ -1,5 +1,10 @@
 using Libtask
 
+foo(x) = sin(cos(x))
+bar(x) = foo(foo(x))
+
+Libtask.is_primitive(::typeof(foo), args...) = false
+
 @testset "tapedfunction" begin
  # Test case 1: stack allocated objects are deep copied.
  @testset "Instruction{typeof(__new__)}" begin
@@ -31,4 +36,10 @@ using Libtask
 
  @test typeof(r) === Float64
  end
+ @testset "recurse into function" begin
+ tf = Libtask.TapedFunction(bar, 5.0)
+ count = 0
+ tf(4.0; callback=() -> (count += 1))
+ @test count == 9
+ end
 end