problem with tullio

[pagnani]

Hi there,

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
end

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)

The tullio version does not work. Below the very long stacktrace that I had to cut because could not fit the limit size of issues. A t the end my versioninfo and my package status).

julia> gradenzyme(provatullio,A,B)
ERROR: Enzyme compilation failed due to illegal type analysis.
Current scope: 
; Function Attrs: mustprogress willreturn
define double @preprocess_julia_Eval_31823({} addrspace(10)* noundef nonnull align 16 dereferenceable(40) %0, {} addrspace(10)* noundef nonnull align 16 dereferenceable(40) %1) local_unnamed_addr #72 !dbg !2617 {
[....]
Stacktrace:
  [1] julia_error(cstr::Cstring, val::Ptr{…}, errtype::Enzyme.API.ErrorType, data::Ptr{…}, data2::Ptr{…}, B::Ptr{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:1696
  [2] EnzymeCreateAugmentedPrimal(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnUsed::Bool, shadowReturnUsed::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, forceAnonymousTape::Bool, width::Int64, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/KJgKj/src/api.jl:177
  [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::Tuple{…}, returnPrimal::Bool, jlrules::Vector{…}, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:3093
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:4767
  [5] codegen
    @ ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:4348 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool) (repeats 2 times)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:5362
  [7] cached_compilation
    @ ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:5396 [inlined]
  [8] (::Enzyme.Compiler.var"#509#510"{…})(ctx::LLVM.Context)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:5462
  [9] JuliaContext(f::Enzyme.Compiler.var"#509#510"{…})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/U36Ed/src/driver.jl:47
 [10] #s1056#508
    @ ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:5414 [inlined]
 [11] 
    @ Enzyme.Compiler ./none:0
 [12] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
    @ Core ./boot.jl:602
 [13] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::Tullio.Eval{…}, df::Nothing, primal_1::Matrix{…}, shadow_1_1::Matrix{…}, primal_2::Matrix{…}, shadow_2_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/KJgKj/src/rules/jitrules.jl:175
 [14] macro expansion
    @ ~/.julia/packages/Tullio/vChiX/src/macro.jl:976 [inlined]
 [15] provatullio
    @ ~/SCRA/TEST_ENZYME/prova.jl:5 [inlined]
 [16] fA
    @ ~/SCRA/TEST_ENZYME/prova.jl:44 [inlined]
 [17] fA
    @ ~/SCRA/TEST_ENZYME/prova.jl:0 [inlined]
 [18] diffejulia_fA_30873_inner_1wrap
    @ ~/SCRA/TEST_ENZYME/prova.jl:0
 [19] macro expansion
    @ ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:5310 [inlined]
 [20] enzyme_call
    @ ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:4988 [inlined]
 [21] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/KJgKj/src/compiler.jl:4930 [inlined]
 [22] autodiff
    @ ~/.julia/packages/Enzyme/KJgKj/src/Enzyme.jl:215 [inlined]
 [23] autodiff
    @ ~/.julia/packages/Enzyme/KJgKj/src/Enzyme.jl:224 [inlined]
 [24] gradient
    @ ~/.julia/packages/Enzyme/KJgKj/src/Enzyme.jl:805 [inlined]
 [25] gradenzyme(f::typeof(provatullio), A::Matrix{Float64}, B::Matrix{Float64})
    @ Main ~/SCRA/TEST_ENZYME/prova.jl:46
 [26] top-level scope
    @ REPL[86]:1
Some type information was truncated. Use `show(err)` to see complete types.

Package 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

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/

[wsmoses]

Since I know tulio does odd things under the hood and has custom derivative rules defined for other AD tools, I think the correct solution here is to define custom enzyme rules within Tulio.jl.

Open an issue on Tulio.jl and cc me?

[wsmoses]

However @pagnani what is your julia/os/enzyme version.

On my system (julia 1.10, ubuntu, latest Enzyme), this succeeds:

julia> gradenzyme(provaloop,A,B)
([0.23380880969329965 0.5227475047630739; 0.48402478399611826 0.48401775694315896], [0.3002268423233645 0.9257902665611332; 0.6552753578836962 0.1012109582137013])

[wsmoses]

never mind, you already uploaded that. Can you upload the full stack trace?

[pagnani]

never mind, you already uploaded that. Can you upload the full stack trace?

Hi @wsmoses thanks for the timely reply.

Here the complete (and scary) stacktrace.

However @pagnani what is your julia/os/enzyme version.

On my system (julia 1.10, ubuntu, latest Enzyme), this succeeds:

julia> gradenzyme(provaloop,A,B)
([0.23380880969329965 0.5227475047630739; 0.48402478399611826 0.48401775694315896], [0.3002268423233645 0.9257902665611332; 0.6552753578836962 0.1012109582137013])

Yes, the loopy version, as in my original issue, works for me too, is the Tullio version which is not working. The original problem is in a much more complicated loss function, but I found out that also this minimal case is not working.

Il will also open an issue on Tullio and cross-link it here later.

Also, unrelated to this issue, is this way of computing the gradient over two sets of variables ( i.e. by making two closures over A, and B) the correct thing to do?

Zygote allows to make lambdas over variable as

function gradzygote(f,A,B)
    Zygote.gradient((x,y)->f(x,y), A,B)
end

but I did not manage to find an equivalent in Enzyme.

Thanks for your work!

[pagnani]

Added this issue on Tullio's repo.

A

[wsmoses]

@pagnani this code succeeds for me on present Enzyme.

julia> gradenzyme(provaloop,A,B)
([0.5663164027331671 0.3256755318530826; 0.6901298018433815 0.2663614388591653], [0.5089800932986028 0.9230120823778869; 0.6783049935954609 0.13943622375552933])

julia> gradenzyme(provatullio,A,B)

┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/kqxyC/src/utils.jl:59
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/kqxyC/src/utils.jl:59
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/kqxyC/src/utils.jl:59
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/kqxyC/src/utils.jl:59
([0.5663164027331671 0.3256755318530826; 0.6901298018433815 0.2663614388591653], [0.5089800932986028 0.9230120823778869; 0.6783049935954609 0.13943622375552933])

It is still likely desirable that Tullio add EnzymeRules, but in any case everything works otherwise.

Closing, please reopen if it persists.