Bad interaction with Enzyme?

[pagnani]

Hi there!
Cross linking here the original issue on Enzyme following the advice of @wsmoses.

Repeating here:

I experience a problem with Enzyme interacting with Tullio. The MWE is the computation of the trace of the product with two matrices $tr(A,B)$ . I wrote a method with loops and one using Tullio.

function provatullio(A,B)
    @tullio mytrace := A[i,j]*B[j,i] 
    return mytrace
end

function provaloop(A,B)
    tr = zero(eltype(A))
    for i in axes(A,1)
        for j in axes(A,2)
            tr += A[i,j]*B[j,i]
        end
    end
    tr
en

Then I constructed a wrapper to compute the gradient of the two methods

function gradenzyme(f,A,B)
    fA(x)=f(x,B)
    fB(x)=f(A,x)
    return Enzyme.gradient(Reverse, fA, A), Enzyme.gradient(Reverse, fB, B)
end

Test:

julia> n = 2; A=rand(n,n); B = rand(n,n);

julia> gradenzyme(provaloop,A,B)
([0.9338361114622028 0.9780719543182385; 0.7051893579184637 0.5412218303659214], [0.5131207164626734 0.5299631790389514; 0.0956376425423926 0.7901845832722304])

julia> gradenzyme(provatullio,A,B)

As you can see, the loopy version is ok. The Tullio's one fails with this scary stacktrace

Thanks a lot

[pagnani]

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

[mcabbott]

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

[pagnani]

Pinging @wsmoses which is the expert here ...

[mcabbott]

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?

[wsmoses]

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

[mcabbott]

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.