remove Amb, add example test

Author: shashi

src/controlflow.jl

@@ -1,4 +1,110 @@
-using Amb
+#### Path
+
+# First just do it for the case where there we assume
+# tainted gotoifnots do not go in a loop!
+# TODO: write a thing to detect this! (overdub predicates only in tainted ifs)
+# implement snapshotting function state as an optimization for branch exploration
+mutable struct Path
+    path::BitVector
+    cursor::Int
+end
+
+Path() = Path([], 1)
+
+function increment!(bitvec)
+    for i=1:length(bitvec)
+        if bitvec[i] === true
+            bitvec[i] = false
+        else
+            bitvec[i] = true
+            break
+        end
+    end
+end
+
+function reset!(p::Path)
+    p.cursor=1
+    increment!(p.path)
+    nothing
+end
+
+function alldone(p::Path) # must be called at the end of the function!
+    all(identity, p.path)
+end
+
+function current_predicate!(p::Path)
+    if p.cursor > length(p.path)
+        push!(p.path, false)
+    else
+        p.path[p.cursor]
+    end
+    val = p.path[p.cursor]
+    p.cursor+=1
+    val
+end
+
+alldone(c) = alldone(c.metadata[2])
+reset!(c) = reset!(c.metadata[2])
+current_predicate!(c) = current_predicate!(c.metadata[2])
+
+#=
+julia> p=Path()
+Path(Bool[], 1)
+
+julia> alldone(p) # must be called at the end of a full run
+true
+
+julia> current_predicate!(p)
+false
+
+julia> alldone(p) # must be called at the end of a full run
+false
+
+julia> current_predicate!(p)
+false
+
+julia> p
+Path(Bool[false, false], 3)
+
+julia> alldone(p)
+false
+
+julia> reset!(p)
+
+julia> p
+Path(Bool[true, false], 1)
+
+julia> current_predicate!(p)
+true
+
+julia> current_predicate!(p)
+false
+
+julia> alldone(p)
+false
+
+julia> reset!(p)
+
+julia> p
+Path(Bool[false, true], 1)
+
+julia> current_predicate!(p)
+false
+
+julia> current_predicate!(p)
+true
+
+julia> reset!(p)
+
+julia> current_predicate!(p)
+true
+
+julia> current_predicate!(p)
+true
+
+julia> alldone(p)
+true
+=#
 
 """
 `abstract_run(g, ctx, overdubbed_fn, args...)`
@@ -32,9 +138,17 @@ end
 # do something to merge metadata from all the runs
 ```
 """
-function abstract_run(acc, ctx::Cassette.Context, overdub_fn, args...)
-    pass_ctx = Cassette.similarcontext(ctx, pass=AbsintPass)
-        acc(Cassette.overdub(pass_ctx, overdub_fn, args...))
+function abstract_run(acc, ctx::Cassette.Context, overdub_fn, args...; verbose=true)
+    path = Path()
+    pass_ctx = Cassette.similarcontext(ctx, metadata=(ctx.metadata, path), pass=AbsintPass)
+
+    while true
+        acc(Cassette.recurse(pass_ctx, ()->overdub_fn(args...)))
+
+        verbose && println("Explored path: ", path)
+        alldone(path) && break
+        reset!(path)
+    end
 end
 
 """
@@ -47,8 +161,6 @@ function istainted(ctx, cond)
           " See docs for `istainted`.")
 end
 
-_choice() = (@amb true false)
-
 # Must return 7 exprs
 function rewrite_branch(ctx, stmt, extraslot, i)
     # turn
@@ -70,9 +182,9 @@ function rewrite_branch(ctx, stmt, extraslot, i)
     # not tainted? jump to the penultimate statement
     push!(exprs, Expr(:gotoifnot, istainted_ssa, i+5))
 
-    # tainted? then use this_here_predicate!(SSAValue(1))
+    # tainted? then use current_predicate!(SSAValue(1))
     current_pred = i+2
-    push!(exprs, :($_choice()))
+    push!(exprs, :($(Expr(:nooverdub, current_predicate!))($(Expr(:contextslot)))))
 
     # Store the interpreter-provided predicate in the slot
     push!(exprs, Expr(:(=), extraslot, SSAValue(i+2)))
@@ -93,7 +205,7 @@ function rewrite_ir(ctx, ref)
     # turn
     #   <val> ? t : f
     # into
-    #   istainted(<val>) ? this_here_predicate!(p) : <val> ? t : f
+    #   istainted(<val>) ? current_predicate!(p) : <val> ? t : f
 
     ir = ref.code_info
     ir = copy(ir)

src/jacobian.jl

@@ -25,22 +25,16 @@ function Base.push!(S::Sparsity, i::Int, j::Int)
     push!(S.J, j)
 end
 
-struct ProvinanceSet{T}
-    set::T # Set, Array, Int, Tuple, anything!
+struct ProvinanceSet
+    set::Set{Int}
+    ProvinanceSet(s::Set) = new(s)
+    ProvinanceSet(s) = new(Set(s))
 end
 
 # note: this is not strictly set union, just some efficient way of concating
 Base.union(p::ProvinanceSet, ::Cassette.NoMetaData) = p
 Base.union(::Cassette.NoMetaData, p::ProvinanceSet) = p
 
-Base.union(p::ProvinanceSet{<:Tuple},
-           q::ProvinanceSet{<:Integer}) = ProvinanceSet((p.set..., q.set,))
-Base.union(p::ProvinanceSet{<:Integer},
-           q::ProvinanceSet{<:Tuple}) = ProvinanceSet((p.set, q.set...,))
-Base.union(p::ProvinanceSet{<:Integer},
-           q::ProvinanceSet{<:Integer}) = ProvinanceSet((p.set, q.set,))
-Base.union(p::ProvinanceSet{<:Tuple},
-           q::ProvinanceSet{<:Tuple}) = ProvinanceSet((p.set..., q.set...,))
 Base.union(p::ProvinanceSet,
            q::ProvinanceSet) = ProvinanceSet(union(p.set, q.set))
 Base.union(p::ProvinanceSet,
@@ -115,7 +109,7 @@ function Cassette.overdub(ctx::JacobianSparsityContext,
                           Y::Tagged,
                           val::Tagged,
                           idx::Int...)
-    S = ctx.metadata
+    S = ctx.metadata[1]
     if metatype(Y, ctx) <: JacOutput
         set = metadata(val, ctx)
         if set isa ProvinanceSet
@@ -159,7 +153,7 @@ function Cassette.overdub(ctx::JacobianSparsityContext,
                           Y::Tagged,
                           ystart,
                           len)
-    S = ctx.metadata
+    S = ctx.metadata[1]
     if metatype(Y, ctx) <: JacInput && metatype(X, ctx) <: JacOutput
         # Write directly to the output sparsity
         val = Cassette.fallback(ctx, f, X, xstart, Y, ystart, len)

src/propagate_tags.jl

@@ -1,11 +1,16 @@
 
+@inline anytagged() = false
+@inline anytagged(x::Tagged, args...) = true
+@inline anytagged(x, args...) = anytagged(args...)
+
+
 macro proptagcontext(name)
     quote
         Cassette.@context($name)
 
         function Cassette.overdub(ctx::$name, f, args...)
             # this check can be inferred (in theory)
-            if any(x->x isa Tagged, args)
+            if anytagged(args...)
                 # This is a slower check
                 if !any(x->!(metatype(x, ctx) <: Cassette.NoMetaData), args)
                     return Cassette.recurse(ctx, f, args...)

test/common.jl

@@ -10,15 +10,15 @@ function jactester(f, Y, X, args...)
     ctx, val = jacobian_sparsity(f, Y, X, args...; raw=true)
 end
 
-jactestmeta(args...) = jactester(args...)[1].metadata
+jactestmeta(args...) = jactester(args...)[1].metadata[1]
 jactestval(args...) = jactester(args...) |> ((ctx,val),) -> untag(val, ctx)
 jactesttag(args...) = jactester(args...) |> ((ctx,val),) -> metadata(val, ctx)
 
 function hesstester(f, X, args...)
     ctx, val = hessian_sparsity(f, X, args...; raw=true)
 end
 
-hesstestmeta(args...) = hesstester(args...)[1].metadata
+hesstestmeta(args...) = hesstester(args...)[1].metadata[1]
 hesstestval(args...)  = hesstester(args...) |> ((ctx,val),) -> untag(val, ctx)
 hesstesttag(args...)  = hesstester(args...) |> ((ctx,val),) -> metadata(val, ctx)
 

test/examples.jl

@@ -0,0 +1,16 @@
+fcalls = 0
+function f(dx,x)
+  global fcalls += 1
+  for i in 2:length(x)-1
+    dx[i] = x[i-1] - 2x[i] + x[i+1]
+  end
+  dx[1] = -2x[1] + x[2]
+  dx[end] = x[end-1] - 2x[end]
+  nothing
+end
+
+using SparsityDetection, SparseArrays
+input = rand(10)
+output = similar(input)
+sparsity_pattern = sparsity!(f,output,input)
+jac = Float64.(sparse(sparsity_pattern))

test/jacobian.jl

@@ -18,11 +18,15 @@ let
     g(y,x) = y[:] .= x .+ 1
     #g(y,x) = y .= x .+ 1 -- memove
 
-    @test sparse(jactestmeta(g, [1], [2])[1]) == sparse([1], [1], true)
+    println("Broadcast timings")
+    println("  y .= x")
     # test path of unsafe_copy from Input to Output
-    @test sparsity!((y,x) -> y .= x, [1,2,3], [1,2,3]) == sparse([1,2,3], [1,2,3], true)
+    @test @time jacobian_sparsity((y,x) -> y .= x, [1,2,3], [1,2,3]) == sparse([1,2,3], [1,2,3], true)
+    println("  y[:] .= x .+ 1")
+    @test @time sparse(jactestmeta(g, [1], [2])) == sparse([1], [1], true)
+    println("  y[1:2] .= x[2:3]")
     # test path of unsafe_copy from Input to an intermediary
-    @test sparsity!((y,x) -> y[1:2] .= x[2:3], [1,2,3], [1,2,3]) == sparse([1,2],[2,3],true, 3,3)
+    @test @time jacobian_sparsity((y,x) -> y[1:2] .= x[2:3], [1,2,3], [1,2,3]) == sparse([1,2],[2,3],true, 3,3)
 
     using LinearAlgebra, SparsityDetection
 
@@ -31,7 +35,7 @@ let
         mul!(out, A, x)
     end
     x = [1:4;]; out = similar(x);
-    @test sparsity!(testsparse!, out, x) == sparse([1,2,1,2,3,2,3,4,3,4],
+    @test jacobian_sparsity(testsparse!, out, x) == sparse([1,2,1,2,3,2,3,4,3,4],
                                                    [1,1,2,2,2,3,3,3,4,4], true)
 end
 
@@ -44,5 +48,5 @@ end
 
     x = [1.0:10;]
     out = similar(x)
-    @test all(sparsity!(f, out, x) .== 1)
+    @test all(jacobian_sparsity(f, out, x) .== 1)
 end

test/runtests.jl

@@ -5,3 +5,4 @@ include("common.jl")
 @testset "Paraboloid example" begin include("paraboloid.jl") end
 
 @testset "Exploration" begin include("ifsandbuts.jl") end
+@testset "Examples" begin include("examples.jl") end