Enzyme modifies entries of an array in an inactive struct

jlk9 commented 3 days ago

This is a continuation of issue 485 in KernelAbstractions.jl, which now appears to be more of an Enzyme issue.

Given this kernel called on CPU and function wrapper, where the mesh of our model should be an inactive type:

# Setting meshes to inactive types:
Enzyme.EnzymeRules.inactive_type(::Type{T} where T <:HorzMesh) = true

@kernel function GradientOnEdgeModified(@Const(dcEdge), GradEdge)
    # global indices over nEdges
    iEdge, k = @index(Global, NTuple)

    @inbounds GradEdge[k, iEdge] = GradEdge[k, iEdge] / dcEdge[iEdge]

    @synchronize()
end

function gradient_normSq(grad, mesh::HorzMesh; backend=KA.CPU())

    nEdges = size(grad)[2]
    vert_levels = 1

    kernel! = GradientOnEdgeModified(backend)
    kernel!(mesh.Edges.dcEdge, grad, workgroupsize=64, ndrange=(nEdges, vert_levels))

    KA.synchronize(backend)

    #@show grad

    normSq = 0.0
    for i = 1:nEdges
        normSq += grad[i]^2
    end

    return normSq
end

we stow away the pre-AD mesh: old_mesh = deepcopy(mesh). Then we call autodiff on a given set of gradient values and Horzmesh object:

d_normSq = autodiff(Enzyme.Reverse,
                    gradient_normSq,
                    Duplicated(gradNum, d_gradNum),
                    Duplicated(mesh, d_mesh))

and see that one of the arrays of the mesh is changed:

isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge) = false

mesh.Edges.dcEdge is used in our gradient kernel but not modified, so it should not be changed with a call to autodiff. But two seemingly unrelated changes to the code eliminate the bug:

Adding the @show grad statement to our gradient norm function.
Commenting out the inactive type rule for mesh objects.

This can be replicated with this code: https://github.com/jlk9/MPAS-Ocean.jl/tree/bug-reduce. I've tried reducing further but had difficulties with the bug disappearing. I'll update if further reductions or observations are done.

wsmoses commented 3 days ago

So just because something is inactive doesn't mean its immutable. you could have other unrelated computations modified inplace, which aren't differentiated and Enzyme would also run those same updates as in the orignial function

wsmoses commented 3 days ago

Modifying your code to test this, I don't believe this is a bug per above:

using MOKA: HorzMesh, ReadHorzMesh
using KernelAbstractions
using Enzyme

Enzyme.API.printall!(true)

import Downloads
import KernelAbstractions as KA

backend = KA.CPU()
#=
"""
    HorzMesh

A struct, comprised of SoA, describing a 2-D TRiSK mesh
"""
struct HorzMesh{ET}
    Edges::ET
end

# these are line segments
@kwdef struct Edges{I,FV} 
    dcEdge::FV
    nEdges::I
end

edgeArray = KA.zeros(backend, Float64, 2306)
for i = 1:2306
    edgeArray[i] = i + rand()
end

edges = Edges(edgeArray, 2306)
mesh = HorzMesh(edges)
=#
# Setting meshes to inactive types:
Enzyme.EnzymeRules.inactive_type(::Type{<:HorzMesh}) = true

@kernel function GradientOnEdgeModified(@Const(dcEdge), GradEdge)
    # global indices over nEdges
    iEdge, k = @index(Global, NTuple)

    @inbounds GradEdge[k, iEdge] = GradEdge[k, iEdge] / dcEdge[iEdge]

    @synchronize()
end

# As a cleaner / easier to read test, let's create an outer function that measures the norm of the gradient computed by kernel:
function gradient_normSq(grad, mesh::HorzMesh; backend=KA.CPU())

    nEdges = size(grad)[2]
    vert_levels = 1

    # New modified kernel:
    kernel! = GradientOnEdgeModified(backend)
    kernel!(mesh.Edges.dcEdge, grad, workgroupsize=64, ndrange=(nEdges, vert_levels))

    KA.synchronize(backend)

    @show grad

    normSq = 0.0
    for i = 1:nEdges
        normSq += grad[i]^2
    end

    return normSq
end

mesh_url = "https://gist.github.com/mwarusz/f8caf260398dbe140d2102ec46a41268/raw/e3c29afbadc835797604369114321d93fd69886d/PlanarPeriodic48x48.nc"
mesh_fn  = "MokaMesh.nc"

Downloads.download(mesh_url, mesh_fn)

mesh = ReadHorzMesh(mesh_fn; backend=backend)

# Let's recreate all the variables:
gradNum = KA.zeros(backend, Float64, (1, mesh.Edges.nEdges))
for i = 1:mesh.Edges.nEdges
    gradNum[1,i] = gradNum[1,i] + i
end

d_gradNum = KA.zeros(backend, Float64, (1, mesh.Edges.nEdges))
d_mesh    = Enzyme.make_zero(mesh)

old_mesh = deepcopy(mesh)
gradient_normSq(gradNum, mesh)

@show isequal(mesh.PrimaryCells, old_mesh.PrimaryCells)
@show isequal(mesh.DualCells, old_mesh.DualCells)

@show isequal(mesh.Edges.nEdges, old_mesh.Edges.nEdges)
@show isequal(mesh.Edges.xᵉ, old_mesh.Edges.xᵉ)
@show isequal(mesh.Edges.yᵉ, old_mesh.Edges.yᵉ)
@show isequal(mesh.Edges.zᵉ, old_mesh.Edges.zᵉ)
@show isequal(mesh.Edges.fᵉ, old_mesh.Edges.fᵉ)
@show isequal(mesh.Edges.nEdgesOnEdge, old_mesh.Edges.nEdgesOnEdge)
@show isequal(mesh.Edges.cellsOnEdge, old_mesh.Edges.cellsOnEdge)
@show isequal(mesh.Edges.verticesOnEdge, old_mesh.Edges.verticesOnEdge)
@show isequal(mesh.Edges.edgesOnEdge, old_mesh.Edges.edgesOnEdge)
@show isequal(mesh.Edges.weightsOnEdge, old_mesh.Edges.weightsOnEdge)
@show isequal(mesh.Edges.dvEdge, old_mesh.Edges.dvEdge)
@show isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge)
@show isequal(mesh.Edges.angleEdge, old_mesh.Edges.angleEdge)

@show old_mesh.Edges.dcEdge
@show mesh.Edges.dcEdge
mesh = deepcopy(old_mesh)
d_normSq = autodiff(Enzyme.Reverse,
                    gradient_normSq,
                    Duplicated(gradNum, d_gradNum),
                    Duplicated(mesh, d_mesh))

@show isequal(mesh.PrimaryCells, old_mesh.PrimaryCells)
@show isequal(mesh.DualCells, old_mesh.DualCells)

@show isequal(mesh.Edges.nEdges, old_mesh.Edges.nEdges)
@show isequal(mesh.Edges.xᵉ, old_mesh.Edges.xᵉ)
@show isequal(mesh.Edges.yᵉ, old_mesh.Edges.yᵉ)
@show isequal(mesh.Edges.zᵉ, old_mesh.Edges.zᵉ)
@show isequal(mesh.Edges.fᵉ, old_mesh.Edges.fᵉ)
@show isequal(mesh.Edges.nEdgesOnEdge, old_mesh.Edges.nEdgesOnEdge)
@show isequal(mesh.Edges.cellsOnEdge, old_mesh.Edges.cellsOnEdge)
@show isequal(mesh.Edges.verticesOnEdge, old_mesh.Edges.verticesOnEdge)
@show isequal(mesh.Edges.edgesOnEdge, old_mesh.Edges.edgesOnEdge)
@show isequal(mesh.Edges.weightsOnEdge, old_mesh.Edges.weightsOnEdge)
@show isequal(mesh.Edges.dvEdge, old_mesh.Edges.dvEdge)
@show isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge)
@show isequal(mesh.Edges.angleEdge, old_mesh.Edges.angleEdge)

@show old_mesh.Edges.dcEdge
@show mesh.Edges.dcEdge

Non-differentiated code

isequal(mesh.PrimaryCells, old_mesh.PrimaryCells) = false
isequal(mesh.DualCells, old_mesh.DualCells) = false
isequal(mesh.Edges.nEdges, old_mesh.Edges.nEdges) = true
isequal(mesh.Edges.xᵉ, old_mesh.Edges.xᵉ) = true
isequal(mesh.Edges.yᵉ, old_mesh.Edges.yᵉ) = true
isequal(mesh.Edges.zᵉ, old_mesh.Edges.zᵉ) = true
isequal(mesh.Edges.fᵉ, old_mesh.Edges.fᵉ) = true
isequal(mesh.Edges.nEdgesOnEdge, old_mesh.Edges.nEdgesOnEdge) = true
isequal(mesh.Edges.cellsOnEdge, old_mesh.Edges.cellsOnEdge) = true
isequal(mesh.Edges.verticesOnEdge, old_mesh.Edges.verticesOnEdge) = true
isequal(mesh.Edges.edgesOnEdge, old_mesh.Edges.edgesOnEdge) = true
isequal(mesh.Edges.weightsOnEdge, old_mesh.Edges.weightsOnEdge) = true
isequal(mesh.Edges.dvEdge, old_mesh.Edges.dvEdge) = true
isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge) = true
isequal(mesh.Edges.angleEdge, old_mesh.Edges.angleEdge) = true

Enzyme results

isequal(mesh.PrimaryCells, old_mesh.PrimaryCells) = false
isequal(mesh.DualCells, old_mesh.DualCells) = false
isequal(mesh.Edges.nEdges, old_mesh.Edges.nEdges) = true
isequal(mesh.Edges.xᵉ, old_mesh.Edges.xᵉ) = true
isequal(mesh.Edges.yᵉ, old_mesh.Edges.yᵉ) = true
isequal(mesh.Edges.zᵉ, old_mesh.Edges.zᵉ) = true
isequal(mesh.Edges.fᵉ, old_mesh.Edges.fᵉ) = true
isequal(mesh.Edges.nEdgesOnEdge, old_mesh.Edges.nEdgesOnEdge) = true
isequal(mesh.Edges.cellsOnEdge, old_mesh.Edges.cellsOnEdge) = true
isequal(mesh.Edges.verticesOnEdge, old_mesh.Edges.verticesOnEdge) = true
isequal(mesh.Edges.edgesOnEdge, old_mesh.Edges.edgesOnEdge) = true
isequal(mesh.Edges.weightsOnEdge, old_mesh.Edges.weightsOnEdge) = true
isequal(mesh.Edges.dvEdge, old_mesh.Edges.dvEdge) = true
isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge) = true
isequal(mesh.Edges.angleEdge, old_mesh.Edges.angleEdge) = true

These are the same (in essence Enzyme behaves as if it calls the original function).

jlk9 commented 3 days ago

Enzyme is definitely affecting mesh.Edges.dcEdge. The code block you posted includes the @show grad statement in gradient_normSq, which I already found makes the change to dcEdge disappear.

using MOKA: HorzMesh, ReadHorzMesh
using KernelAbstractions
using Enzyme

import Downloads
import KernelAbstractions as KA

backend = KA.CPU()

# Setting meshes to inactive types:
Enzyme.EnzymeRules.inactive_type(::Type{T} where T <:HorzMesh) = true

@kernel function GradientOnEdgeModified(@Const(dcEdge), GradEdge)
    # global indices over nEdges
    iEdge, k = @index(Global, NTuple)

    @inbounds GradEdge[k, iEdge] = GradEdge[k, iEdge] / dcEdge[iEdge]

    @synchronize()
end

function gradient_normSq(grad, mesh::HorzMesh; backend=KA.CPU())

    nEdges = size(grad)[2]
    vert_levels = 1

    # New modified kernel:
    kernel! = GradientOnEdgeModified(backend)
    kernel!(mesh.Edges.dcEdge, grad, workgroupsize=64, ndrange=(nEdges, vert_levels))

    KA.synchronize(backend)

    #@show grad

    normSq = 0.0
    for i = 1:nEdges
        normSq += grad[i]^2
    end

    return normSq
end

mesh_url = "https://gist.github.com/mwarusz/f8caf260398dbe140d2102ec46a41268/raw/e3c29afbadc835797604369114321d93fd69886d/PlanarPeriodic48x48.nc"
mesh_fn  = "MokaMesh.nc"

Downloads.download(mesh_url, mesh_fn)

mesh = ReadHorzMesh(mesh_fn; backend=backend)

# Let's recreate all the variables:
gradNum = KA.zeros(backend, Float64, (1, mesh.Edges.nEdges))
for i = 1:mesh.Edges.nEdges
    gradNum[1,i] = gradNum[1,i] + i
end

d_gradNum = KA.zeros(backend, Float64, (1, mesh.Edges.nEdges))
d_mesh    = Enzyme.make_zero(mesh)

old_mesh = deepcopy(mesh)

d_normSq = autodiff(Enzyme.Reverse,
                    gradient_normSq,
                    Duplicated(gradNum, d_gradNum),
                    Duplicated(mesh, d_mesh))

#gradient_normSq(gradNum, mesh)

@show isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge)

Here's what happens when I run the above script (available in the link above) comparing mesh.Edges.dcEdge pre and post autodiff with (1) mesh specified as inactive and (2) no @show grad statement in gradient_normSq:

isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge) = false

when I comment out the call to autodiff and uncomment the call of just the forward function gradient_normSq, still without the show statement and with mesh an inactive type, this happens:

isequal(mesh.Edges.dcEdge, old_mesh.Edges.dcEdge) = true

In this circumstance Enzyme affects the values of dcEdge when the original function, on its own, does not.

jlk9 commented 3 days ago

Reduced the example in https://github.com/jlk9/MPAS-Ocean.jl/tree/bug-reduce further. Now it doesn't actually use the MOKA.jl package, instead using a minimal mesh struct as a fill-in (found in bugMesh.jl), and since the datafile MokaMesh is already downloaded we removed that portion too. The change of mesh.Edges.dcEdge values is still present when calling autodiff.