Bad interaction with Enzyme? about tullio.jl HOT 6 OPEN

Forgot to add

versioninfo

julia> versioninfo()
Julia Version 1.10.0
Commit 3120989f39b (2023-12-25 18:01 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (x86_64-apple-darwin22.4.0)
  CPU: 12 × Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
  Threads: 11 on 12 virtual cores
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 8
  LD_LIBRARY_PATH = /Users/pagnani/.julia/conda/3/lib/:/opt/local/lib/mariadb/mysql/

and pkg status

TEST_ENZYME) pkg> st
Status `~/SCRA/TEST_ENZYME/Project.toml`
  [7da242da] Enzyme v0.11.14
  [587475ba] Flux v0.14.11
  [f6369f11] ForwardDiff v0.10.36
  [bdcacae8] LoopVectorization v0.12.166
  [37e2e3b7] ReverseDiff v1.15.1
  [bc48ee85] Tullio v0.3.7
  [e88e6eb3] Zygote v0.6.69

from tullio.jl.

I don't know what to make of this error. The recursive code for scalar reductions is here, and perhaps something about this looks problematic? E.g. the number of recursions depends on an Int, from Threads.nthreads().

https://github.com/mcabbott/Tullio.jl/blob/master/src/threads.jl#L217-L256

from tullio.jl.

Pinging @wsmoses which is the expert here ...

from tullio.jl.

Reading the linked issue EnzymeAD/Enzyme.jl#1279 I see that perhaps the suggestion is that Tullio should have Enzyme gradient rules -- as it currently does for Zygote & Tracker.

That would be fine, someone just has to write them! I was defeated when last I tried to write Enzyme rules (for some NNlib funcitons), but maybe docs & examples have improved since then. Comments:

• The existing rules for Zygote allocate a new array for each gradient. I believe Enzyme usually allocates these up front, so it would be more efficient to write into those.
• At present the gradient computation writes into all the arrays at once. So a forward pass which has one loop nest reading 3 arrays & writing 1, gets a gradient which still has one loop nest, and writes into 3. This turns out to be sub-optimal for performance, but would be quite a big job to change. I presume that, if Enzyme handled things without custom rules, it would work the same way -- one loop nest?

from tullio.jl.

Folliwng up here since I closed the linked PR (it now works without a rule on present Enzyme). However I do think a tullio rule here is appropriate.

Happy to help however I can.

The rough semantics required for rules are:

Forward Mode:
x[i,j,k] = f(y, z, ...)
dx[i,j,k] = df(y, dy, z, dz, ...)

Enzyme will provide x, dx, y, dy, etc but I dont know how to do those ops in tullio.

For reverse mode

x[i,j,k] = f(y[u,v], z, ...)

dy[u,v] += df(dx[i,j,k], y, z, ...)
dz[u,v] += df(dx[i,j,k], y, z, ...)
dx[i,j,k] = 0

Same thing here

from tullio.jl.

Many thanks for taking a look.

I'm really swamped but I believe that adding a simple reverse-mode rule would be quite simple. Tullio already writes functions for computing all the reverse gradients, and uses them for Zygote etc.

Adding similar functions for forward mode would not be very hard, but would require some new transformations.

Fully exploiting what Enzyme knows about activity would be ideal, but is unlikely. This package always writes all code at macro-expansion time, so it would need to speculatively write code for several possibilities. In some cases, knowing that a particular array is const would change what loops can be parallelised in the gradient. But Tullio really isn't set up to think about multiple such possibilities at once.

from tullio.jl.