Merge pull request #11 from SymbolicML/snoopcompile

MilesCranmer · web-flow · commit 34271cb68c72 · 2022-11-27T10:04:15.000-05:00
Improve precompilation with SnoopCompile.jl
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -27,27 +27,25 @@ jobs:
       fail-fast: false
       matrix:
         julia-version:
-          - '1.8.2'
+          - '1.8.3'
         os:
           - ubuntu-latest
+          - windows-latest
+          - macOS-latest
+        include:
+          - os: ubuntu-latest
+            julia-version: '1.7.2'
+          - os: ubuntu-latest
+            julia-version: '1.6.7'
 
     steps:
-      - uses: actions/checkout@v1.0.0
+      - uses: actions/checkout@v2
       - name: "Set up Julia"
-        uses: julia-actions/setup-julia@v1.6.0
+        uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
-      - name: Cache dependencies
-        uses: actions/cache@v1 # Thanks FromFile.jl
-        env:
-          cache-name: cache-artifacts
-        with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-build-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-build-${{ env.cache-name }}-
-            ${{ runner.os }}-build-
-            ${{ runner.os }}-
+      - name: "Cache artifacts"
+        uses: julia-actions/cache@v1
       - name: "Build package"
         uses: julia-actions/julia-buildpkg@v1
       - name: "Run tests"
diff --git a/Project.toml b/Project.toml
@@ -1,21 +1,23 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 LoopVectorization = "0.12"
 Reexport = "1"
+SnoopPrecompile = "1"
 SymbolicUtils = "0.19"
 Zygote = "0.6"
 julia = "1.6"
diff --git a/src/DynamicExpressions.jl b/src/DynamicExpressions.jl
@@ -43,4 +43,7 @@ macro ignore(args...) end
 # To get LanguageServer to register library within tests
 @ignore include("../test/runtests.jl")
 
+include("precompile.jl")
+do_precompilation(; mode=:precompile)
+
 end
diff --git a/src/OperatorEnumConstruction.jl b/src/OperatorEnumConstruction.jl
@@ -221,51 +221,20 @@ function OperatorEnum(;
     binary_operators = Function[op for op in binary_operators]
     unary_operators = Function[op for op in unary_operators]
 
-    if enable_autodiff
-        diff_binary_operators = Function[]
-        diff_unary_operators = Function[]
+    diff_binary_operators = Function[]
+    diff_unary_operators = Function[]
 
-        test_inputs = Float32.(LinRange(-100, 100, 99))
-        # Create grid over [-100, 100]^2:
-        test_inputs_xy = Array{Float32}(undef, 2, 99^2)
-        row = 1
-        for x in test_inputs, y in test_inputs
-            test_inputs_xy[:, row] .= [x, y]
-            row += 1
-        end
+    if enable_autodiff
         for op in binary_operators
             diff_op(x, y) = gradient(op, x, y)
-
-            test_output = diff_op.(test_inputs_xy[1, :], test_inputs_xy[2, :])
-            gradient_exists = all((x) -> x !== nothing, Iterators.flatten(test_output))
-            if gradient_exists
-                push!(diff_binary_operators, diff_op)
-            else
-                @warn "Automatic differentiation has been turned off, since operator $(op) does not have well-defined gradients."
-                enable_autodiff = false
-                break
-            end
+            push!(diff_binary_operators, diff_op)
         end
-
         for op in unary_operators
             diff_op(x) = gradient(op, x)[1]
-            test_output = diff_op.(test_inputs)
-            gradient_exists = all((x) -> x !== nothing, test_output)
-            if gradient_exists
-                push!(diff_unary_operators, diff_op)
-            else
-                @warn "Automatic differentiation has been turned off, since operator $(op) does not have well-defined gradients."
-                enable_autodiff = false
-                break
-            end
+            push!(diff_unary_operators, diff_op)
         end
     end
 
-    if !enable_autodiff
-        diff_binary_operators = Function[]
-        diff_unary_operators = Function[]
-    end
-
     operators = OperatorEnum(
         binary_operators, unary_operators, diff_binary_operators, diff_unary_operators
     )
diff --git a/src/precompile.jl b/src/precompile.jl
@@ -0,0 +1,202 @@
+import SnoopPrecompile: @precompile_all_calls, @precompile_setup
+
+macro ignore_domain_error(ex)
+    return esc(
+        quote
+            try
+                $ex
+            catch e
+                if !(e isa DomainError)
+                    rethrow(e)
+                end
+            end
+        end,
+    )
+end
+
+"""
+    test_all_combinations(; binary_operators, unary_operators, turbo, types)
+
+Test all combinations of the given operators and types. Useful for precompilation.
+"""
+function test_all_combinations(; binary_operators, unary_operators, turbo, types)
+    for binops in binary_operators,
+        unaops in unary_operators,
+        use_turbo in turbo,
+        T in types
+
+        length(binops) == 0 && length(unaops) == 0 && continue
+        T == Float16 && use_turbo && continue
+
+        X = rand(T, 3, 10)
+        operators = OperatorEnum(;
+            binary_operators=binops,
+            unary_operators=unaops,
+            define_helper_functions=false,
+            enable_autodiff=true,
+        )
+        x = Node(T; feature=1)
+        c = Node(T; val=one(T))
+
+        # Trivial:
+        for l in (x, c)
+            @ignore_domain_error eval_tree_array(l, X, operators; turbo=use_turbo)
+            for variable in (true, false)
+                @ignore_domain_error eval_grad_tree_array(
+                    l, X, operators; turbo=use_turbo, variable
+                )
+            end
+        end
+
+        # Binary operators
+        for i in eachindex(binops), l in (x, c), r in (x, c)
+            tree = Node(i, l, r)
+            tree = convert(Node{T}, tree)
+            @ignore_domain_error eval_tree_array(tree, X, operators; turbo=use_turbo)
+            for variable in (true, false)
+                @ignore_domain_error eval_grad_tree_array(
+                    l, X, operators; turbo=use_turbo, variable
+                )
+            end
+        end
+
+        # Unary operators
+        for j in eachindex(unaops), k in eachindex(unaops), l in (x, c)
+            tree = Node(j, l)
+            tree = convert(Node{T}, tree)
+            @ignore_domain_error eval_tree_array(tree, X, operators; turbo=use_turbo)
+            for variable in (true, false)
+                @ignore_domain_error eval_grad_tree_array(
+                    l, X, operators; turbo=use_turbo, variable
+                )
+            end
+
+            tree = Node(j, Node(k, l))
+            tree = convert(Node{T}, tree)
+            @ignore_domain_error eval_tree_array(tree, X, operators; turbo=use_turbo)
+            for variable in (true, false)
+                @ignore_domain_error eval_grad_tree_array(
+                    l, X, operators; turbo=use_turbo, variable
+                )
+            end
+        end
+
+        # Both operators
+        for i in eachindex(binary_operators),
+            j1 in eachindex(unary_operators),
+            j2 in eachindex(unary_operators),
+            l in (x, c),
+            r in (x, c)
+
+            tree = Node(i, Node(j1, l), Node(j2, r))
+            tree = convert(Node{T}, tree)
+            @ignore_domain_error eval_tree_array(tree, X, operators; turbo=use_turbo)
+            for variable in (true, false)
+                @ignore_domain_error eval_grad_tree_array(
+                    l, X, operators; turbo=use_turbo, variable
+                )
+            end
+
+            tree = Node(j1, Node(i, l, r))
+            tree = convert(Node{T}, tree)
+            @ignore_domain_error eval_tree_array(tree, X, operators; turbo=use_turbo)
+            for variable in (true, false)
+                @ignore_domain_error eval_grad_tree_array(
+                    l, X, operators; turbo=use_turbo, variable
+                )
+            end
+        end
+    end
+    return nothing
+end
+
+function test_functions_on_trees(::Type{T}, operators) where {T}
+    local x, c, tree
+    for T1 in [Float16, Float32, Float64]
+        x = Node(T1; feature=1)
+        c = Node(T1; val=T1(1.0))
+        tree = Node(
+            2,
+            Node(1, Node(1, Node(2, x, c), Node(3, c, Node(1, x)))),
+            Node(3, Node(1, Node(4, x, x))),
+        )
+    end
+    tree = convert(Node{T}, tree)
+    for preserve_topology in [true, false]
+        tree = copy_node(tree; preserve_topology)
+        set_node!(tree, copy_node(tree; preserve_topology))
+    end
+
+    string_tree(tree, operators)
+    count_nodes(tree)
+    count_constants(tree)
+    count_depth(tree)
+    index_constants(tree)
+    has_operators(tree)
+    has_constants(tree)
+    get_constants(tree)
+    set_constants(tree, get_constants(tree))
+    combine_operators(tree, operators)
+    simplify_tree(tree, operators)
+    return nothing
+end
+
+macro maybe_precompile_setup(mode, ex)
+    precompile_ex = Expr(
+        :macrocall, Symbol("@precompile_setup"), LineNumberNode(@__LINE__), ex
+    )
+    return quote
+        if $(esc(mode)) == :compile
+            $(esc(ex))
+        elseif $(esc(mode)) == :precompile
+            $(esc(precompile_ex))
+        else
+            error("Invalid value for mode: " * show($(esc(mode))))
+        end
+    end
+end
+
+macro maybe_precompile_all_calls(mode, ex)
+    precompile_ex = Expr(
+        :macrocall, Symbol("@precompile_all_calls"), LineNumberNode(@__LINE__), ex
+    )
+    return quote
+        if $(esc(mode)) == :compile
+            $(esc(ex))
+        elseif $(esc(mode)) == :precompile
+            $(esc(precompile_ex))
+        else
+            error("Invalid value for mode: " * show($(esc(mode))))
+        end
+    end
+end
+
+"""`mode=:precompile` will use `@precompile_*` directives; `mode=:compile` runs."""
+function do_precompilation(; mode=:precompile)
+    @maybe_precompile_setup mode begin
+        binary_operators = [[+, -, *, /, ^]]
+        unary_operators = [[sin, cos, exp, log, sqrt, abs]]
+        turbo = [true, false]
+        types = [Float32, Float64]
+        @maybe_precompile_all_calls mode begin
+            test_all_combinations(;
+                binary_operators=binary_operators,
+                unary_operators=unary_operators,
+                turbo=turbo,
+                types=types,
+            )
+        end
+        operators = OperatorEnum(;
+            binary_operators=binary_operators[1],
+            unary_operators=unary_operators[1],
+            define_helper_functions=false,
+        )
+        # Want to precompile all above calls.
+        types = [Float16, Float32, Float64]
+        for T in types
+            @maybe_precompile_all_calls mode begin
+                test_functions_on_trees(T, operators)
+            end
+        end
+    end
+end
diff --git a/test/test_evaluation.jl b/test/test_evaluation.jl
@@ -111,30 +111,40 @@ for turbo in [false, true], T in [Float16, Float32, Float64]
     @test isnan(tree(X; turbo=turbo)[1])
 end
 
-# And, with generic operator enum, this should be an actual error:
-operators = GenericOperatorEnum(; binary_operators=[+, -, *, /], unary_operators=[cos, sin])
-x1 = Node(Float64; feature=1)
-tree = sin(x1 / 0.0)
-X = randn(Float32, 10);
-@noinline stack = try
-    tree(X)[1]
-    @test false
-catch e
-    @test isa(e, ErrorException)
-    # Check that "Failed to evaluate" is in the message:
-    @test occursin("Failed to evaluate", e.msg)
-    current_exceptions()
-end;
-@test length(stack) == 2
-@test isa(stack[1].exception, DomainError)
-
-# If a method is not defined, we should get a nothing:
-X = randn(Float32, 1, 10);
-@test tree(X; throw_errors=false) === nothing
-# or a MethodError:
-try
-    tree(X; throw_errors=true)
-    @test false
-catch e
-    @test isa(current_exceptions()[1].exception, MethodError)
+# Check if julia version >= 1.7:
+if VERSION >= v"1.7"
+    # And, with generic operator enum, this should be an actual error:
+    operators = GenericOperatorEnum(;
+        binary_operators=[+, -, *, /], unary_operators=[cos, sin]
+    )
+    x1 = Node(Float64; feature=1)
+    tree = sin(x1 / 0.0)
+    X = randn(Float32, 10)
+    local stack
+    try
+        tree(X)[1]
+        @test false
+    catch e
+        @test e isa ErrorException
+        # Check that "Failed to evaluate" is in the message:
+        @test occursin("Failed to evaluate", e.msg)
+        stack = current_exceptions()
+    end
+    @test length(stack) == 2
+    @test stack[1].exception isa DomainError
+
+    # If a method is not defined, we should get a nothing:
+    X = randn(Float32, 1, 10)
+    @test tree(X; throw_errors=false) === nothing
+    # or a MethodError:
+    try
+        tree(X; throw_errors=true)
+        @test false
+    catch e
+        @test e isa ErrorException
+        @test occursin("Failed to evaluate", e.msg)
+        stack = current_exceptions()
+    end
+    @test length(stack) == 2
+    @test stack[1].exception isa MethodError
 end
diff --git a/test/test_precompilation.jl b/test/test_precompilation.jl
diff --git a/test/unittest.jl b/test/unittest.jl
