tape instruction support

Author: KDr2

src/tapedfunction.jl

@@ -1,30 +1,46 @@
-mutable struct Instruction{F}
+abstract type AbstractInstruction end
+
+mutable struct Tape
+    tape::Vector{AbstractInstruction}
+    counter::Int
+    owner
+end
+
+mutable struct Instruction{F} <: AbstractInstruction
     fun::F
     input::Tuple
     output
-    tape
+    tape::Tape
 end
 
-
-mutable struct Tape
-    tape::Vector{Instruction}
-    owner
+mutable struct TapeInstruction <: AbstractInstruction
+    subtape::Tape
+    tape::Tape
 end
 
-Tape() = Tape(Vector{Instruction}(), nothing)
-Tape(owner) = Tape(Vector{Instruction}(), owner)
+Tape() = Tape(Vector{AbstractInstruction}(), 1, nothing)
+Tape(owner) = Tape(Vector{AbstractInstruction}(), 1, owner)
 MacroTools.@forward Tape.tape Base.iterate, Base.length
 MacroTools.@forward Tape.tape Base.push!, Base.getindex, Base.lastindex
 const NULL_TAPE = Tape()
 
+function setowner!(tape::Tape, owner)
+    tape.owner = owner
+    for ins in tape
+        isa(ins, TapeInstruction) && setowner!(ins.subtape, owner)
+    end
+end
+
 mutable struct Box{T}
     val::T
 end
 
 val(x) = x
 val(x::Box) = x.val
 box(x) = Box(x)
+box(x::Box) = x
 any_box(x) = Box{Any}(x)
+any_box(x::Box) = Box{Any}(x.val)
 
 gettape(x) = nothing
 gettape(x::Instruction) = x.tape
@@ -45,6 +61,12 @@ function Base.show(io::IO, instruction::Instruction)
     println(io, "Instruction($(fun)$(map(val, instruction.input)), tape=$(objectid(tape)))")
 end
 
+function Base.show(io::IO, ti::TapeInstruction)
+    subtape = ti.subtape
+    tape = ti.tape
+    println(io, "TapeInstruction($(subtape)), tape=$(objectid(tape)))")
+end
+
 function Base.show(io::IO, tp::Tape)
     buf = IOBuffer()
     print(buf, "$(length(tp))-element Tape")
@@ -63,25 +85,73 @@ function (instr::Instruction{F})() where F
     instr.output.val = output
 end
 
+function (instr::TapeInstruction)()
+    run(instr.subtape)
+end
+
 function run(tape::Tape, args...)
-    input = map(box, args)
-    tape[1].input = input
+    if length(args) > 0
+        input = map(box, args)
+        tape[1].input = input
+    end
     for instruction in tape
         instruction()
+        tape.counter += 1
     end
 end
 
+# if we should trace into a function
+# TODO:
+#     overload (instr::Instruction{F})() to specify
+#     which function to trace into
+function trace_into end
+trace_into(x) = false
+
 function run_and_record!(tape::Tape, f, args...)
     f = val(f) # f maybe a Boxed closure
-    output = try
-        box(f(map(val, args)...))
-    catch e
-        @warn e
-        any_box(nothing)
+    should_trace = trace_into(f)
+    if !should_trace
+        output = try
+            box(f(map(val, args)...))
+        catch e
+            @warn e
+            any_box(nothing)
+        end
+        ins = Instruction(f, args, output, tape)
+        push!(tape, ins)
+        return output
+    else
+        real_args = map(val, args)
+        ir = IRTools.@code_ir f(real_args...)
+        ir = intercept(ir; recorder=:run_and_record!)
+        # 1. we should distinguish fixed args and varargs here
+        arg_len = ir |> IRTools.arguments |> length
+        arg_len -= 2 # 1 for f, 1 for vargs
+        p_args = real_args[1:arg_len]
+        v_args = arg_len < 1 ? real_args : real_args[arg_len+1:end]
+        # detect if there is an vararg for f
+        if length(v_args) == 1
+            m = which(f, typeof(real_args))
+            no_vararg = @static if VERSION >= v"1.7"
+                (m.sig <: Tuple) && hasproperty(m.sig, :types) &&
+                    !isa(m.sig.types[end], Core.TypeofVararg)
+            else
+                (m.sig <: Tuple) && hasproperty(m.sig, :types) &&
+                    !(m.sig.types[end] <: Vararg{Any})
+            end
+            if no_vararg
+                v_args = v_args[1]
+            end
+        end
+        subtape = IRTools.evalir(ir, f, p_args..., v_args)
+        # 2. we should recover the args after getting the tape
+        #    to keep the chain complete
+        subtape[1].input = args
+        ins = TapeInstruction(subtape, tape)
+        output = subtape[end].output
+        push!(tape, ins)
+        return output
     end
-    ins = Instruction(f, args, output, tape)
-    push!(tape, ins)
-    return output
 end
 
 function dry_record!(tape::Tape, f, args...)
@@ -188,7 +258,7 @@ function (tf::TapedFunction)(args...)
         tape = IRTools.evalir(ir, tf.func, args...)
         tf.ir = ir
         tf.tape = tape
-        tape.owner = tf
+        setowner!(tape, tf)
         return result(tape)
     end
     # TODO: use cache

src/tapedtask.jl

@@ -1,33 +1,33 @@
 struct TapedTaskException
     exc
+    backtrace
 end
 
 struct TapedTask
     task::Task
     tf::TapedFunction
-    counter::Ref{Int}
     produce_ch::Channel{Any}
     consume_ch::Channel{Int}
     produced_val::Vector{Any}
 
     function TapedTask(
-        t::Task, tf::TapedFunction, counter, pch::Channel{Any}, cch::Channel{Int})
-        new(t, tf, counter, pch, cch, Any[])
+        t::Task, tf::TapedFunction, pch::Channel{Any}, cch::Channel{Int})
+        new(t, tf, pch, cch, Any[])
     end
 end
 
 function TapedTask(tf::TapedFunction, args...)
     tf.owner != nothing && error("TapedFunction is owned to another task.")
     # dry_run(tf)
     isempty(tf.tape) && tf(args...)
-    counter = Ref{Int}(1)
     produce_ch = Channel()
     consume_ch = Channel{Int}()
     task = @task try
-        step_in(tf, counter, args)
+        step_in(tf.tape, args)
     catch e
-        put!(produce_ch, TapedTaskException(e))
-        # @error "TapedTask Error: " exception=(e, catch_backtrace())
+        bt = catch_backtrace()
+        put!(produce_ch, TapedTaskException(e, bt))
+        # @error "TapedTask Error: " exception=(e, bt)
         rethrow()
     finally
         @static if VERSION >= v"1.4"
@@ -40,7 +40,7 @@ function TapedTask(tf::TapedFunction, args...)
         close(produce_ch)
         close(consume_ch)
     end
-    t = TapedTask(task, tf, counter, produce_ch, consume_ch)
+    t = TapedTask(task, tf, produce_ch, consume_ch)
     task.storage === nothing && (task.storage = IdDict())
     task.storage[:tapedtask] = t
     tf.owner = t
@@ -53,25 +53,42 @@ TapedTask(f, args...) = TapedTask(TapedFunction(f, arity=length(args)), args...)
 TapedTask(t::TapedTask, args...) = TapedTask(func(t), args...)
 func(t::TapedTask) = t.tf.func
 
-function step_in(tf::TapedFunction, counter::Ref{Int}, args)
-    len = length(tf.tape)
-    if(counter[] <= 1 && length(args) > 0)
+function step_in(t::Tape, args)
+    len = length(t)
+    if(t.counter <= 1 && length(args) > 0)
         input = map(box, args)
-        tf.tape[1].input = input
+        t[1].input = input
     end
-    while counter[] <= len
-        tf.tape[counter[]]()
+    while t.counter <= len
+        ins = t[t.counter]
+        if isa(ins, TapeInstruction)
+            step_in(ins.subtape, ())
+        else
+            ins()
+        end
         # produce and wait after an instruction is done
-        ttask = tf.owner
+        ttask = t.owner.owner
         if length(ttask.produced_val) > 0
             val = pop!(ttask.produced_val)
             put!(ttask.produce_ch, val)
             take!(ttask.consume_ch) # wait for next consumer
         end
-        counter[] += 1
+        t.counter += 1
     end
 end
 
+function increase_counter(t::Tape)
+    t.counter > length(t) && return
+    instr = t[t.counter]
+    if isa(instr, TapeInstruction)
+        increase_counter(instr.subtape)
+    else
+        # must be a produce instruction?
+        t.counter += 1
+    end
+end
+next_step(t::TapedTask) = increase_counter(t.tf.tape)
+
 #=
 # ** Approach (A) to implement `produce`:
 # Make`produce` a standalone instturction. This approach does NOT
@@ -94,14 +111,15 @@ function (instr::Instruction{typeof(produce)})()
 end
 =#
 
-
 # ** Approach (B) to implement `produce`:
 # This way has its caveat:
 # `produce` may deeply hide in an instruction, but not be an instruction
 # itself, and when we copy a task, the newly copied task will resume from
 # the instruction after the one which contains this `produce` call. If the
 # call to `produce` is not the last expression in the instuction, that
 # instruction will not be whole executed in the copied task.
+# With the abilty to trace into nested function call, we can minimize the
+# limitation of this caveat.
 @inline function is_in_tapedtask()
     ct = current_task()
     ct.storage === nothing && return false
@@ -114,6 +132,8 @@ end
 function produce(val)
     is_in_tapedtask() || return nothing
     ttask = current_task().storage[:tapedtask]
+    # put!(ttask.produce_ch, val)
+    # take!(ttask.consume_ch) # wait for next consumer
     length(ttask.produced_val) > 1 &&
         error("There is a produced value which is not consumed.")
     push!(ttask.produced_val, val)
@@ -186,18 +206,26 @@ function copy_box(old_box::Box{T}, roster::Dict{UInt64, Any}) where T
 end
 copy_box(o, roster::Dict{UInt64, Any}) = o
 
-function Base.copy(t::Tape)
+function Base.copy(x::Instruction, on_tape::Tape, roster::Dict{UInt64, Any})
+    input = map(x.input) do ob
+        copy_box(ob, roster)
+    end
+    output = copy_box(x.output, roster)
+    Instruction(x.fun, input, output, on_tape)
+end
+
+function Base.copy(x::TapeInstruction, on_tape::Tape, roster::Dict{UInt64, Any})
+    subtape = copy(x.subtape, roster)
+    TapeInstruction(subtape, on_tape)
+end
+
+function Base.copy(t::Tape, roster::Dict{UInt64, Any})
     old_data = t.tape
-    new_data = Vector{Instruction}()
-    new_tape = Tape(new_data, t.owner)
+    new_data = Vector{AbstractInstruction}()
+    new_tape = Tape(new_data, t.counter, t.owner)
 
-    roster = Dict{UInt64, Any}()
     for x in old_data
-        input = map(x.input) do ob
-            copy_box(ob, roster)
-        end
-        output = copy_box(x.output, roster)
-        new_ins = Instruction(x.fun, input, output, new_tape)
+        new_ins = copy(x, new_tape, roster)
         push!(new_data, new_ins)
     end
 
@@ -207,16 +235,18 @@ end
 function Base.copy(tf::TapedFunction)
     new_tf = TapedFunction(tf.func; arity=tf.arity)
     new_tf.ir = tf.ir
-    new_tape = copy(tf.tape)
-    new_tape.owner = new_tf
+    roster = Dict{UInt64, Any}()
+    new_tape = copy(tf.tape, roster)
+    setowner!(new_tape, new_tf)
     new_tf.tape = new_tape
     return new_tf
 end
 
 function Base.copy(t::TapedTask)
-    # t.counter[] <= 1 && error("Can't copy a TapedTask which is not running.")
     tf = copy(t.tf)
     new_t = TapedTask(tf)
-    new_t.counter[] = t.counter[] + 1
+    new_t.task.storage = copy(t.task.storage)
+    new_t.task.storage[:tapedtask] = new_t
+    next_step(new_t)
     return new_t
 end

test/runtests.jl

@@ -2,6 +2,7 @@ using Libtask
 using Test
 
 include("ctask.jl")
+include("special-instuctions.jl")
 include("tarray.jl")
 include("tref.jl")
 

test/special-instuctions.jl

@@ -0,0 +1,28 @@
+@testset "Special Instructions" begin
+
+    @testset "TapeInstruction" begin
+        i1(x) = i2(x)
+        i2(x) = produce(x)
+
+        Libtask.trace_into(::typeof(i1)) = true
+        Libtask.trace_into(::typeof(i2)) = true
+
+        function f()
+            t = 0
+            while t < 4
+                i1(t)
+                t = 1 + t
+            end
+        end
+
+        ctask = CTask(f)
+        @test consume(ctask) == 0
+        @test consume(ctask) == 1
+        a = copy(ctask)
+        @test consume(a) == 2
+        @test consume(a) == 3
+        @test consume(ctask) == 2
+        @test consume(ctask) == 3
+    end
+
+end