sqrt behavior at 0 about enzyme HOT 30 OPEN

@tgymnich we already have a flag already within the checkedMul/checkedDiv, which I presume we could just use here.

from enzyme.

Hi @Bike,
just for your question about the documentation, if you are interested into the rules which we have, this is the base place to check:

For older or more complex rules we also have https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/Enzyme/AdjointGenerator.h, but they might be a lot harder to understand.

from enzyme.

If I recall correctly that is there for correctness around edge points.

Incidentally we may now have a better way of handling this now (the checkeddiv inst), which is conditionally activated upon a flag saying that is necessary.

As Manuel mentioned our derivative rules for things like this are simply specified in that file, which if you'd like to get into Enzyme dev, wouldn't be hard to open a PR and update!

from enzyme.

Also that test which you mention is actually the derivative generated by tapenade, another AD tool, which also has similar handling of these points.

from enzyme.

@Bike some of the reasoning behind this is outlined in this paper: https://proceedings.neurips.cc/paper/2020/file/4aaa76178f8567e05c8e8295c96171d8-Paper.pdf

from enzyme.

@tgymnich Thank you, I will read. I suspect that the considerations for ML are not so relevant for us (we're doing physics simulation) but I'm glad to know the ideas.

@ZuseZ4 @wsmoses Thanks for the tip. I patched this code to see how things would go without the zero test, and after optimization it ends up just as fast as my manual version. I assume Enzyme would not accept a PR to remove the test as that would result in bad behavior on the edge. As a more configurable solution for us, is it perhaps possible to manually specify the derivative of a function? Then we could have our code use my_sqrt, and tell Enzyme that the derivative of my_sqrt is the untested division.

from enzyme.

@Bike This functionality could be introduced behind some kind of flag. Given that the changes are not too invasive and some notion of completeness (what about FDiv and so on).

from enzyme.

Since the problem with sqrt is also due to division by zero, could the sqrt pattern just be changed to use CheckedDiv, rather than using the flag separately? There is a slight difference in that the existing sqrt pattern uses a ueq comparison rather than CheckedDiv's oeq - I don't know if that's important, depends on what a NaN input should do. I don't know how to enable or disable this EnzymeStrongZero flag, but I could try patching this and submitting a PR.

from enzyme.

@Bike you can enable the flag depending on the pass manager you use like so:

1. clang legacy new PM:
   -fpass-plugin=/opt/compiler-explorer/main/ClangEnzyme-XX.so -Xclang -load -Xclang /opt/compiler-explorer/main/ClangEnzyme-XX.so -mllvm -enzyme-strong-zero
2. clang new legacy PM:
   -Xclang -load -Xclang /opt/compiler-explorer/main/ClangEnzyme-XX.so -mllvm -enzyme-strong-zero

from enzyme.

Just a note that turning on trapping floating point is an important debugging tool so I'd prefer not to see fp exceptions for code that is considered unexceptional. I don't know if this case is, but the way it's being talked about up above suggests maybe.

from enzyme.

As a side note that you may be interested in @jedbrown a few weeks back we added support for running custom code on all the intermediate derivative values. For example, performing a nan-check (and then getting a backtrace to what triggered then nan).

from enzyme.

Cool, is that for forward or reverse or both?

from enzyme.

Both! Here's the (undocumented) Julia flag for enabling the nan checker as an example: https://github.com/EnzymeAD/Enzyme.jl/blob/02715a8bbac185342fe427f0090a8893d3a8af1d/src/compiler.jl#L5854

If you have other use cases and/or want to play with (from whatever input language, lmk)

from enzyme.

@tgymnich Oh, thanks for the tip. I'd love to be able to use the new pass manager. But I'm a bit confused - the FAQ linked in that other issue (https://enzyme.mit.edu/getting_started/Faq/#opt-cant-find--enzyme-option) says the opposite, that -fpass-plugin is the new one. But in any case it doesn't seem to work: clang++-15 -fpass-plugin=/path/to/my/LLVMEnzyme-15.so enzyme-test.cc -o enzyme-test gets me a linker error about undefined reference to `__enzyme_autodiff(void*, double)', so presumably Enzyme is not being run. (I can run things through opt manually but being able to do everything in one go would be nice.)

from enzyme.

@Bike uou should use ClangEnzyme not LLVMEnzyme for loading into clang

from enzyme.

@Bike my bad I mixed up the pass managers. It should be the other way around.

from enzyme.

Thank you, I got it working. The -enzyme-strong-zero flag with CheckedDiv does not seem to be working as I would expect, so I'll have to play around with it.

from enzyme.