Skip to content

Commit

Permalink
add a hack to avoid inference recursion limit to allow more optimization
Browse files Browse the repository at this point in the history
For this target:
```julia
using JSONBase, BenchmarkTools
struct A
    a::Int
    b::Int
    c::Int
    d::Int
end
@benchmark JSONBase.materialize("""{ "a": 1, "b": 2, "c": 3, "d": 4}""", A)
```

Before:
```julia
BenchmarkTools.Trial: 10000 samples with 331 evaluations.
 Range (min … max):  258.937 ns …   4.373 μs  ┊ GC (min … max): 0.00% … 91.32%
 Time  (median):     279.202 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   291.054 ns ± 185.340 ns  ┊ GC (mean ± σ):  3.24% ±  4.74%

           ▃█▃    ▂
  ▁▂▂▂▁▁▁▁▃███▃▂▂▄█▇▅▅▅▄▂▂▂▃▃▄▃▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  259 ns           Histogram: frequency by time          334 ns <

 Memory estimate: 384 bytes, allocs estimate: 10.
```

After:
```julia
BenchmarkTools.Trial: 10000 samples with 750 evaluations.
 Range (min … max):  173.333 ns …  1.618 μs  ┊ GC (min … max): 0.00% … 86.43%
 Time  (median):     179.389 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   184.518 ns ± 58.847 ns  ┊ GC (mean ± σ):  1.40% ±  3.86%

          ▂▄▆██▇▅▄▁▁ ▂▃▄▅▄▄▄▄▂▁▁▁▂▂▂▂▂▂▂▂▂▂▁                   ▂
  ▂▅▄▅▅▃▃█████████████████████████████████████████████▇▇▅▄▅▆▅▅ █
  173 ns        Histogram: log(frequency) by time       200 ns <

 Memory estimate: 128 bytes, allocs estimate: 2.
```

closes quinnj#2
  • Loading branch information
aviatesk committed Mar 19, 2023
1 parent a56acb6 commit f004e51
Show file tree
Hide file tree
Showing 5 changed files with 97 additions and 5 deletions.
5 changes: 3 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,9 @@ Parsers = "2.5.1"
julia = "1.6"

[extras]
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[targets]
test = ["Test", "OrderedCollections"]
test = ["Test", "JET", "OrderedCollections"]
66 changes: 65 additions & 1 deletion src/JSONBase.jl
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,70 @@ include("binary.jl")
include("materialize.jl")
include("json.jl")

# HACK to avoid inference recursion limit and the de-optimization:
# This works since know the inference recursion will terminate due to the fact that this
# method is only called when materializing a struct with definite number of fields, i.e.
# that is not self-referencing, so it is guaranteed that there are no cycles in a recursive
# `materialize` call. Especially, the `fieldcount` call in the struct fallback case within
# the `materialize` should have errored for this case.
# TODO we should revisit this hack when we start to support https://github.com/quinnj/JSONBase.jl/issues/3
function validate_recursion_relation_sig(f, nargs::Int, sig)
@nospecialize f sig
sig = Base.unwrap_unionall(sig)
@assert sig isa DataType "unexpected `recursion_relation` call"
@assert sig.name === Tuple.name "unexpected `recursion_relation` call"
@assert length(sig.parameters) == nargs "unexpected `recursion_relation` call"
@assert sig.parameters[1] == typeof(f) "unexpected `recursion_relation` call"
return sig
end
@static if hasfield(Method, :recursion_relation)
let applyobject_recursion_relation = function (
method::Method, topmost::Union{Nothing,Method},
@nospecialize(sig), @nospecialize(topmostsig))
# Core.println("applyobject")
# Core.println(" method = ", method)
# Core.println(" topmost = ", topmost)
# Core.println(" sig = ", sig)
# Core.println(" topmostsig = ", topmostsig)
sig = validate_recursion_relation_sig(applyobject, 3, sig)
topmostsig = validate_recursion_relation_sig(applyobject, 3, topmostsig)
return sig.parameters[2] topmostsig.parameters[2]
end
method = only(methods(applyobject, (Any,LazyValues,)))
method.recursion_relation = applyobject_recursion_relation
end
let applyfield_recursion_relation = function (
method::Method, topmost::Union{Nothing,Method},
@nospecialize(sig), @nospecialize(topmostsig))
# Core.println("applyfield")
# Core.println(" method = ", method)
# Core.println(" topmost = ", topmost)
# Core.println(" sig = ", sig)
# Core.println(" topmostsig = ", topmostsig)
sig = validate_recursion_relation_sig(applyfield, 6, sig)
topmostsig = validate_recursion_relation_sig(applyfield, 6, topmostsig)
return sig.parameters[2] topmostsig.parameters[2]
end
method = only(methods(applyfield, (Type,Type,Any,Any,Any)))
method.recursion_relation = applyfield_recursion_relation
end
let _materialize_recursion_relation = function (
method::Method, topmost::Union{Nothing,Method},
@nospecialize(sig), @nospecialize(topmostsig))
# Core.println("_materialize")
# Core.println(" method = ", method)
# Core.println(" topmost = ", topmost)
# Core.println(" sig = ", sig)
# Core.println(" topmostsig = ", topmostsig)
sig = validate_recursion_relation_sig(_materialize, 5, sig)
topmostsig = validate_recursion_relation_sig(_materialize, 5, topmostsig)
return sig.parameters[4] topmostsig.parameters[4]
end
method = only(methods(_materialize, (Any,LazyValue,Type,Type)))
method.recursion_relation = _materialize_recursion_relation
end
end

# a helper higher-order function that converts an
# API.applyeach function that operates potentially on a
# PtrString to one that operates on a String
Expand Down Expand Up @@ -84,4 +148,4 @@ end # module
# package docs
# topretty
# allow materialize on any ObjectLike? i.e. Dicts? (would need applyobject on Dict)
# checkout JSON5, Amazon Ion?
# checkout JSON5, Amazon Ion?
2 changes: 1 addition & 1 deletion src/lazy.jl
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ _applyobject(f::F, x) where {F} = applyobject(f, x)
# we're now positioned at the start of the value
val = lazy(buf, pos, len, b, opts)
ret = keyvalfunc(key, val)
# if ret is not an Continue, then we're
# if ret is not an Continue, then we're
# short-circuiting parsing via e.g. selection syntax
# so return immediately
ret isa Continue || return ret
Expand Down
24 changes: 24 additions & 0 deletions test/optimization.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
using JSONBase, JET

struct OptimizationFailureChecker end
function JET.configured_reports(::OptimizationFailureChecker, reports::Vector{JET.InferenceErrorReport})
return filter(reports) do @nospecialize report::JET.InferenceErrorReport
isa(report, JET.OptimizationFailureReport)
end
end

# https://github.com/quinnj/JSONBase.jl/issues/2
struct RecursiveInferenceTest1
a::Int
b::Int
end
@test_opt annotate_types=true report_config=OptimizationFailureChecker() JSONBase.materialize("""{ "a": 1, "b": 2 }""", RecursiveInferenceTest1)

struct RecursiveInferenceTest2Inner
b::Int
end
struct RecursiveInferenceTest2
a::Int
b::RecursiveInferenceTest2Inner
end
@test_opt annotate_types=true report_config=OptimizationFailureChecker() JSONBase.materialize("""{ "a": 1, "b": { "b": 2 } }""", RecursiveInferenceTest2)
5 changes: 4 additions & 1 deletion test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ struct File end
["Gilbert", "2013", 24, true]
["Alexa", "2013", 29, true]
["May", "2012B", 14, false]
["Deloise", "2012A", 19, true]
["Deloise", "2012A", 19, true]
"""; jsonlines=true, float64=true) ==
[["Name", "Session", "Score", "Completed"],
["Gilbert", "2013", 24.0, true],
Expand Down Expand Up @@ -294,3 +294,6 @@ end
include("struct.jl")
include("json.jl")
include("numbers.jl")
@static if VERSION v"1.8"
@testset "Optimization test with JET" include("optimization.jl")
end

0 comments on commit f004e51

Please sign in to comment.