Precompilation can lead to Dual tag ordering problems

[Technici4n]

Hi, I found an issue when working with precompilation and nested derivatives.

The tag count system is setup such that the inner derivative receives a tag with a higher tagcount than the outer derivative. When precompiling, it can happen that only the outer derivative is precompiled, for example because of runtime dispatch that depends on the value of some parameters. When running the code later, the atomic counter for tags starts from 0 again, which can cause the inner derivative to be given a tag with a lower count. This leads to a Cannot determine ordering of Dual tags error.

Here is a reproduction example: https://github.com/Technici4n/FDPrecompilationIssue.
Precompilation should succeed, and running the test should fail with the following output:

     Testing Running tests...
Tag count of Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}}(0,1) is 1
Tag count of Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#1#2"{ForwardDiff.Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}, Int64, 1}}, ForwardDiff.Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}, Int64, 1}}}(Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}}(0,0),Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}}(0,2)) is 0
ERROR: LoadError: Cannot determine ordering of Dual tags ForwardDiff.Tag{FDPrecompilationIssue.var"#1#2"{ForwardDiff.Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}, Int64, 1}}, ForwardDiff.Dual{ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}, Int64, 1}} and ForwardDiff.Tag{FDPrecompilationIssue.var"#3#4"{Int64}, Int64}

(I encountered this issue was encountered in the context of https://github.com/mfherbst/AiidaDFTK.jl but the setup is more complicated than in this minimal reproduction example. AiidaDFTK precompiles a stress computation (outer derivative) using the LDA functional (does not require a nested derivative). I encountered the issue when running a stress computation through AiidaDFTK using a PBE functional (which requires a nested gradient). You can think of α from the reproduction example as the parameter that selects the functional, LDA as α = 0, PBE as α = 1.)

[gdalle]

Is there a way to detect from inside the package if we're currently pre-compiling it? If so, maybe the tag counter could be initialized at typemin(Int) instead of 0?

[gdalle]

Although I'm not sure I understand why the counter restarts at zero, shouldn't it be the same state as it was at the end of precompilation?

https://www.oxinabox.net/2023/06/16/top-level-code-in-julia.html

[KristofferC]

TAGCOUNT is zero when ForwardDiff finishes precompiling. Mutations to it during precompilation of another package cannot change that.