PreallocationTools.jl icon indicating copy to clipboard operation
PreallocationTools.jl copied to clipboard

Add rules for Enzyme

Open vchuravy opened this issue 5 months ago • 3 comments

What feature report was addressed?

We are evaluating the use of PreallocationTools for Trixi, but need to ensure that adding it does not break Enzyme support.

Checklist

  • [x] Appropriate tests were added
  • [x] The new feature was added in a way that does not break public API
  • [] New documentation related to the new feature was added
  • [x] The new code follows the contributor guidelines, in particular the SciML Style Guide and COLPRAC.

Additional context

Enzyme requires the shadow heap values to have the same type and shape as the primals. Therefore dual_du can actually be smaller than with ForwardDiff.

Supporting Enzyme batch mode is only possible on v1.11 or by allocating memory / doing unsafe operations.

My biggest worry would be user who attempt to use Enzyme over ForwardDiff. Then in the case of Const(stod) Enzyme and ForwardDiff would compete for the same memory. The case of Duplicated(stod, Enzyme.make_zero(stod)) does the right thing in this instance.

vchuravy avatar Aug 31 '25 16:08 vchuravy