search

WaterLily-jl / WaterLily.jl

Fast and simple fluid simulator in Julia
Other
616 stars 75 forks source link

Using DifferentiationInterface #144

Open gdalle opened 1 month ago

gdalle commented 1 month ago

Hi @b-fg! Just opening an issue to chat if you have any more questions or bugs related to DifferentiationInterface.jl. I didn't really understand from https://github.com/gdalle/DifferentiationInterface.jl/issues/355 how you want to use it: outside of WaterLily to try differentiating through it, or inside WaterLily at the few spots where you currently use ForwardDiff?

b-fg commented 1 month ago

Hi @gdalle! I appreciate the follow up :) We were planning on using DI.jl to differentiate through WaterLily. For example, users typically define a (parametrized) geometry and then compute flow-induced metrics such as lift or drag forces. To compute the derivative of "d(lift)/d(geometry)" we must differentiate through the whole solver.

Using ForwardDiff.jl directly, or through DI.jl, this currently works on CPU backends. But to be able to differentiate on the GPU backend we need Enzyme.jl. As you suggested, I opened an issue there since this is currently not working. @vchuravy also opened and issue related to this some time ago. https://github.com/EnzymeAD/Enzyme.jl/issues/1547 https://github.com/EnzymeAD/Enzyme.jl/issues/1647

Once this works as expected, it would be neat to switch to DI.jl because it's a more friendly interface for users. So that's the plan!

vchuravy commented 1 month ago

I don't think you won't get around using Enzyme internally, but you might be able to define a ChainRule to define something that DI can access.

This is what we did for https://github.com/Argonne-National-Laboratory/Checkpointing.jl

Besides GPU the core challenge is mutation support, which would require writing a CFD simulation in a very different style.