Enzyme compilation failed due to illegal type analysis - in a case that supposed to work.

[GiggleLiu]

When I try to differentiate OMEinsum, the following error were thrown. The code is shown bellow, where the variable code is a mutable, callable object. Since it does not carry gradient, I suppose it should work. But it does not.

using Enzyme.EnzymeRules, OMEinsum, Enzyme
using OMEinsum: get_size_dict!

function EnzymeRules.augmented_primal(
        config::EnzymeRules.RevConfigWidth{1},
        func::Const{typeof(einsum!)}, ::Type, 
        code::Const, xs::Duplicated, ys::Duplicated, sx::Const, sy::Const, size_dict::Const)
    @assert sx.val == 1 && sy.val == 0 "Only α = 1 and β = 0 is supported, got: $sx, $sy"
    # Compute primal
    if EnzymeRules.needs_primal(config)
        primal = func.val(code.val, xs.val, ys.val, sx.val, sy.val, size_dict.val)
    else
        primal = nothing
    end
    # Save x in tape if x will be overwritten
    if EnzymeRules.overwritten(config)[3]
        tape = copy(xs.val)
    else
        tape = nothing
    end
    shadow = ys.dval
    return EnzymeRules.AugmentedReturn(primal, shadow, tape)
end

function EnzymeRules.reverse(config::EnzymeRules.RevConfigWidth{1},
               func::Const{typeof(einsum!)}, dret::Type{<:Annotation}, tape,
               code::Const,
               xs::Duplicated, ys::Duplicated, sx::Const, sy::Const, size_dict::Const)
   xval = EnzymeRules.overwritten(config)[3] ? tape : xs.val
   for i=1:length(xs.val)
       xs.dval[i] .+= OMEinsum.einsum_grad(OMEinsum.getixs(code.val),
             xval, OMEinsum.getiy(code.val), size_dict.val, conj(ys.dval), i)
   end
   return (nothing, nothing, nothing, nothing, nothing, nothing)
end

using Test
@testset "EnzymeExt bp check 2" begin
    A, B, C = randn(2, 3), randn(3, 4), randn(4, 2)
    code = optimize_code(ein"ij, jk, ki->", uniformsize(ein"ij, jk, ki->", 2), TreeSA())
    cost0 = code(A, B, C)[]
    gA = zero(A); gB = zero(B); gC = zero(C);
    f(code, a, b, c) = code(a, b, c)[]
    Enzyme.autodiff(set_runtime_activity(Reverse), f, Active, Const(code), Duplicated(A, gA), Duplicated(B, gB), Duplicated(C, gC))
    cost, mg = OMEinsum.cost_and_gradient(code, (A, B, C))
    @test cost[] ≈ cost0
    @test all(gA .≈ mg[1])
    @test all(gB .≈ mg[2])
    @test all(gC .≈ mg[3])
end

The error message is:

EnzymeExt bp check 2: Error During Test at REPL[6]:1
  Got exception outside of a @test
  Enzyme compilation failed due to illegal type analysis.
  Current scope: 
  ; Function Attrs: mustprogress willreturn
  define internal fastcc nonnull "enzyme_type"="{}" {} addrspace(10)* @preprocess_julia_collect_to__54382({} addrspace(10)* noundef nonnull align 8 dereferenceable(24) "enzyme_inactive" "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,8]:Pointer, [-1,8,0]:Integer, [-1,8,1]:Integer, [-1,8,2]:Integer, [-1,8,3]:Integer, [-1,8,4]:Integer, [-1,8,5]:Integer, [-1,8,6]:Integer, [-1,8,7]:Integer, [-1,8,8]:Pointer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer}" "enzymejl_parmtype"="6204202512" "enzymejl_parmtype_ref"="2" %0, { [1 x {} addrspace(10)*], {} addrspace(10)* } addrspace(11)* nocapture noundef nonnull readonly align 8 dereferenceable(16) "enzyme_inactive" "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Pointer, [-1,0,0,0]:Integer, [-1,0,0,1]:Integer, [-1,0,0,2]:Integer, [-1,0,0,3]:Integer, [-1,0,0,4]:Integer, [-1,0,0,5]:Integer, [-1,0,0,6]:Integer, [-1,0,0,7]:Integer, [-1,0,0,8]:Pointer, [-1,0,8]:Pointer, [-1,0,8,0]:Integer, [-1,0,8,1]:Integer, [-1,0,8,2]:Integer, [-1,0,8,3]:Integer, [-1,0,8,4]:Integer, [-1,0,8,5]:Integer, [-1,0,8,6]:Integer, [-1,0,8,7]:Integer, [-1,0,8,8]:Pointer, [-1,0,8,8,-1]:Integer, [-1,0,16]:Pointer, [-1,0,16,0]:Integer, [-1,0,16,1]:Integer, [-1,0,16,2]:Integer, [-1,0,16,3]:Integer, [-1,0,16,4]:Integer, [-1,0,16,5]:Integer, [-1,0,16,6]:Integer, [-1,0,16,7]:Integer, [-1,0,16,8]:Pointer, [-1,0,16,8,-1]:Integer, [-1,0,24]:Integer, [-1,0,25]:Integer, [-1,0,26]:Integer, [-1,0,27]:Integer, [-1,0,28]:Integer, [-1,0,29]:I
...
    %iv.next = add nuw nsw i64 %iv, 1, !dbg !118: {[-1]:Integer}, intvals: {1,}
    %iv.next5 = add nuw nsw i64 %iv4, 1, !dbg !201: {[-1]:Integer}, intvals: {1,}
    %138 = add i64 %iv, 2, !dbg !203: {}, intvals: {2,}
    %87 = add i64 %138, %iv4, !dbg !201: {[-1]:Integer}, intvals: {2,}
    %139 = add i64 %iv, 3, !dbg !203: {}, intvals: {3,}
    %86 = add i64 %139, %iv4, !dbg !201: {[-1]:Integer}, intvals: {3,}
    %140 = add i64 %iv, 2, !dbg !203: {}, intvals: {2,}
    %85 = add i64 %140, %iv4, !dbg !201: {[-1]:Integer}, intvals: {2,}
  </analysis>
  
  Illegal updateAnalysis prev:{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Integer, [-1,8]:Pointer, [-1,8,0]:Integer, [-1,8,1]:Integer, [-1,8,2]:Integer, [-1,8,3]:Integer, [-1,8,4]:Integer, [-1,8,5]:Integer, [-1,8,6]:Integer, [-1,8,7]:Integer, [-1,8,8]:Pointer, [-1,8,8,-1]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer} new: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Pointer, [-1,0,8]:Pointer, [-1,8]:Pointer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer}
  val: {} addrspace(10)* %0 origin=  %common.ret.op = phi {} addrspace(10)* [ %0, %top ], [ %61, %guard_exit ], [ %61, %common.ret.loopexit ], [ %0, %common.ret.loopexit3 ]
  
  Stacktrace:
    [1] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder})
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:1576
    [2] EnzymeCreateAugmentedPrimal(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{Enzyme.API.CDIFFE_TYPE}, TA::Enzyme.TypeAnalysis, returnUsed::Bool, shadowReturnUsed::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{Bool}, forceAnonymousTape::Bool, runtimeActivity::Bool, width::Int64, atomicAdd::Bool)
      @ Enzyme.API ~/.julia/packages/Enzyme/RTS5U/src/api.jl:389
    [3] enzyme!(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{4, Bool}, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{Int64}, boxedArgs::Set{Int64})
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:4095
    [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7338
    [5] codegen
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:6146 [inlined]
    [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8468
    [7] _thunk
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8468 [inlined]
    [8] cached_compilation
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8509 [inlined]
    [9] thunkbase(ctx::LLVM.Context, mi::Core.MethodInstance, ::Val{0x0000000000006860}, ::Type{Const{typeof(einsum)}}, ::Type{Duplicated{Any}}, tt::Type{Tuple{Duplicated{SlicedEinsum{Char, DynamicNestedEinsum{Char}}}, Duplicated{Tuple{Matrix{Float64}, Matrix{Float64}, Matrix{Float64}}}, Const{Dict{Char, Int64}}}}, ::Val{Enzyme.API.DEM_ReverseModePrimal}, ::Val{1}, ::Val{(true, true, true, true)}, ::Val{true}, ::Val{false}, ::Type{FFIABI}, ::Val{false}, ::Val{true})
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8641
   [10] #s2105#19135
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8778 [inlined]
   [11] var"#s2105#19135"(FA::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, ReturnPrimal::Any, ShadowInit::Any, World::Any, ABI::Any, ErrIfFuncWritten::Any, RuntimeActivity::Any, ::Any, ::Any, ::Any, ::Any, tt::Any, ::Any, ::Any, ::Any, ::Any, ::Any, ::Any, ::Any, ::Any)
      @ Enzyme.Compiler ./none:0
   [12] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
      @ Core ./boot.jl:707
   [13] runtime_generic_augfwd(activity::Type{Val{(false, true, true, false)}}, runtimeActivity::Val{true}, width::Val{1}, ModifiedBetween::Val{(true, true, true, true)}, RT::Val{@NamedTuple{1, 2, 3}}, f::typeof(einsum), df::Nothing, primal_1::SlicedEinsum{Char, DynamicNestedEinsum{Char}}, shadow_1_1::SlicedEinsum{Char, DynamicNestedEinsum{Char}}, primal_2::Tuple{Matrix{Float64}, Matrix{Float64}, Matrix{Float64}}, shadow_2_1::Tuple{Matrix{Float64}, Matrix{Float64}, Matrix{Float64}}, primal_3::Dict{Char, Int64}, shadow_3_1::Nothing)
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/RTS5U/src/rules/jitrules.jl:469
   [14] #_#155
      @ ~/.julia/dev/OMEinsum/src/slicing.jl:109 [inlined]
   [15] SlicedEinsum
      @ ~/.julia/dev/OMEinsum/src/slicing.jl:104 [inlined]
   [16] f
      @ ./REPL[6]:6 [inlined]
   [17] f
      @ ./REPL[6]:0 [inlined]
   [18] augmented_julia_f_16439_inner_1wrap
      @ ./REPL[6]:0
   [19] macro expansion
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8398 [inlined]
   [20] enzyme_call
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7950 [inlined]
   [21] AugmentedForwardThunk
      @ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7787 [inlined]
   [22] autodiff
      @ ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:384 [inlined]
   [23] autodiff(::ReverseMode{false, true, FFIABI, false, false}, ::var"#f#1", ::Type{Active}, ::Const{SlicedEinsum{Char, DynamicNestedEinsum{Char}}}, ::Duplicated{Matrix{Float64}}, ::Duplicated{Matrix{Float64}}, ::Duplicated{Matrix{Float64}})
      @ Enzyme ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:512
   [24] macro expansion
      @ REPL[6]:7 [inlined]
   [25] macro expansion
      @ ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/Test/src/Test.jl:1700 [inlined]
   [26] top-level scope
      @ REPL[6]:2
   [27] eval
      @ ./boot.jl:430 [inlined]
   [28] eval_user_input(ast::Any, backend::REPL.REPLBackend, mod::Module)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/REPL/src/REPL.jl:245
   [29] repl_backend_loop(backend::REPL.REPLBackend, get_module::Function)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/REPL/src/REPL.jl:342
   [30] start_repl_backend(backend::REPL.REPLBackend, consumer::Any; get_module::Function)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/REPL/src/REPL.jl:327
   [31] run_repl(repl::REPL.AbstractREPL, consumer::Any; backend_on_current_task::Bool, backend::Any)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/REPL/src/REPL.jl:483
   [32] run_repl(repl::REPL.AbstractREPL, consumer::Any)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/REPL/src/REPL.jl:469
   [33] (::Base.var"#1139#1141"{Bool, Symbol, Bool})(REPL::Module)
      @ Base ./client.jl:446
   [34] #invokelatest#2
      @ ./essentials.jl:1055 [inlined]
   [35] invokelatest
      @ ./essentials.jl:1052 [inlined]
   [36] run_main_repl(interactive::Bool, quiet::Bool, banner::Symbol, history_file::Bool, color_set::Bool)
      @ Base ./client.jl:430
   [37] repl_main
      @ ./client.jl:567 [inlined]
   [38] _start()
      @ Base ./client.jl:541
Test Summary:        | Error  Total   Time
EnzymeExt bp check 2 |     1      1  35.8s
ERROR: Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.

[GiggleLiu]

If anyone can help me get this PR work, I will be greatful! under-Peter/OMEinsum.jl#176

[wsmoses]

can you paste the whole log?

[GiggleLiu]

can you paste the whole log?

Good point. I tried, but the log exceeded the length limit that GitHub allow. Is there any way that I can circumvent this limit? Or which part do you want to see?

[wsmoses]

you can upload a file?

[GiggleLiu]

It works. Please check the following file:

error.txt