EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
439 stars 62 forks source link

Internal handler hits `AssertionError: x.indices == dx.indices Stacktrace:` #1781

Open vchuravy opened 1 week ago

vchuravy commented 1 week ago
ERROR: AssertionError: x.indices == dx.indices
Stacktrace:
  [1] Duplicated
    @ ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:69 [inlined]
  [2] Duplicated
    @ ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:68 [inlined]
  [3] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(getindex), df::Nothing, primal_1::SubArray{…}, shadow_1_1::SubArray{…}, primal_2::Colon, shadow_2_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Tb3Iu/src/rules/jitrules.jl:149
  [4] Res!
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:51 [inlined]
  [5] Res!
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:0 [inlined]
  [6] diffejulia_Res__8056_inner_1wrap
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:0
  [7] macro expansion
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:7151 [inlined]
  [8] enzyme_call
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:6760 [inlined]
  [9] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/compiler.jl:6637 [inlined]
 [10] autodiff
    @ ~/.julia/packages/Enzyme/Tb3Iu/src/Enzyme.jl:320 [inlined]
 [11] autodiff(::ReverseMode{…}, ::typeof(Res!), ::Type{…}, ::Duplicated{…}, ::Duplicated{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…})
    @ Enzyme ~/.julia/packages/Enzyme/Tb3Iu/src/Enzyme.jl:332
 [12] top-level scope
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:102
Some type information was truncated. Use `show(err)` to see complete types.

@wsmoses said:

That’s an intentional error message telling you that you have a duplicated of a view with different offset for primal and shadow

But the call to Enzyme is:

autodiff(Enzyme.Reverse, Res!, Const, Duplicated(Rvec, dR), Duplicated(Cvec, dC), Const(C_bc), Const(Χ), Const(arch), Const(grid), Const(launch))

and

julia> typeof(Rvec)
Vector{Float64} (alias for Array{Float64, 1})

julia> typeof(dR)
Vector{Float64} (alias for Array{Float64, 1})

julia> typeof(dC)
Vector{Float64} (alias for Array{Float64, 1})

julia> typeof(Cvec)
Vector{Float64} (alias for Array{Float64, 1})
# example_2.jl # # We solve the 2D elliptic PDE: # ∂( Χ * ∂C/∂x )/∂x + ∂( Χ * ∂C/∂z )/∂z = 0 # with boundary conditions: # C(0) = 10, C(L) = 100 # and assume that Χ is constant using LinearAlgebra, ForwardDiff, SparseArrays using KernelAbstractions using Chmy, Chmy.Architectures, Chmy.Grids, Chmy.Fields, Chmy.BoundaryConditions, Chmy.GridOperators, Chmy.KernelLaunch using GLMakie # gradients, q.x = -Χ*∂/∂x(C), q.y = -Χ*∂/∂y(C) @kernel inbounds=true function compute_q!(q, C, Χ, g::StructuredGrid, O) I = @index(Global, NTuple) I = I + O q.x[I...] = -Χ * ∂x(C, g, I...) q.y[I...] = -Χ * ∂y(C, g, I...) end # ∂/∂x(qx) + ∂/∂y(qy) @kernel inbounds=true function update_C!(C, q, g::StructuredGrid, O) I = @index(Global, NTuple) I = I + O C[I...] = divg(q, g, I...) end # residual routine with fields function Res!(R::Field, C::Field, C_bc::Tuple, q::NamedTuple, Χ::Float64, arch::SingleDeviceArchitecture, grid::StructuredGrid, launch::Launcher) bc!(arch, grid, C_bc) launch(arch, grid, compute_q! => (q,C,Χ, grid)) launch(arch, grid, update_C! => (R,q, grid)) return nothing end function Res!(Rvec::Vector, Cvec::Vector, C_bc, Χ, arch, grid, launch) # this is certainly inefficient, we should allocate this earlier... C = Field(arch, grid, Center(), eltype(Cvec)) R = Field(arch, grid, Center(), eltype(Rvec)) # residual q = VectorField(arch,grid, eltype(Cvec)) # conductivity set!(C, Cvec) Res!(R, C, C_bc, q, Χ, arch, grid, launch) Rvec .= interior(R)[:] return nothing end # function Enz_Res!(R, C, C_bc, q, Χ, arch, grid, launch) function Res1!(Rvec::Vector, Cvec::Vector, C_bc, C, R, q, Χ, arch, grid, launch) set!(C, Cvec) Res!(R, C, C_bc, q, Χ, arch, grid, launch) Rvec .= interior(R)[:] return nothing end # Setup backend = CPU() arch = Arch(backend) grid = UniformGrid(arch; origin=(0, 0), extent=(2,2), dims=(11,12)) launch = Launcher(arch, grid) # Arrays C = Field(arch, grid, Center()) R = Field(arch, grid, Center()) # residual q = VectorField(arch,grid) # conductivity Cd = Field(arch, grid, Center(), ForwardDiff.Dual) Rd = Field(arch, grid, Center(), ForwardDiff.Dual) qd = VectorField(arch,grid, ForwardDiff.Dual) P = Field(arch, grid, Center()) # set BC (note that you need to use Cd): C_bc = batch(grid, Cd => (y=(Dirichlet(ForwardDiff.Dual(9.0)), Dirichlet(ForwardDiff.Dual(100.0))),x=Neumann())) Χ = 1.0 # residual routine in vector format Rvec = interior(R)[:] Cvec = interior(C)[:] # "normal" residual routine Res!(Rvec, Cvec, C_bc, Χ, arch, grid, launch) dC = zero(Cvec) dR = zero(Rvec) using Enzyme Enzyme.API.runtimeActivity!(true) autodiff(Enzyme.Reverse, Res!, Const, Duplicated(Rvec, dR), Duplicated(Cvec, dC), Const(C_bc), Const(Χ), Const(arch), Const(grid), Const(launch)) #autodiff(Enzyme.Forward, Res!, Const, Duplicated(Rvec, dR), Duplicated(Cvec, dC), Const(C_bc), Const(Χ), Const(arch), Const(grid), Const(launch)) ``` (chmy_ad) pkg> st Status `~/src/snippets/chmy_ad/Project.toml` [33a72cf0] Chmy v0.1.19 `~/src/Chmy` [7da242da] Enzyme v0.12.35 [f6369f11] ForwardDiff v0.10.36 [e9467ef8] GLMakie v0.10.9 [63c18a36] KernelAbstractions v0.9.25 [37e2e46d] LinearAlgebra [2f01184e] SparseArrays v1.10.0 ```
vchuravy commented 1 week ago

After #1780 I now hit this also in ForwardMode

ERROR: AssertionError: x.indices == dx.indices
Stacktrace:
  [1] Duplicated
    @ ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:69 [inlined]
  [2] Duplicated
    @ ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:68 [inlined]
  [3] runtime_generic_fwd(activity::Type{…}, width::Val{…}, RT::Val{…}, f::typeof(getindex), df::Nothing, primal_1::SubArray{…}, shadow_1_1::SubArray{…}, primal_2::Colon, shadow_2_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:135
  [4] Res!
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:51 [inlined]
  [5] Res!
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:0 [inlined]
  [6] fwddiffejulia_Res__19223_inner_1wrap
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:0
  [7] macro expansion
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
  [8] enzyme_call
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
  [9] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6674 [inlined]
 [10] autodiff
    @ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:437 [inlined]
 [11] autodiff(::ForwardMode{…}, ::typeof(Res!), ::Type{…}, ::Duplicated{…}, ::Duplicated{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…})
    @ Enzyme ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:332
 [12] top-level scope
    @ ~/src/snippets/chmy_ad/example_2_ForwardDiff.jl:103
Some type information was truncated. Use `show(err)` to see complete types.