sqrt behavior at 0

[Bike]

My group is getting started trying out Enzyme, and I noticed some unintuitive behavior on square roots: the derivative of sqrt at zero seems to be specially computed as zero. __enzyme_autodiff(my_sqrt, x) ends up with an fcmp and select to essentially do (x = 0.0) ? 0.0 : 1/(sqrt(2.0*x)), instead of what I expected, 1/(sqrt(2.0*x)).

This is, unless I'm missing something, mathematically incorrect - the derivative at zero should be undefined as the function is not continuous there (at least on reals), or positive infinity if you take the limit from above. But I'm less concerned about that and more about the decrease in efficiency from the branch. It seems to slow things around 40% in a very quick benchmark, versus the manual 1/sqrt(2.0*x). Here's the explorer showing a ucomisd -> je sequence.

I am not sure that this is a bug since it appears to be a deliberate design decision, with similar logic apparent e.g. in this benchmark. But I am at least interested in knowing why it was designed this way, and whether there is any way to configure this away. I couldn't find any mention of this behavior on the website, the arxiv preprint, or the code. We would like to use Enzyme rather than the symbolic differentiation we are doing now but this is a noticeable decrease in performance for us.

[ZuseZ4]

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:

Enzyme/enzyme/Enzyme/InstructionDerivatives.td

Line 785 in 3c0014a

def : IntrPattern<(Op $x),

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.

[wsmoses]

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!

[wsmoses]

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

[tgymnich]

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

[Bike]

@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.

[tgymnich]

@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).

[wsmoses]

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

[Bike]

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.

[tgymnich]

@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

[jedbrown]

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.

[wsmoses]

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).

[jedbrown]

Cool, is that for forward or reverse or both?