Running out of memory

[raminammour]

Hello,

Thanks for finally supporting mutation in AD!!

I am on a mac and on julia 1.7 it seems that an example of this size will die (allocates a lot of memory and goes to swap, I guess the kernel kills it). I tried for the function to be non-allocating as below, and didn't help.

function fwd3!!(d,c,pp,pc)
    n1,n2,n3=size(c)
    nt=length(d)
    
    #pp=zeros(n1,n2,n3)
    #pc=zeros(n1,n2,n3)
    pp.=0.
    pc.=0.
    pp[n1÷2,n2÷2,n3÷2]=1
    
    for it in 1:nt
        for iy in 2:n3-1
            for ix in 2:n2-1
                for iz in 2:n1-1
                    pc[iz,ix,iy] = 2pp[iz,ix,iy] - pc[iz,ix,iy] + 
                        1e-8*c[iz,ix,iy]^2*(pp[iz-1,ix,iy]-2pp[iz,ix,iy]+pp[iz+1,ix,iy])+
                        1e-8*c[iz,ix,iy]^2*(pp[iz,ix-1,iy]-2pp[iz,ix,iy]+pp[iz,ix+1,iy])+
                        1e-8*c[iz,ix,iy]^2*(pp[iz,ix,iy-1]-2pp[iz,ix,iy]+pp[iz,ix,iy+1])
                end
            end
        end
        pc,pp=pp,pc
        d[it] = pc[n1÷4,n2÷3,n3÷3]
    end
    return nothing
end

c=1e3*ones(101,111,105)
pp,pc=similar(c),similar(c)
d=zeros(1500)
@time fwd3!!(d,c,pp,pc)

function adj3!(d,c,dd)
    ddc=copy(dd)
    dc=zeros(size(c)...)
    pp,pc=similar(c),similar(c)
    tpp,tpc=zeros(size(c...)),zeros(size(c)...)
    Enzyme.autodiff(fwd3!!, Duplicated(d,ddc), Duplicated(c,dc),Duplicated(pp,tpp),Duplicated(pc,tpc))
    return dc
end

ddr=rand(size(d)...)
@time dc=adj3!(d,c,ddr); ### and dead

Is there a way around that? I understand that the size of the problem is big, but these are the sizes I am interested in for the applications I have in mind :)

Also, while I am here... Is there a way to define pushforwards in enzyme? I asked on discourse, so excuse me for asking here as well!

Cheers!

[raminammour]

An afterthought: having written the pullback by hand a few times, is enzyme effectively saving pp/pc for every it in 1:nt?

[vchuravy]

That might be cc: @wsmoses