From 9b0a5c616067f6e99349b8a12bc62a4789184479 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 18 May 2024 11:27:39 -0700 Subject: [PATCH 1/2] WIP transition to package extensions --- .cirrus.yml | 2 +- .github/workflows/CI.yml | 2 +- Manifest.toml | 390 ++++++++++++++++++ Project.toml | 45 +- ext/ADNLPModelsEnzymeExt.jl | 23 ++ ext/ADNLPModelsSparseDiffToolsExt.jl | 140 +++++++ .../ADNLPModelsSymbolicsExt.jl | 7 + ext/ADNLPModelsZygoteExt.jl | 121 ++++++ src/ADNLPModels.jl | 35 +- src/enzyme.jl | 21 - src/sparse_diff_tools.jl | 145 ------- src/zygote.jl | 119 ------ test/Project.toml | 28 -- test/runtests.jl | 18 +- 14 files changed, 745 insertions(+), 351 deletions(-) create mode 100644 Manifest.toml create mode 100644 ext/ADNLPModelsEnzymeExt.jl create mode 100644 ext/ADNLPModelsSparseDiffToolsExt.jl rename src/sparse_sym.jl => ext/ADNLPModelsSymbolicsExt.jl (96%) create mode 100644 ext/ADNLPModelsZygoteExt.jl delete mode 100644 src/enzyme.jl delete mode 100644 src/sparse_diff_tools.jl delete mode 100644 src/zygote.jl delete mode 100644 test/Project.toml diff --git a/.cirrus.yml b/.cirrus.yml index 48c3fb18..0683f66c 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -5,7 +5,7 @@ task: image_family: freebsd-13-3 env: matrix: - - JULIA_VERSION: 1.6 + - JULIA_VERSION: 1.9 - JULIA_VERSION: 1 - name: musl Linux container: diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 10f060b4..99035b05 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -10,7 +10,7 @@ jobs: fail-fast: false matrix: version: - - '1.6' + - '1.9' - '1' - 'nightly' os: diff --git a/Manifest.toml b/Manifest.toml new file mode 100644 index 00000000..d8f41f14 --- /dev/null +++ b/Manifest.toml @@ -0,0 +1,390 @@ +# This file is machine-generated - editing it directly is not advised + +julia_version = "1.10.3" +manifest_format = "2.0" +project_hash = "bb1d07f7da969f7a910a69afe0616febf618ea69" + +[[deps.AMD]] +deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse_jll"] +git-tree-sha1 = "45a1272e3f809d36431e57ab22703c6896b8908f" +uuid = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e" +version = "0.5.3" + +[[deps.ArgTools]] +uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" +version = "1.1.1" + +[[deps.Artifacts]] +uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" + +[[deps.Base64]] +uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" + +[[deps.ChainRulesCore]] +deps = ["Compat", "LinearAlgebra"] +git-tree-sha1 = "575cd02e080939a33b6df6c5853d14924c08e35b" +uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +version = "1.23.0" +weakdeps = ["SparseArrays"] + + [deps.ChainRulesCore.extensions] + ChainRulesCoreSparseArraysExt = "SparseArrays" + +[[deps.ColPack]] +deps = ["ColPack_jll", "LinearAlgebra", "Random", "SparseArrays"] +git-tree-sha1 = "83a23545e7969d8b21fb85271b9cd04c8e09d08b" +uuid = "ffa27691-3a59-46ab-a8d4-551f45b8d401" +version = "0.3.0" + +[[deps.ColPack_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "2a518018a2b888ba529e944d34d4bd84b54d652d" +uuid = "f218ff0c-cb54-5151-80c4-c0f62c730ce6" +version = "0.3.0+0" + +[[deps.CommonSubexpressions]] +deps = ["MacroTools", "Test"] +git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7" +uuid = "bbf7d656-a473-5ed7-a52c-81e309532950" +version = "0.3.0" + +[[deps.Compat]] +deps = ["TOML", "UUIDs"] +git-tree-sha1 = "b1c55339b7c6c350ee89f2c1604299660525b248" +uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" +version = "4.15.0" +weakdeps = ["Dates", "LinearAlgebra"] + + [deps.Compat.extensions] + CompatLinearAlgebraExt = "LinearAlgebra" + +[[deps.CompilerSupportLibraries_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" +version = "1.1.1+0" + +[[deps.Dates]] +deps = ["Printf"] +uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" + +[[deps.DiffResults]] +deps = ["StaticArraysCore"] +git-tree-sha1 = "782dd5f4561f5d267313f23853baaaa4c52ea621" +uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" +version = "1.1.0" + +[[deps.DiffRules]] +deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"] +git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272" +uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" +version = "1.15.1" + +[[deps.DocStringExtensions]] +deps = ["LibGit2"] +git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d" +uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +version = "0.9.3" + +[[deps.Downloads]] +deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] +uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +version = "1.6.0" + +[[deps.ExprTools]] +git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" +uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" +version = "0.1.10" + +[[deps.FastClosures]] +git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef" +uuid = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" +version = "0.3.2" + +[[deps.FileWatching]] +uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" + +[[deps.ForwardDiff]] +deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions"] +git-tree-sha1 = "cf0fe81336da9fb90944683b8c41984b08793dad" +uuid = "f6369f11-7733-5829-9624-2563aa707210" +version = "0.10.36" +weakdeps = ["StaticArrays"] + + [deps.ForwardDiff.extensions] + ForwardDiffStaticArraysExt = "StaticArrays" + +[[deps.FunctionWrappers]] +git-tree-sha1 = "d62485945ce5ae9c0c48f124a84998d755bae00e" +uuid = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" +version = "1.1.3" + +[[deps.InteractiveUtils]] +deps = ["Markdown"] +uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" + +[[deps.IrrationalConstants]] +git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" +uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" +version = "0.2.2" + +[[deps.JLLWrappers]] +deps = ["Artifacts", "Preferences"] +git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca" +uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" +version = "1.5.0" + +[[deps.LDLFactorizations]] +deps = ["AMD", "LinearAlgebra", "SparseArrays", "Test"] +git-tree-sha1 = "70f582b446a1c3ad82cf87e62b878668beef9d13" +uuid = "40e66cde-538c-5869-a4ad-c39174c6795b" +version = "0.10.1" + +[[deps.LibCURL]] +deps = ["LibCURL_jll", "MozillaCACerts_jll"] +uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" +version = "0.6.4" + +[[deps.LibCURL_jll]] +deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"] +uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" +version = "8.4.0+0" + +[[deps.LibGit2]] +deps = ["Base64", "LibGit2_jll", "NetworkOptions", "Printf", "SHA"] +uuid = "76f85450-5226-5b5a-8eaa-529ad045b433" + +[[deps.LibGit2_jll]] +deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll"] +uuid = "e37daf67-58a4-590a-8e99-b0245dd2ffc5" +version = "1.6.4+0" + +[[deps.LibSSH2_jll]] +deps = ["Artifacts", "Libdl", "MbedTLS_jll"] +uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8" +version = "1.11.0+1" + +[[deps.Libdl]] +uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb" + +[[deps.LinearAlgebra]] +deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] +uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + +[[deps.LinearOperators]] +deps = ["FastClosures", "LDLFactorizations", "LinearAlgebra", "Printf", "Requires", "SparseArrays", "TimerOutputs"] +git-tree-sha1 = "f06df3a46255879cbccae1b5b6dcb16994c31be7" +uuid = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125" +version = "2.7.0" +weakdeps = ["ChainRulesCore"] + + [deps.LinearOperators.extensions] + LinearOperatorsChainRulesCoreExt = "ChainRulesCore" + +[[deps.LogExpFunctions]] +deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] +git-tree-sha1 = "18144f3e9cbe9b15b070288eef858f71b291ce37" +uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +version = "0.3.27" + + [deps.LogExpFunctions.extensions] + LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" + LogExpFunctionsChangesOfVariablesExt = "ChangesOfVariables" + LogExpFunctionsInverseFunctionsExt = "InverseFunctions" + + [deps.LogExpFunctions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" + InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" + +[[deps.Logging]] +uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" + +[[deps.MacroTools]] +deps = ["Markdown", "Random"] +git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df" +uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" +version = "0.5.13" + +[[deps.Markdown]] +deps = ["Base64"] +uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" + +[[deps.MbedTLS_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1" +version = "2.28.2+1" + +[[deps.MozillaCACerts_jll]] +uuid = "14a3606d-f60d-562e-9121-12d972cd8159" +version = "2023.1.10" + +[[deps.NLPModels]] +deps = ["FastClosures", "LinearAlgebra", "LinearOperators", "Printf", "SparseArrays"] +git-tree-sha1 = "2d110433ba53dcf225c1a958b1f91fda6cf3cead" +uuid = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +version = "0.21.0" + +[[deps.NaNMath]] +deps = ["OpenLibm_jll"] +git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" +uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +version = "1.0.2" + +[[deps.NetworkOptions]] +uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" +version = "1.2.0" + +[[deps.OpenBLAS_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] +uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" +version = "0.3.23+4" + +[[deps.OpenLibm_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "05823500-19ac-5b8b-9628-191a04bc5112" +version = "0.8.1+2" + +[[deps.OpenSpecFun_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1" +uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" +version = "0.5.5+0" + +[[deps.Pkg]] +deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] +uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +version = "1.10.0" + +[[deps.PrecompileTools]] +deps = ["Preferences"] +git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f" +uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +version = "1.2.1" + +[[deps.Preferences]] +deps = ["TOML"] +git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6" +uuid = "21216c6a-2e73-6563-6e65-726566657250" +version = "1.4.3" + +[[deps.Printf]] +deps = ["Unicode"] +uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" + +[[deps.REPL]] +deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] +uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" + +[[deps.Random]] +deps = ["SHA"] +uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" + +[[deps.Requires]] +deps = ["UUIDs"] +git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7" +uuid = "ae029012-a4dd-5104-9daa-d747884805df" +version = "1.3.0" + +[[deps.ReverseDiff]] +deps = ["ChainRulesCore", "DiffResults", "DiffRules", "ForwardDiff", "FunctionWrappers", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "Random", "SpecialFunctions", "StaticArrays", "Statistics"] +git-tree-sha1 = "cc6cd622481ea366bb9067859446a8b01d92b468" +uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267" +version = "1.15.3" + +[[deps.SHA]] +uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" +version = "0.7.0" + +[[deps.Serialization]] +uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" + +[[deps.Sockets]] +uuid = "6462fe0b-24de-5631-8697-dd941f90decc" + +[[deps.SparseArrays]] +deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] +uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +version = "1.10.0" + +[[deps.SpecialFunctions]] +deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] +git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14" +uuid = "276daf66-3868-5448-9aa4-cd146d93841b" +version = "2.4.0" +weakdeps = ["ChainRulesCore"] + + [deps.SpecialFunctions.extensions] + SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" + +[[deps.StaticArrays]] +deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] +git-tree-sha1 = "bf074c045d3d5ffd956fa0a461da38a44685d6b2" +uuid = "90137ffa-7385-5640-81b9-e52037218182" +version = "1.9.3" +weakdeps = ["ChainRulesCore", "Statistics"] + + [deps.StaticArrays.extensions] + StaticArraysChainRulesCoreExt = "ChainRulesCore" + StaticArraysStatisticsExt = "Statistics" + +[[deps.StaticArraysCore]] +git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d" +uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +version = "1.4.2" + +[[deps.Statistics]] +deps = ["LinearAlgebra", "SparseArrays"] +uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +version = "1.10.0" + +[[deps.SuiteSparse_jll]] +deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] +uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" +version = "7.2.1+1" + +[[deps.TOML]] +deps = ["Dates"] +uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76" +version = "1.0.3" + +[[deps.Tar]] +deps = ["ArgTools", "SHA"] +uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" +version = "1.10.0" + +[[deps.Test]] +deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] +uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[[deps.TimerOutputs]] +deps = ["ExprTools", "Printf"] +git-tree-sha1 = "f548a9e9c490030e545f72074a41edfd0e5bcdd7" +uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" +version = "0.5.23" + +[[deps.UUIDs]] +deps = ["Random", "SHA"] +uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" + +[[deps.Unicode]] +uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" + +[[deps.Zlib_jll]] +deps = ["Libdl"] +uuid = "83775a58-1f1d-513f-b197-d71354ab007a" +version = "1.2.13+1" + +[[deps.libblastrampoline_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" +version = "5.8.0+1" + +[[deps.nghttp2_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d" +version = "1.52.0+1" + +[[deps.p7zip_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" +version = "17.4.0+2" diff --git a/Project.toml b/Project.toml index f0df9704..96ba1b3d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ADNLPModels" uuid = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a" -version = "0.7.2" +version = "0.8.0" [deps] ColPack = "ffa27691-3a59-46ab-a8d4-551f45b8d401" @@ -11,10 +11,49 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +[weakdeps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" +Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[extensions] +ADNLPModelsEnzymeExt = "Enzyme" +ADNLPModelsSparseDiffToolsExt = "SparseDiffTools" +ADNLPModelsSymbolicsExt = "Symbolics" +ADNLPModelsZygoteExt = "Zygote" + [compat] ColPack = "0.3" ForwardDiff = "0.9.0, 0.10.0" NLPModels = "0.18, 0.19, 0.20, 0.21" -Requires = "1" ReverseDiff = "1" -julia = "^1.6" +julia = "^1.9" +Enzyme = "0.12" +SparseDiffTools = "2.3" +Symbolics = "5.3" +Zygote = "0.6" + +[extras] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" +Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617" +NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = [ + "Test", + "CUDA", + "Enzyme", + "SparseDiffTools", + "Symbolics", + "Zygote", + "ManualNLPModels", + "NLPModelsModifiers", + "NLPModelsTest", +] diff --git a/ext/ADNLPModelsEnzymeExt.jl b/ext/ADNLPModelsEnzymeExt.jl new file mode 100644 index 00000000..13ed9015 --- /dev/null +++ b/ext/ADNLPModelsEnzymeExt.jl @@ -0,0 +1,23 @@ +module ADNLPModelsEnzymeExt + +using Enzyme, ADNLPModels + +struct EnzymeADGradient <: ADNLPModels.ADBackend end + +function EnzymeADGradient( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + x0::AbstractVector = rand(nvar), + kwargs..., +) + return EnzymeADGradient() +end + +function ADNLPModels.gradient!(::EnzymeADGradient, g, f, x) + Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) + return g +end + +end \ No newline at end of file diff --git a/ext/ADNLPModelsSparseDiffToolsExt.jl b/ext/ADNLPModelsSparseDiffToolsExt.jl new file mode 100644 index 00000000..cf204d53 --- /dev/null +++ b/ext/ADNLPModelsSparseDiffToolsExt.jl @@ -0,0 +1,140 @@ +module ADNLPModelsSparseDiffToolsExt + +using SparseDiffTools, SparseArrays, ADNLPModels, NLPModels, ForwardDiff +import ADNLPModels: ADModel, AbstractADNLSModel, InPlaceADbackend, ADBackend + +function sparse_matrix_colors(A, alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm) + return SparseDiffTools.matrix_colors(A, alg) +end + +struct SDTSparseADJacobian{Tv, Ti, T, T2, T3, T4, T5} <: ADNLPModels.ADBackend + cfJ::SparseDiffTools.ForwardColorJacCache{T, T2, T3, T4, T5, SparseMatrixCSC{Tv, Ti}} +end + +function SDTSparseADJacobian( + nvar, + f, + ncon, + c!; + x0::S = rand(nvar), + alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm = SparseDiffTools.GreedyD1Color(), + kwargs..., +) where {S} + T = eltype(S) + output = similar(x0, ncon) + J = compute_jacobian_sparsity(c!, output, x0) + colors = sparse_matrix_colors(J, alg) + jac = SparseMatrixCSC{T, Int}(J.m, J.n, J.colptr, J.rowval, T.(J.nzval)) + + dx = fill!(S(undef, ncon), 0) + cfJ = SparseDiffTools.ForwardColorJacCache(c!, x0, colorvec = colors, dx = dx, sparsity = jac) + SDTSparseADJacobian(cfJ) +end + +function get_nln_nnzj(b::SDTSparseADJacobian, nvar, ncon) + nnz(b.cfJ.sparsity) +end + +function NLPModels.jac_structure!( + b::SDTSparseADJacobian, + nlp::ADModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + rows .= rowvals(b.cfJ.sparsity) + for i = 1:(nlp.meta.nvar) + for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols +end + +function NLPModels.jac_coord!( + b::SDTSparseADJacobian, + nlp::ADModel, + x::AbstractVector, + vals::AbstractVector, +) + SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nlp.c!, x, b.cfJ) + vals .= nonzeros(b.cfJ.sparsity) + return vals +end + +function NLPModels.jac_structure_residual!( + b::SDTSparseADJacobian, + nls::AbstractADNLSModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + rows .= rowvals(b.cfJ.sparsity) + for i = 1:(nls.meta.nvar) + for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols +end + +function NLPModels.jac_coord_residual!( + b::SDTSparseADJacobian, + nls::AbstractADNLSModel, + x::AbstractVector, + vals::AbstractVector, +) + SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nls.F!, x, b.cfJ) + vals .= nonzeros(b.cfJ.sparsity) + return vals +end + +struct SDTForwardDiffADJprod{T} <: InPlaceADbackend + tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} + tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} +end + +function SDTForwardDiffADJprod( + nvar::Integer, + f, + ncon::Integer = 0, + c!::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + tmp_in = + Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, nvar) + tmp_out = + Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, ncon) + return SDTForwardDiffADJprod(tmp_in, tmp_out) +end + +function Jprod!(b::SDTForwardDiffADJprod, Jv, c!, x, v, ::Val) + SparseDiffTools.auto_jacvec!(Jv, c!, x, v, b.tmp_in, b.tmp_out) + return Jv +end + +struct SDTForwardDiffADHvprod{T} <: ADBackend + tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} + tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} +end +function SDTForwardDiffADHvprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., +) where {T} + tmp_in = + Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, nvar) + tmp_out = + Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, nvar) + return SDTForwardDiffADHvprod(tmp_in, tmp_out) +end + +function Hvprod!(b::SDTForwardDiffADHvprod, Hv, x, v, f, args...) + ϕ!(dy, x; f = f) = ForwardDiff.gradient!(dy, f, x) + SparseDiffTools.auto_hesvecgrad!(Hv, (dy, x) -> ϕ!(dy, x), x, v, b.tmp_in, b.tmp_out) + return Hv +end + +end \ No newline at end of file diff --git a/src/sparse_sym.jl b/ext/ADNLPModelsSymbolicsExt.jl similarity index 96% rename from src/sparse_sym.jl rename to ext/ADNLPModelsSymbolicsExt.jl index e90c5514..6eda375f 100644 --- a/src/sparse_sym.jl +++ b/ext/ADNLPModelsSymbolicsExt.jl @@ -1,3 +1,8 @@ +module ADNLPModelsSymbolicsExt + +using Symbolics, ADNLPModels, NLPModels +import ADNLPModels: ADBackend, ADModel, AbstractADNLSModel + function compute_hessian_sparsity(f, nvar, c!, ncon) Symbolics.@variables xs[1:nvar] xsi = Symbolics.scalarize(xs) @@ -172,3 +177,5 @@ function NLPModels.hess_coord!( @eval $(b.cfH)($vals, $x, $(b.y), $obj_weight) return vals end + +end \ No newline at end of file diff --git a/ext/ADNLPModelsZygoteExt.jl b/ext/ADNLPModelsZygoteExt.jl new file mode 100644 index 00000000..9bb8707f --- /dev/null +++ b/ext/ADNLPModelsZygoteExt.jl @@ -0,0 +1,121 @@ +module ADNLPModelsZygoteExt + +using Zygote, ADNLPModels +import ADNLPModels: ADModel, AbstractADNLSModel, ADBackend, ImmutableADbackend + +struct ZygoteADGradient <: ADBackend end +struct ZygoteADJacobian <: ImmutableADbackend + nnzj::Int +end +struct ZygoteADHessian <: ImmutableADbackend + nnzh::Int +end +struct ZygoteADJprod <: ImmutableADbackend end +struct ZygoteADJtprod <: ImmutableADbackend end +# See https://fluxml.ai/Zygote.jl/latest/limitations/ +function get_immutable_c(nlp::ADModel) + function c(x; nnln = nlp.meta.nnln) + c = Zygote.Buffer(x, nnln) + nlp.c!(c, x) + return copy(c) + end + return c +end +get_c(nlp::ADModel, ::ImmutableADbackend) = get_immutable_c(nlp) + +function get_immutable_F(nls::AbstractADNLSModel) + function F(x; nequ = nls.nls_meta.nequ) + Fx = Zygote.Buffer(x, nequ) + nls.F!(Fx, x) + return copy(Fx) + end + return F +end +get_F(nls::AbstractADNLSModel, ::ImmutableADbackend) = get_immutable_F(nls) + +function ZygoteADGradient( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + return ZygoteADGradient() +end +function gradient(::ZygoteADGradient, f, x) + g = Zygote.gradient(f, x)[1] + return g === nothing ? zero(x) : g +end +function gradient!(::ZygoteADGradient, g, f, x) + _g = Zygote.gradient(f, x)[1] + g .= _g === nothing ? 0 : _g +end + +function ZygoteADJacobian( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + @assert nvar > 0 + nnzj = nvar * ncon + return ZygoteADJacobian(nnzj) +end +function jacobian(::ZygoteADJacobian, f, x) + return Zygote.jacobian(f, x)[1] +end + +function ZygoteADHessian( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + @assert nvar > 0 + nnzh = nvar * (nvar + 1) / 2 + return ZygoteADHessian(nnzh) +end +function hessian(b::ZygoteADHessian, f, x) + return jacobian( + ForwardDiffADJacobian(length(x), f, x0 = x), + x -> gradient(ZygoteADGradient(), f, x), + x, + ) +end + +function ZygoteADJprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + return ZygoteADJprod() +end +function Jprod!(::ZygoteADJprod, Jv, f, x, v, ::Val) + Jv .= vec(Zygote.jacobian(t -> f(x + t * v), 0)[1]) + return Jv +end + +function ZygoteADJtprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + kwargs..., +) + return ZygoteADJtprod() +end +function Jtprod!(::ZygoteADJtprod, Jtv, f, x, v, ::Val) + g = Zygote.gradient(x -> dot(f(x), v), x)[1] + if g === nothing + Jtv .= zero(x) + else + Jtv .= g + end + return Jtv +end + +end \ No newline at end of file diff --git a/src/ADNLPModels.jl b/src/ADNLPModels.jl index 481724bb..84b72430 100644 --- a/src/ADNLPModels.jl +++ b/src/ADNLPModels.jl @@ -6,7 +6,6 @@ using LinearAlgebra, SparseArrays using ColPack, ForwardDiff, ReverseDiff # JSO using NLPModels -using Requires abstract type AbstractADNLPModel{T, S} <: AbstractNLPModel{T, S} end abstract type AbstractADNLSModel{T, S} <: AbstractNLSModel{T, S} end @@ -45,11 +44,25 @@ end include("sparse_jacobian.jl") include("sparse_hessian.jl") +# Attempt to load some symbols from the Symbolics extensions +symbolics_ext = Base.get_extension(@__MODULE__, :ADNLPModelsSymbolicsExt) +if !isnothing(symbolics_ext) + SparseSymbolicsADJacobian = symbolics_ext.SparseSymbolicsADJacobian + SparseSymbolicsADHessian = symbolics_ext.SparseSymbolicsADHessian + SDTSparseADJacobian = symbolics_ext.SDTSparseADJacobian + + # These backends should only be included if the module is loaded + predefined_backend[:default][:jacobian_backend] = SparseADJacobian + predefined_backend[:default][:jacobian_residual_backend] = SparseADJacobian + predefined_backend[:optimized][:jacobian_backend] = SparseADJacobian + predefined_backend[:optimized][:jacobian_residual_backend] = SparseADJacobian + + predefined_backend[:default][:hessian_backend] = SparseADHessian + predefined_backend[:optimized][:hessian_backend] = SparseReverseADHessian +end + include("forward.jl") include("reverse.jl") -include("enzyme.jl") -include("sparse_diff_tools.jl") -include("zygote.jl") include("predefined_backend.jl") include("nlp.jl") @@ -181,20 +194,6 @@ function ADNLSModel!(model::AbstractNLSModel; kwargs...) end end -@init begin - @require Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" begin - include("sparse_sym.jl") - - predefined_backend[:default][:jacobian_backend] = SparseADJacobian - predefined_backend[:default][:jacobian_residual_backend] = SparseADJacobian - predefined_backend[:optimized][:jacobian_backend] = SparseADJacobian - predefined_backend[:optimized][:jacobian_residual_backend] = SparseADJacobian - - predefined_backend[:default][:hessian_backend] = SparseADHessian - predefined_backend[:optimized][:hessian_backend] = SparseReverseADHessian - end -end - export get_adbackend, set_adbackend! """ diff --git a/src/enzyme.jl b/src/enzyme.jl deleted file mode 100644 index db5133fe..00000000 --- a/src/enzyme.jl +++ /dev/null @@ -1,21 +0,0 @@ -struct EnzymeADGradient <: ADNLPModels.ADBackend end - -function EnzymeADGradient( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - x0::AbstractVector = rand(nvar), - kwargs..., -) - return EnzymeADGradient() -end - -@init begin - @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin - function ADNLPModels.gradient!(::EnzymeADGradient, g, f, x) - Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Duplicated(x, g)) # gradient!(Reverse, g, f, x) - return g - end - end -end diff --git a/src/sparse_diff_tools.jl b/src/sparse_diff_tools.jl deleted file mode 100644 index 8ec2917d..00000000 --- a/src/sparse_diff_tools.jl +++ /dev/null @@ -1,145 +0,0 @@ -@init begin - @require SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" begin - function sparse_matrix_colors(A, alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm) - return SparseDiffTools.matrix_colors(A, alg) - end - - struct SDTSparseADJacobian{Tv, Ti, T, T2, T3, T4, T5} <: ADNLPModels.ADBackend - cfJ::SparseDiffTools.ForwardColorJacCache{T, T2, T3, T4, T5, SparseMatrixCSC{Tv, Ti}} - end - - function SDTSparseADJacobian( - nvar, - f, - ncon, - c!; - x0::S = rand(nvar), - alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm = SparseDiffTools.GreedyD1Color(), - kwargs..., - ) where {S} - T = eltype(S) - output = similar(x0, ncon) - J = compute_jacobian_sparsity(c!, output, x0) - colors = sparse_matrix_colors(J, alg) - jac = SparseMatrixCSC{T, Int}(J.m, J.n, J.colptr, J.rowval, T.(J.nzval)) - - dx = fill!(S(undef, ncon), 0) - cfJ = SparseDiffTools.ForwardColorJacCache(c!, x0, colorvec = colors, dx = dx, sparsity = jac) - SDTSparseADJacobian(cfJ) - end - - function get_nln_nnzj(b::SDTSparseADJacobian, nvar, ncon) - nnz(b.cfJ.sparsity) - end - - function NLPModels.jac_structure!( - b::SDTSparseADJacobian, - nlp::ADModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, - ) - rows .= rowvals(b.cfJ.sparsity) - for i = 1:(nlp.meta.nvar) - for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) - cols[j] = i - end - end - return rows, cols - end - - function NLPModels.jac_coord!( - b::SDTSparseADJacobian, - nlp::ADModel, - x::AbstractVector, - vals::AbstractVector, - ) - SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nlp.c!, x, b.cfJ) - vals .= nonzeros(b.cfJ.sparsity) - return vals - end - - function NLPModels.jac_structure_residual!( - b::SDTSparseADJacobian, - nls::AbstractADNLSModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, - ) - rows .= rowvals(b.cfJ.sparsity) - for i = 1:(nls.meta.nvar) - for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) - cols[j] = i - end - end - return rows, cols - end - - function NLPModels.jac_coord_residual!( - b::SDTSparseADJacobian, - nls::AbstractADNLSModel, - x::AbstractVector, - vals::AbstractVector, - ) - SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nls.F!, x, b.cfJ) - vals .= nonzeros(b.cfJ.sparsity) - return vals - end - - struct SDTForwardDiffADJprod{T} <: InPlaceADbackend - tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} - tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} - end - - function SDTForwardDiffADJprod( - nvar::Integer, - f, - ncon::Integer = 0, - c!::Function = (args...) -> []; - x0::AbstractVector{T} = rand(nvar), - kwargs..., - ) where {T} - tmp_in = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( - undef, - nvar, - ) - tmp_out = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( - undef, - ncon, - ) - return SDTForwardDiffADJprod(tmp_in, tmp_out) - end - - function Jprod!(b::SDTForwardDiffADJprod, Jv, c!, x, v, ::Val) - SparseDiffTools.auto_jacvec!(Jv, c!, x, v, b.tmp_in, b.tmp_out) - return Jv - end - - struct SDTForwardDiffADHvprod{T} <: ADBackend - tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} - tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} - end - function SDTForwardDiffADHvprod( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - x0::AbstractVector{T} = rand(nvar), - kwargs..., - ) where {T} - tmp_in = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( - undef, - nvar, - ) - tmp_out = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( - undef, - nvar, - ) - return SDTForwardDiffADHvprod(tmp_in, tmp_out) - end - - function Hvprod!(b::SDTForwardDiffADHvprod, Hv, x, v, f, args...) - ϕ!(dy, x; f = f) = ForwardDiff.gradient!(dy, f, x) - SparseDiffTools.auto_hesvecgrad!(Hv, (dy, x) -> ϕ!(dy, x), x, v, b.tmp_in, b.tmp_out) - return Hv - end - end -end diff --git a/src/zygote.jl b/src/zygote.jl deleted file mode 100644 index 63358a7e..00000000 --- a/src/zygote.jl +++ /dev/null @@ -1,119 +0,0 @@ -struct ZygoteADGradient <: ADBackend end -struct ZygoteADJacobian <: ImmutableADbackend - nnzj::Int -end -struct ZygoteADHessian <: ImmutableADbackend - nnzh::Int -end -struct ZygoteADJprod <: ImmutableADbackend end -struct ZygoteADJtprod <: ImmutableADbackend end - -@init begin - @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin - # See https://fluxml.ai/Zygote.jl/latest/limitations/ - function get_immutable_c(nlp::ADModel) - function c(x; nnln = nlp.meta.nnln) - c = Zygote.Buffer(x, nnln) - nlp.c!(c, x) - return copy(c) - end - return c - end - get_c(nlp::ADModel, ::ImmutableADbackend) = get_immutable_c(nlp) - - function get_immutable_F(nls::AbstractADNLSModel) - function F(x; nequ = nls.nls_meta.nequ) - Fx = Zygote.Buffer(x, nequ) - nls.F!(Fx, x) - return copy(Fx) - end - return F - end - get_F(nls::AbstractADNLSModel, ::ImmutableADbackend) = get_immutable_F(nls) - - function ZygoteADGradient( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - kwargs..., - ) - return ZygoteADGradient() - end - function gradient(::ZygoteADGradient, f, x) - g = Zygote.gradient(f, x)[1] - return g === nothing ? zero(x) : g - end - function gradient!(::ZygoteADGradient, g, f, x) - _g = Zygote.gradient(f, x)[1] - g .= _g === nothing ? 0 : _g - end - - function ZygoteADJacobian( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - kwargs..., - ) - @assert nvar > 0 - nnzj = nvar * ncon - return ZygoteADJacobian(nnzj) - end - function jacobian(::ZygoteADJacobian, f, x) - return Zygote.jacobian(f, x)[1] - end - - function ZygoteADHessian( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - kwargs..., - ) - @assert nvar > 0 - nnzh = nvar * (nvar + 1) / 2 - return ZygoteADHessian(nnzh) - end - function hessian(b::ZygoteADHessian, f, x) - return jacobian( - ForwardDiffADJacobian(length(x), f, x0 = x), - x -> gradient(ZygoteADGradient(), f, x), - x, - ) - end - - function ZygoteADJprod( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - kwargs..., - ) - return ZygoteADJprod() - end - function Jprod!(::ZygoteADJprod, Jv, f, x, v, ::Val) - Jv .= vec(Zygote.jacobian(t -> f(x + t * v), 0)[1]) - return Jv - end - - function ZygoteADJtprod( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - kwargs..., - ) - return ZygoteADJtprod() - end - function Jtprod!(::ZygoteADJtprod, Jtv, f, x, v, ::Val) - g = Zygote.gradient(x -> dot(f(x), v), x)[1] - if g === nothing - Jtv .= zero(x) - else - Jtv .= g - end - return Jtv - end - end -end diff --git a/test/Project.toml b/test/Project.toml deleted file mode 100644 index d254894f..00000000 --- a/test/Project.toml +++ /dev/null @@ -1,28 +0,0 @@ -[deps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -ManualNLPModels = "30dfa513-9b2f-4fb3-9796-781eabac1617" -NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" -NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" -NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" -ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" -Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[compat] -CUDA = "4, 5" -Enzyme = "0.10, 0.11, 0.12" -ForwardDiff = "0.10" -ManualNLPModels = "0.1" -NLPModels = "0.21" -NLPModelsModifiers = "0.7" -NLPModelsTest = "0.10" -ReverseDiff = "1" -SparseDiffTools = "2.3" -Symbolics = "5.3" -Zygote = "0.6" diff --git a/test/runtests.jl b/test/runtests.jl index 3eea51f6..99729969 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,23 +1,11 @@ +# Load extension dependencies first +using SparseDiffTools, Symbolics + using CUDA, LinearAlgebra, SparseArrays, Test using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! -@testset "Error without loading package for sparsity pattern" begin - f(x) = sum(x) - c!(cx, x) = begin - cx .= 1 - return x - end - nvar, ncon = 2, 1 - x0 = ones(nvar) - cx = rand(ncon) - @test_throws ArgumentError ADNLPModels.compute_jacobian_sparsity(c!, cx, x0) - @test_throws ArgumentError ADNLPModels.compute_hessian_sparsity(f, nvar, c!, ncon) -end - -using SparseDiffTools, Symbolics - @testset "Test using a NLPModel instead of AD-backend" begin include("manual.jl") end From 606984f9c12e900ab5631ecba6f67508cd0ad248 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 18 May 2024 17:53:09 -0700 Subject: [PATCH 2/2] Only use package extensions for AD backends --- Project.toml | 2 - ext/ADNLPModelsSparseDiffToolsExt.jl | 140 ----------------- src/ADNLPModels.jl | 29 ++-- src/sparse_diff_tools.jl | 145 ++++++++++++++++++ .../sparse_sym.jl | 7 - test/runtests.jl | 18 ++- 6 files changed, 174 insertions(+), 167 deletions(-) delete mode 100644 ext/ADNLPModelsSparseDiffToolsExt.jl create mode 100644 src/sparse_diff_tools.jl rename ext/ADNLPModelsSymbolicsExt.jl => src/sparse_sym.jl (96%) diff --git a/Project.toml b/Project.toml index 96ba1b3d..8c53e1c9 100644 --- a/Project.toml +++ b/Project.toml @@ -19,8 +19,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] ADNLPModelsEnzymeExt = "Enzyme" -ADNLPModelsSparseDiffToolsExt = "SparseDiffTools" -ADNLPModelsSymbolicsExt = "Symbolics" ADNLPModelsZygoteExt = "Zygote" [compat] diff --git a/ext/ADNLPModelsSparseDiffToolsExt.jl b/ext/ADNLPModelsSparseDiffToolsExt.jl deleted file mode 100644 index cf204d53..00000000 --- a/ext/ADNLPModelsSparseDiffToolsExt.jl +++ /dev/null @@ -1,140 +0,0 @@ -module ADNLPModelsSparseDiffToolsExt - -using SparseDiffTools, SparseArrays, ADNLPModels, NLPModels, ForwardDiff -import ADNLPModels: ADModel, AbstractADNLSModel, InPlaceADbackend, ADBackend - -function sparse_matrix_colors(A, alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm) - return SparseDiffTools.matrix_colors(A, alg) -end - -struct SDTSparseADJacobian{Tv, Ti, T, T2, T3, T4, T5} <: ADNLPModels.ADBackend - cfJ::SparseDiffTools.ForwardColorJacCache{T, T2, T3, T4, T5, SparseMatrixCSC{Tv, Ti}} -end - -function SDTSparseADJacobian( - nvar, - f, - ncon, - c!; - x0::S = rand(nvar), - alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm = SparseDiffTools.GreedyD1Color(), - kwargs..., -) where {S} - T = eltype(S) - output = similar(x0, ncon) - J = compute_jacobian_sparsity(c!, output, x0) - colors = sparse_matrix_colors(J, alg) - jac = SparseMatrixCSC{T, Int}(J.m, J.n, J.colptr, J.rowval, T.(J.nzval)) - - dx = fill!(S(undef, ncon), 0) - cfJ = SparseDiffTools.ForwardColorJacCache(c!, x0, colorvec = colors, dx = dx, sparsity = jac) - SDTSparseADJacobian(cfJ) -end - -function get_nln_nnzj(b::SDTSparseADJacobian, nvar, ncon) - nnz(b.cfJ.sparsity) -end - -function NLPModels.jac_structure!( - b::SDTSparseADJacobian, - nlp::ADModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) - rows .= rowvals(b.cfJ.sparsity) - for i = 1:(nlp.meta.nvar) - for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) - cols[j] = i - end - end - return rows, cols -end - -function NLPModels.jac_coord!( - b::SDTSparseADJacobian, - nlp::ADModel, - x::AbstractVector, - vals::AbstractVector, -) - SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nlp.c!, x, b.cfJ) - vals .= nonzeros(b.cfJ.sparsity) - return vals -end - -function NLPModels.jac_structure_residual!( - b::SDTSparseADJacobian, - nls::AbstractADNLSModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) - rows .= rowvals(b.cfJ.sparsity) - for i = 1:(nls.meta.nvar) - for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) - cols[j] = i - end - end - return rows, cols -end - -function NLPModels.jac_coord_residual!( - b::SDTSparseADJacobian, - nls::AbstractADNLSModel, - x::AbstractVector, - vals::AbstractVector, -) - SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nls.F!, x, b.cfJ) - vals .= nonzeros(b.cfJ.sparsity) - return vals -end - -struct SDTForwardDiffADJprod{T} <: InPlaceADbackend - tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} - tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} -end - -function SDTForwardDiffADJprod( - nvar::Integer, - f, - ncon::Integer = 0, - c!::Function = (args...) -> []; - x0::AbstractVector{T} = rand(nvar), - kwargs..., -) where {T} - tmp_in = - Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, nvar) - tmp_out = - Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, ncon) - return SDTForwardDiffADJprod(tmp_in, tmp_out) -end - -function Jprod!(b::SDTForwardDiffADJprod, Jv, c!, x, v, ::Val) - SparseDiffTools.auto_jacvec!(Jv, c!, x, v, b.tmp_in, b.tmp_out) - return Jv -end - -struct SDTForwardDiffADHvprod{T} <: ADBackend - tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} - tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} -end -function SDTForwardDiffADHvprod( - nvar::Integer, - f, - ncon::Integer = 0, - c::Function = (args...) -> []; - x0::AbstractVector{T} = rand(nvar), - kwargs..., -) where {T} - tmp_in = - Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, nvar) - tmp_out = - Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}(undef, nvar) - return SDTForwardDiffADHvprod(tmp_in, tmp_out) -end - -function Hvprod!(b::SDTForwardDiffADHvprod, Hv, x, v, f, args...) - ϕ!(dy, x; f = f) = ForwardDiff.gradient!(dy, f, x) - SparseDiffTools.auto_hesvecgrad!(Hv, (dy, x) -> ϕ!(dy, x), x, v, b.tmp_in, b.tmp_out) - return Hv -end - -end \ No newline at end of file diff --git a/src/ADNLPModels.jl b/src/ADNLPModels.jl index 84b72430..83a2c1b8 100644 --- a/src/ADNLPModels.jl +++ b/src/ADNLPModels.jl @@ -6,6 +6,7 @@ using LinearAlgebra, SparseArrays using ColPack, ForwardDiff, ReverseDiff # JSO using NLPModels +using Requires abstract type AbstractADNLPModel{T, S} <: AbstractNLPModel{T, S} end abstract type AbstractADNLSModel{T, S} <: AbstractNLSModel{T, S} end @@ -44,25 +45,23 @@ end include("sparse_jacobian.jl") include("sparse_hessian.jl") -# Attempt to load some symbols from the Symbolics extensions -symbolics_ext = Base.get_extension(@__MODULE__, :ADNLPModelsSymbolicsExt) -if !isnothing(symbolics_ext) - SparseSymbolicsADJacobian = symbolics_ext.SparseSymbolicsADJacobian - SparseSymbolicsADHessian = symbolics_ext.SparseSymbolicsADHessian - SDTSparseADJacobian = symbolics_ext.SDTSparseADJacobian - - # These backends should only be included if the module is loaded - predefined_backend[:default][:jacobian_backend] = SparseADJacobian - predefined_backend[:default][:jacobian_residual_backend] = SparseADJacobian - predefined_backend[:optimized][:jacobian_backend] = SparseADJacobian - predefined_backend[:optimized][:jacobian_residual_backend] = SparseADJacobian - - predefined_backend[:default][:hessian_backend] = SparseADHessian - predefined_backend[:optimized][:hessian_backend] = SparseReverseADHessian +@init begin + @require Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" begin + include("sparse_sym.jl") + + predefined_backend[:default][:jacobian_backend] = SparseADJacobian + predefined_backend[:default][:jacobian_residual_backend] = SparseADJacobian + predefined_backend[:optimized][:jacobian_backend] = SparseADJacobian + predefined_backend[:optimized][:jacobian_residual_backend] = SparseADJacobian + + predefined_backend[:default][:hessian_backend] = SparseADHessian + predefined_backend[:optimized][:hessian_backend] = SparseReverseADHessian + end end include("forward.jl") include("reverse.jl") +include("sparse_diff_tools.jl") include("predefined_backend.jl") include("nlp.jl") diff --git a/src/sparse_diff_tools.jl b/src/sparse_diff_tools.jl new file mode 100644 index 00000000..8ec2917d --- /dev/null +++ b/src/sparse_diff_tools.jl @@ -0,0 +1,145 @@ +@init begin + @require SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" begin + function sparse_matrix_colors(A, alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm) + return SparseDiffTools.matrix_colors(A, alg) + end + + struct SDTSparseADJacobian{Tv, Ti, T, T2, T3, T4, T5} <: ADNLPModels.ADBackend + cfJ::SparseDiffTools.ForwardColorJacCache{T, T2, T3, T4, T5, SparseMatrixCSC{Tv, Ti}} + end + + function SDTSparseADJacobian( + nvar, + f, + ncon, + c!; + x0::S = rand(nvar), + alg::SparseDiffTools.SparseDiffToolsColoringAlgorithm = SparseDiffTools.GreedyD1Color(), + kwargs..., + ) where {S} + T = eltype(S) + output = similar(x0, ncon) + J = compute_jacobian_sparsity(c!, output, x0) + colors = sparse_matrix_colors(J, alg) + jac = SparseMatrixCSC{T, Int}(J.m, J.n, J.colptr, J.rowval, T.(J.nzval)) + + dx = fill!(S(undef, ncon), 0) + cfJ = SparseDiffTools.ForwardColorJacCache(c!, x0, colorvec = colors, dx = dx, sparsity = jac) + SDTSparseADJacobian(cfJ) + end + + function get_nln_nnzj(b::SDTSparseADJacobian, nvar, ncon) + nnz(b.cfJ.sparsity) + end + + function NLPModels.jac_structure!( + b::SDTSparseADJacobian, + nlp::ADModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + rows .= rowvals(b.cfJ.sparsity) + for i = 1:(nlp.meta.nvar) + for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols + end + + function NLPModels.jac_coord!( + b::SDTSparseADJacobian, + nlp::ADModel, + x::AbstractVector, + vals::AbstractVector, + ) + SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nlp.c!, x, b.cfJ) + vals .= nonzeros(b.cfJ.sparsity) + return vals + end + + function NLPModels.jac_structure_residual!( + b::SDTSparseADJacobian, + nls::AbstractADNLSModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + rows .= rowvals(b.cfJ.sparsity) + for i = 1:(nls.meta.nvar) + for j = b.cfJ.sparsity.colptr[i]:(b.cfJ.sparsity.colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols + end + + function NLPModels.jac_coord_residual!( + b::SDTSparseADJacobian, + nls::AbstractADNLSModel, + x::AbstractVector, + vals::AbstractVector, + ) + SparseDiffTools.forwarddiff_color_jacobian!(b.cfJ.sparsity, nls.F!, x, b.cfJ) + vals .= nonzeros(b.cfJ.sparsity) + return vals + end + + struct SDTForwardDiffADJprod{T} <: InPlaceADbackend + tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} + tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} + end + + function SDTForwardDiffADJprod( + nvar::Integer, + f, + ncon::Integer = 0, + c!::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., + ) where {T} + tmp_in = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( + undef, + nvar, + ) + tmp_out = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( + undef, + ncon, + ) + return SDTForwardDiffADJprod(tmp_in, tmp_out) + end + + function Jprod!(b::SDTForwardDiffADJprod, Jv, c!, x, v, ::Val) + SparseDiffTools.auto_jacvec!(Jv, c!, x, v, b.tmp_in, b.tmp_out) + return Jv + end + + struct SDTForwardDiffADHvprod{T} <: ADBackend + tmp_in::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} + tmp_out::Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}} + end + function SDTForwardDiffADHvprod( + nvar::Integer, + f, + ncon::Integer = 0, + c::Function = (args...) -> []; + x0::AbstractVector{T} = rand(nvar), + kwargs..., + ) where {T} + tmp_in = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( + undef, + nvar, + ) + tmp_out = Vector{SparseDiffTools.Dual{ForwardDiff.Tag{SparseDiffTools.DeivVecTag, T}, T, 1}}( + undef, + nvar, + ) + return SDTForwardDiffADHvprod(tmp_in, tmp_out) + end + + function Hvprod!(b::SDTForwardDiffADHvprod, Hv, x, v, f, args...) + ϕ!(dy, x; f = f) = ForwardDiff.gradient!(dy, f, x) + SparseDiffTools.auto_hesvecgrad!(Hv, (dy, x) -> ϕ!(dy, x), x, v, b.tmp_in, b.tmp_out) + return Hv + end + end +end diff --git a/ext/ADNLPModelsSymbolicsExt.jl b/src/sparse_sym.jl similarity index 96% rename from ext/ADNLPModelsSymbolicsExt.jl rename to src/sparse_sym.jl index 6eda375f..e90c5514 100644 --- a/ext/ADNLPModelsSymbolicsExt.jl +++ b/src/sparse_sym.jl @@ -1,8 +1,3 @@ -module ADNLPModelsSymbolicsExt - -using Symbolics, ADNLPModels, NLPModels -import ADNLPModels: ADBackend, ADModel, AbstractADNLSModel - function compute_hessian_sparsity(f, nvar, c!, ncon) Symbolics.@variables xs[1:nvar] xsi = Symbolics.scalarize(xs) @@ -177,5 +172,3 @@ function NLPModels.hess_coord!( @eval $(b.cfH)($vals, $x, $(b.y), $obj_weight) return vals end - -end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 99729969..3eea51f6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,11 +1,23 @@ -# Load extension dependencies first -using SparseDiffTools, Symbolics - using CUDA, LinearAlgebra, SparseArrays, Test using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest using ADNLPModels: gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod! +@testset "Error without loading package for sparsity pattern" begin + f(x) = sum(x) + c!(cx, x) = begin + cx .= 1 + return x + end + nvar, ncon = 2, 1 + x0 = ones(nvar) + cx = rand(ncon) + @test_throws ArgumentError ADNLPModels.compute_jacobian_sparsity(c!, cx, x0) + @test_throws ArgumentError ADNLPModels.compute_hessian_sparsity(f, nvar, c!, ncon) +end + +using SparseDiffTools, Symbolics + @testset "Test using a NLPModel instead of AD-backend" begin include("manual.jl") end