JuliaDiff / SparseDiffTools.jl

Fast jacobian computation through sparsity exploitation and matrix coloring
MIT License
237 stars 41 forks source link

Sparsity detection where input access pattern depends on Array data does not work #136

Closed moyner closed 3 years ago

moyner commented 3 years ago

Hi,

I'm currently testing some of the Julia AD features on some very large sparse problems and I wanted to try the sparsity detection/coloring. I'm not entirely sure if this issue should go into SciML/SparsityDetection.jl instead, but that module is deprecated according to the README. This module depends on the deprecated module, but seems to be active? I see that ModelingToolkit.jl does sparsity detection, but I would prefer to make my simple examples work without introducing an entire modeling framework.

I have done variants of the "hello world" Poisson-FDM type example and that works fine with excellent performance. I run into errors when I have a custom function where the access pattern is determined by the entries of an Array instead of hard coded in the function body. I know that the sparsity pattern would change if the Array changes, but it would be not change for a given sparsity context. The array is wrapped in a closure, so I did what I could to tell the compiler that the pattern doesn't change based on the inputs exposed to jacobian_sparsity. The code runs through regular (dense) ForwardDiff without issue.

Here is a short MWE that demonstrates the issue:

using ForwardDiff
using SparsityDetection, SparseArrays
using SparseDiffTools
using Test

function testEq(p::Vector, N::Vector, useIndex::Bool)
    eq = similar(p)
    testEq!(eq, p, N, useIndex)
    return eq
end

function testEq!(eq::Vector, p::Vector, N::Vector, useIndex::Bool)
    # Function that accesses two elements for each entry
    for i = 1:2
        r = N[2];
        l = N[1];
        if useIndex
            eq[i] = p[r] - p[l] # This does not work for sparsity_pattern
        else
            eq[i] = p[2] - p[1] # This works for sparsity_pattern
        end
    end
end
# Generate some data
eq = rand(2)
p = rand(2)
I = [1, 2]

# Out-of-place
F(p) = testEq(p, I, false)
# In-place
F!(eq, p) = testEq!(eq, p, I, false)

# Out-of-place
G(p) = testEq(p, I, true)
# In-place
G!(eq, p) = testEq!(eq, p, I, true)

# Test both versions
@test all(F(p) .== G(p))

# Output the dense Jacobian to verify ForwarDiff compatibility
display(eq)
dEq = ForwardDiff.jacobian(F, p);
display(dEq)

# Works just fine
sparsity_pattern = jacobian_sparsity(F!, eq, p)
# ...but this doesn't, even when closures gives same access pattern in p?
sparsity_pattern_2 = jacobian_sparsity(G!, eq, p)

I have attached the full output + stack trace on a fresh Julia 1.6 install on Windows 10, run in a fresh REPL instance under VSCode. The error is a Casette error so it is pretty long. The key bit is inlined here:

Explored path: SparsityDetection.Path(Bool[], 1)
ERROR: ArgumentError: invalid index: Tagged(Tag{nametype(JacobianSparsityContext),3569839777,Nothing}(), 2, _) of type Cassette.Tagged{Cassette.Tag{nametype(JacobianSparsityContext), 0x00000000d4c76ea1, Nothing}, Int64, SparsityDetection.ProvinanceSet, Cassette.NoMetaMeta, Cassette.Context{nametype(JacobianSparsityContext), Tuple{Sparsity, SparsityDetection.Path}, Cassette.Tag{nametype(JacobianSparsityContext), 0x00000000d4c76ea1, Nothing}, SparsityDetection.var"##PassType#257", IdDict{Module, Dict{Symbol, Cassette.BindingMeta}}, Cassette.DisableHooks}}
Stacktrace:
  [1] to_index(i::Cassette.Tagged{Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, Int64, SparsityDetection.ProvinanceSet, Cassette.NoMetaMeta, Cassette.Context{SparsityDetection.var"##JacobianSparsityContext#Name", Tuple{Sparsity, SparsityDetection.Path}, Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, SparsityDetection.var"##PassType#257", IdDict{Module, Dict{Symbol, Cassette.BindingMeta}}, Cassette.DisableHooks}})
    @ Base .\indices.jl:300
  [2] to_index(A::LinearIndices{1, Tuple{Base.OneTo{Int64}}}, i::Cassette.Tagged{Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, Int64, SparsityDetection.ProvinanceSet, Cassette.NoMetaMeta, Cassette.Context{SparsityDetection.var"##JacobianSparsityContext#Name", Tuple{Sparsity, SparsityDetection.Path}, Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, SparsityDetection.var"##PassType#257", IdDict{Module, Dict{Symbol, Cassette.BindingMeta}}, Cassette.DisableHooks}})
    @ Base .\indices.jl:277
  [3] to_indices
    @ .\indices.jl:333 [inlined]
  [4] to_indices
    @ .\indices.jl:325 [inlined]
  [5] getindex
    @ .\abstractarray.jl:1170 [inlined]
  [6] overdub(ctx::Cassette.Context{SparsityDetection.var"##JacobianSparsityContext#Name", Tuple{Sparsity, SparsityDetection.Path}, Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, SparsityDetection.var"##PassType#257", IdDict{Module, Dict{Symbol, Cassette.BindingMeta}}, Cassette.DisableHooks}, f::typeof(getindex), X::Cassette.Tagged{Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, Vector{Float64}, Union{SparsityDetection.JacInput, SparsityDetection.JacOutput, SparsityDetection.ProvinanceSet}, Vector{Cassette.Meta{SparsityDetection.ProvinanceSet, Cassette.NoMetaMeta}}, Cassette.Context{SparsityDetection.var"##JacobianSparsityContext#Name", Sparsity, Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, Cassette.var"##PassType#259", IdDict{Module, Dict{Symbol, Cassette.BindingMeta}}, Cassette.DisableHooks}}, idx::Cassette.Tagged{Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, Int64, SparsityDetection.ProvinanceSet, Cassette.NoMetaMeta, Cassette.Context{SparsityDetection.var"##JacobianSparsityContext#Name", Tuple{Sparsity, SparsityDetection.Path}, Cassette.Tag{SparsityDetection.var"##JacobianSparsityContext#Name", 0x00000000d4c76ea1, Nothing}, SparsityDetection.var"##PassType#257", IdDict{Module, Dict{Symbol, Cassette.BindingMeta}}, Cassette.DisableHooks}})
    @ SparsityDetection ~\.julia\packages\SparsityDetection\E7o7R\src\jacobian.jl:98
  [7] recurse
    @ ~\git\testrepo\test_sparsity.jl:18 [inlined]

The two identical specializations of the functions only differ in the way the indices 1 and 2 are computed. I can manually enter the sparsity pattern and do coloring etc which works perfectly, but any tips to get this working would be greatly appreciated. Is there some overview of what works and doesn't work for the sparsity detection?

stacktrace.txt

ChrisRackauckas commented 3 years ago

This has nothing to do with SparseDiffTools.jl so this issue is going to get closed. SparseDiffTools.jl assumes you know the sparsity pattern. The other packages, aren't this package and are in a different org so I can't even transfer the issue hence the close.

SparsityDetection.jl was deprecated because when Symbolics.jl gets a tracer it will be powerful enough to capture any function that this analysis can work on. It's built on Cassette.jl which is getting replaced with new tools, so keeping it alive is something that's a bit of work with a time limit. Instead, SymbolicTracing.jl will get functional in the meantime and will be a an easy way to build a full replacement.

I see that ModelingToolkit.jl does sparsity detection, but I would prefer to make my simple examples work without introducing an entire modeling framework.

You don't need to use a modeling framework to use it, so this comment doesn't make much sense. Why not just apply it to Julia code, just like SparsityDetection.jl?

I have done variants of the "hello world" Poisson-FDM type example and that works fine with excellent performance. I run into errors when I have a custom function where the access pattern is determined by the entries of an Array instead of hard coded in the function body. I know that the sparsity pattern would change if the Array changes, but it would be not change for a given sparsity context. The array is wrapped in a closure, so I did what I could to tell the compiler that the pattern doesn't change based on the inputs exposed to jacobian_sparsity. The code runs through regular (dense) ForwardDiff without issue.

This is something that needs runtime information to be transported to compile time information, which is what symbolic tracers like ModelingToolkit/Symbolics.jl can handle which a purely Cassette-based tool cannot.

function testEq!(eq::Vector, p::Vector, N::Vector, useIndex::Bool)
    # Function that accesses two elements for each entry
    for i = 1:2
        r = N[2];
        l = N[1];
        eq[i] = IfElse.ifelse(useIndex,p[r] - p[l],p[2] - p[1])
    end
end

should be perfectly fine for Symbolics.jl with both bool choices, with if statements being able to be used directly after the symbolic tracer is completed.

Hopefully that makes more sense.

moyner commented 3 years ago

This has nothing to do with SparseDiffTools.jl so this issue is going to get closed. SparseDiffTools.jl assumes you know the sparsity pattern. The other packages, aren't this package and are in a different org so I can't even transfer the issue hence the close.

SparsityDetection.jl was deprecated because when Symbolics.jl gets a tracer it will be powerful enough to capture any function that this analysis can work on. It's built on Cassette.jl which is getting replaced with new tools, so keeping it alive is something that's a bit of work with a time limit. Instead, SymbolicTracing.jl will get functional in the meantime and will be a an easy way to build a full replacement

Thank you for the reply! That definitely makes sense. I could raise the issue over at SparsityDetection, but maybe that isn't especially useful for a deprecated package.

There isn't any documentation in SymbolicTracing.jl so I am I understanding you correctly that the new tracer is ongoing/future work? I already know my sparsity pattern so it isn't a big issue, just trying to get my bearings the packages to see what's state of the art and what's not actively developed.

You don't need to use a modeling framework to use it, so this comment doesn't make much sense. Why not just apply it to Julia code, just like SparsityDetection.jl?

Then the question is how to apply it. I have gone over the documentation for ModelingToolkit.jl and the examples I found for sparsity detection involved setting up e.g. ODEFunction and performing modelingtoolkitize and symbolic reduction on it before performing an ODE solve. It is not at all clear to me how this relates to the much more constrained problem of tracing the access pattern of a simple function, even after grep-ing through the source for sparse keywords. The README describes the package as a modeling framework and a acausal modeling language. This, together with feature comparisons to SIMULINK made me hesitant to include it in my simple assessment since my time budget is pretty limited.

A small related question on the use of SparseDiffTools: I have a simple expression that results in a 3,329,000 by 1,122,000 Jacobian, where every row only contains -1 and 1 in two positions in a pre-determined sparsity structure (the function itself is in my initial post). Naturally the system has two colors and with precomputed ForwardColorJacCache it takes about ~100 ms to update the Jacobian when the function itself takes ~5 ms. I know CPUs vary etc, but does this sound reasonable? The function is type stable and no allocations/compilation occurs in the call I'm @btime-ing. I'm a bit surprised at the big ratio between the function call and the Dual call. If it sounds unreasonable I can set up a MWE and file an issue.

ChrisRackauckas commented 3 years ago

There are some advantages to SparsityDetection that SymbolicTracing won't be able to have. You don't need to do the full tracing to do detection, so I am curious what @shashi thinks of the performance on that. The reason for the deprecation was because Cassette didn't update for a bit, but @simeonschaub ended up handling Cassette so IIRC SparsityDetection.jl is fine on v1.6. I think a real AbstractInterpreter version is probably necessary and can be pretty powerful. That might be incentive to un-deprecate it, but I think one of those two would really have to take it on.

Then the question is how to apply it. I have gone over the documentation for ModelingToolkit.jl and the examples I found for sparsity detection involved setting up e.g. ODEFunction and performing modelingtoolkitize and symbolic reduction on it before performing an ODE solve. It is not at all clear to me how this relates to the much more constrained problem of tracing the access pattern of a simple function, even after grep-ing through the source for sparse keywords. The README describes the package as a modeling framework and a acausal modeling language. This, together with feature comparisons to SIMULINK made me hesitant to include it in my simple assessment since my time budget is pretty limited.

Thanks for the feedback.

A small related question on the use of SparseDiffTools: I have a simple expression that results in a 3,329,000 by 1,122,000 Jacobian, where every row only contains -1 and 1 in two positions in a pre-determined sparsity structure (the function itself is in my initial post). Naturally the system has two colors and with precomputed ForwardColorJacCache it takes about ~100 ms to update the Jacobian when the function itself takes ~5 ms. I know CPUs vary etc, but does this sound reasonable? The function is type stable and no allocations/compilation occurs in the call I'm @btime-ing. I'm a bit surprised at the big ratio between the function call and the Dual call. If it sounds unreasonable I can set up a MWE and file an issue.

That doesn't sound reasonable. Open an issue.

moyner commented 3 years ago

There are some advantages to SparsityDetection that SymbolicTracing won't be able to have. You don't need to do the full tracing to do detection, so I am curious what @shashi thinks of the performance on that. The reason for the deprecation was because Cassette didn't update for a bit, but @simeonschaub ended up handling Cassette so IIRC SparsityDetection.jl is fine on v1.6. I think a real AbstractInterpreter version is probably necessary and can be pretty powerful. That might be incentive to un-deprecate it, but I think one of those two would really have to take it on.

Much appreciated. If anyone wants me to open the issue over in that package or need anything else, please let me know.

That doesn't sound reasonable. Open an issue.

Thanks, I have posted a MWE in issue #138!

shashi commented 3 years ago

For this to be possible to trace you need to call jacobian_sparsity with a fixed N, i.e. it's only possible if your "indexer" array N is always the same. Or you are ok with generating the pattern every time it changes.

So

jacobian_sparsity(testEq!, eq, p, SparsityDetection.Fixed(N), useIndex)

might work.

moyner commented 3 years ago

For this to be possible to trace you need to call jacobian_sparsity with a fixed N, i.e. it's only possible if your "indexer" array N is always the same. Or you are ok with generating the pattern every time it changes.

My N is a neighborhood definition that determines the sparsity, so it is static. Thanks for the tip, I'll test that!

ChrisRackauckas commented 3 years ago

Talking with @shashi it seems SymbolicTracing.jl -> jacobian_sparsity would be faster, considering the compile time required by the abstract interpretation tooling, and so that will be how it progresses.