Merge pull request #13 from shashi/s/try-all-paths

sparsity detection: a pass to try all tainted branches

Author: ChrisRackauckas

.travis.yml

@@ -4,7 +4,7 @@ os:
   - linux
   - osx
 julia:
-  - 1.1
+  - 1.2
   - nightly
 matrix:
   allow_failures:

Project.toml

@@ -6,6 +6,7 @@ version = "0.1.0"
 [deps]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
+Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 DiffEqDiffTools = "01453d9d-ee7c-5054-8395-0335cb756afa"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"

src/SparseDiffTools.jl

@@ -22,11 +22,13 @@ export  contract_color,
         autoback_hesvec,autoback_hesvec!,
         JacVec,HesVec,HesVecGrad
 
+
 include("coloring/high_level.jl")
 include("coloring/contraction_coloring.jl")
 include("coloring/greedy_d1_coloring.jl")
 include("coloring/matrix2graph.jl")
 include("differentiation/compute_jacobian_ad.jl")
 include("differentiation/jaches_products.jl")
+include("program_sparsity/program_sparsity.jl")
 
 end # module

src/program_sparsity/path.jl

@@ -0,0 +1,107 @@
+#### 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 this_here_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::SparsityContext) = alldone(c.metadata[2])
+reset!(c::SparsityContext) = reset!(c.metadata[2])
+this_here_predicate!(c::SparsityContext) = this_here_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> this_here_predicate!(p)
+false
+
+julia> alldone(p) # must be called at the end of a full run
+false
+
+julia> this_here_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> this_here_predicate!(p)
+true
+
+julia> this_here_predicate!(p)
+false
+
+julia> alldone(p)
+false
+
+julia> reset!(p)
+
+julia> p
+Path(Bool[false, true], 1)
+
+julia> this_here_predicate!(p)
+false
+
+julia> this_here_predicate!(p)
+true
+
+julia> reset!(p)
+
+julia> this_here_predicate!(p)
+true
+
+julia> this_here_predicate!(p)
+true
+
+julia> alldone(p)
+true
+=#

src/program_sparsity/program_sparsity.jl

@@ -0,0 +1,44 @@
+include("sparsity_tracker.jl")
+include("path.jl")
+include("take_all_branches.jl")
+
+struct Fixed
+    value
+end
+
+"""
+`sparsity!(f, Y, X, args...; sparsity=Sparsity(length(X), length(Y)))`
+
+Execute the program that figures out the sparsity pattern of
+the jacobian of the function `f`.
+
+# Arguments:
+- `f`: the function
+- `Y`: the output array
+- `X`: the input array
+- `args`: trailing arguments to `f`. They are considered subject to change, unless wrapped as `Fixed(arg)`
+- `S`: (optional) the sparsity pattern
+
+Returns a `Sparsity`
+"""
+function sparsity!(f!, Y, X, args...; sparsity=Sparsity(length(Y), length(X)))
+    path = Path()
+    ctx = SparsityContext(metadata=(sparsity, path), pass=BranchesPass)
+    ctx = Cassette.enabletagging(ctx, f!)
+    ctx = Cassette.disablehooks(ctx)
+
+    while true
+        Cassette.recurse(ctx,
+                         f!,
+                         tag(Y, ctx, Output()),
+                         tag(X, ctx, Input()),
+                         # TODO: make this recursive
+                         map(arg -> arg isa Fixed ?
+                             arg.value : tag(arg, ctx, ProvinanceSet(())), args)...)
+
+        println("Explored path: ", path)
+        alldone(path) && break
+        reset!(path)
+    end
+    sparsity
+end

src/program_sparsity/sparsity_tracker.jl

@@ -0,0 +1,188 @@
+using Cassette
+import Cassette: tag, untag, Tagged, metadata, hasmetadata, istagged, canrecurse
+import Core: SSAValue
+using SparseArrays
+
+export Sparsity, sparsity!
+
+"""
+The sparsity pattern.
+
+- `I`: Input index
+- `J`: Ouput index
+
+`(i, j)` means the `j`th element of the output depends on
+the `i`th element of the input. Therefore `length(I) == length(J)`
+"""
+struct Sparsity
+    m::Int
+    n::Int
+    I::Vector{Int} # Input
+    J::Vector{Int} # Output
+end
+
+SparseArrays.sparse(s::Sparsity) = sparse(s.I, s.J, true, s.m, s.n)
+
+Sparsity(m, n) = Sparsity(m, n, Int[], Int[])
+
+function Base.push!(S::Sparsity, i::Int, j::Int)
+    push!(S.I, i)
+    push!(S.J, j)
+end
+
+# Tags:
+struct Input end
+struct Output end
+
+struct ProvinanceSet{T}
+    set::T # Set, Array, Int, Tuple, anything!
+end
+
+# note: this is not strictly set union, just some efficient way of concating
+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,
+           q::ProvinanceSet,
+           rs::ProvinanceSet...) = union(union(p, q), rs...)
+Base.union(p::ProvinanceSet) = p
+
+function Base.push!(S::Sparsity, i::Int, js::ProvinanceSet)
+    for j in js.set
+        push!(S, i, j)
+    end
+end
+
+Cassette.@context SparsityContext
+
+const TagType = Union{Input, Output, ProvinanceSet}
+Cassette.metadatatype(::Type{<:SparsityContext}, ::DataType) = TagType
+
+metatype(x, ctx) = hasmetadata(x, ctx) && istagged(x, ctx) && typeof(metadata(x, ctx))
+function ismetatype(x, ctx, T)
+    hasmetadata(x, ctx) && istagged(x, ctx) && (metadata(x, ctx) isa T)
+end
+
+# Dummy type when you getindex
+struct Tainted end
+
+# getindex on the input
+function Cassette.overdub(ctx::SparsityContext,
+                          f::typeof(getindex),
+                          X::Tagged,
+                          idx::Int...)
+    if ismetatype(X, ctx, Input)
+        i = LinearIndices(untag(X, ctx))[idx...]
+        val = Cassette.fallback(ctx, f, X, idx...)
+        tag(val, ctx, ProvinanceSet(i))
+    else
+        Cassette.recurse(ctx, f, X, idx...)
+    end
+end
+
+# setindex! on the output
+function Cassette.overdub(ctx::SparsityContext,
+                          f::typeof(setindex!),
+                          Y::Tagged,
+                          val::Tagged,
+                          idx::Int...)
+    S, path = ctx.metadata
+    if ismetatype(Y, ctx, Output)
+        set = metadata(val, ctx)
+        if set isa ProvinanceSet
+            i = LinearIndices(untag(Y, ctx))[idx...]
+            push!(S, i, set)
+        end
+        Cassette.fallback(ctx, f, Y, val, idx...)
+    else
+        Cassette.recurse(ctx, f, Y, val, idx...)
+    end
+end
+
+function get_provinance(ctx, arg::Tagged)
+    if metadata(arg, ctx) isa ProvinanceSet
+        metadata(arg, ctx)
+    else
+        ProvinanceSet(())
+    end
+end
+
+get_provinance(ctx, arg) = ProvinanceSet(())
+
+# Any function acting on a value tagged with ProvinanceSet
+function _overdub_union_provinance(::Val{eval}, ctx::SparsityContext, f, args...) where {eval}
+    idxs = findall(x->ismetatype(x, ctx, ProvinanceSet), args)
+    if isempty(idxs)
+        Cassette.fallback(ctx, f, args...)
+    else
+        provinance = union(map(arg->get_provinance(ctx, arg), args[idxs])...)
+        if eval
+            val = Cassette.fallback(ctx, f, args...)
+            tag(val, ctx, provinance)
+        else
+            tag(Tainted(), ctx, provinance)
+        end
+    end
+end
+
+function Cassette.overdub(ctx::SparsityContext, f, args...)
+    haspsets = any(x->ismetatype(x, ctx, ProvinanceSet), args)
+    hasinput = any(x->ismetatype(x, ctx, Input), args)
+    if haspsets && !hasinput # && !canrecurse(ctx, f, args...)
+        _overdub_union_provinance(Val{true}(), ctx, f, args...)
+    else
+        Cassette.recurse(ctx, f, args...)
+    end
+end
+
+#=
+# Examples:
+#
+using UnicodePlots
+
+sspy(s::Sparsity) = spy(sparse(s))
+
+julia> sparsity!([0,0,0], [23,53,83]) do Y, X
+           Y[:] .= X
+           Y == X
+       end
+(true, Sparsity([1, 2, 3], [1, 2, 3]))
+
+julia> sparsity!([0,0,0], [23,53,83]) do Y, X
+           for i=1:3
+               for j=i:3
+                   Y[j] += X[i]
+               end
+           end; Y
+       end
+([23, 76, 159], Sparsity(3, 3, [1, 2, 3, 2, 3, 3], [1, 1, 1, 2, 2, 3]))
+
+julia> sspy(ans[2])
+     Sparsity Pattern
+     ┌─────┐
+   1 │⠀⠄⠀⠀⠀│ > 0
+   3 │⠀⠅⠨⠠⠀│ < 0
+     └─────┘
+     1     3
+     nz = 6
+
+julia> sparsity!(f, zeros(Int, 3,3), [23,53,83])
+([23, 53, 83], Sparsity(9, 3, [2, 5, 8], [1, 2, 3]))
+
+julia> sspy(ans[2])
+     Sparsity Pattern
+     ┌─────┐
+   1 │⠀⠄⠀⠀⠀│ > 0
+     │⠀⠀⠠⠀⠀│ < 0
+   9 │⠀⠀⠀⠐⠀│
+     └─────┘
+     1     3
+     nz = 3
+=#