Added option to save memory

This commit is contained in:
Alberto Ramos 2021-09-24 15:49:53 +02:00
parent 0587e5ffea
commit 879ca7dab4
5 changed files with 152 additions and 153 deletions

View file

@ -37,19 +37,36 @@ function krnl_plaq!(plx, U, lp::SpaceParm)
return nothing
end
function krnl_force_wilson_pln!(frc1, frc2, U, ipl, lp::SpaceParm)
function krnl_force_wilson_pln!(frc1, frc2, U::AbstractArray{T}, ipl, lp::SpaceParm{N,M,D}) where {T,N,M,D}
b, r = CUDA.threadIdx().x, CUDA.blockIdx().x
Ush = @cuStaticSharedMem(T, (D,2))
@inbounds begin
id1, id2 = lp.plidx[ipl]
bu1, ru1 = up((b, r), id1, lp)
bu2, ru2 = up((b, r), id2, lp)
Ush[b,1] = U[b,id1,r]
Ush[b,2] = U[b,id2,r]
sync_threads()
if ru2 == r
gt2 = Ush[bu2,1]
else
gt2 = U[bu2,id1,ru2]
end
if ru1 == r
gt1 = Ush[bu1,2]
else
gt1 = U[bu1,id2,ru1]
end
g1 = U[bu1,id2,ru1]/U[bu2,id1,ru2]
g2 = U[b,id2,r]\U[b,id1,r]
g1 = gt1/gt2
g2 = Ush[b,2]\Ush[b,1]
F1 = projalg(U[b,id1,r]*g1/U[b,id2,r])
F1 = projalg(Ush[b,1]*g1/Ush[b,2])
F2 = projalg(g1*g2)
F3 = projalg(g2*g1)
@ -62,71 +79,78 @@ function krnl_force_wilson_pln!(frc1, frc2, U, ipl, lp::SpaceParm)
return nothing
end
function krnl_force_wilson_nw!(fpl, U, lp::SpaceParm)
b = CUDA.threadIdx().x
ipl = mod1(CUDA.blockIdx().x, lp.npls)
r = div(CUDA.blockIdx().x-1, lp.npls)+1
@inbounds begin #for ipl in 1:lp.npls
id1, id2 = lp.plidx[ipl]
bu1, ru1 = up((b, r), id1, lp)
bu2, ru2 = up((b, r), id2, lp)
g1 = U[bu1,id2,ru1]/U[bu2,id1,ru2]
g2 = U[b,id2,r]\U[b,id1,r]
F1 = projalg(U[b,id1,r]*g1/U[b,id2,r])
F2 = projalg(g1*g2)
F3 = projalg(g2*g1)
fpl[b ,ipl, r ,1,1] = -F1
fpl[b ,ipl, r ,2,1] = F1
fpl[bu1,ipl, ru1,2,2] = -F2
fpl[bu2,ipl, ru2,1,2] = F3
end
return nothing
end
function krnl_force_wilson!(frc, U, lp::SpaceParm)
function krnl_force_wilson_nw!(fpl, U::AbstractArray{T}, lp::SpaceParm{N,M,D}) where {T,N,M,D}
b, r = CUDA.threadIdx().x, CUDA.blockIdx().x
@inbounds for id1 in 1:lp.ndim
bu1, ru1, bd1, rd1 = updw((b, r), id1, lp)
@inbounds for id2 in id1+1:lp.ndim
bu2, ru2, bd2, rd2 = updw((b, r), id2, lp)
bud, rud = dw((bu1,ru1), id2, lp)
bdu, rdu = dw((bu2,ru2), id1, lp)
F1 = projalg(U[b,id1,r]*U[bu1,id2,ru1]/(U[b,id2,r]*U[bu2,id1,ru2]))
F2 = projalg((U[b,id2,r]/(U[bd1,id2,rd1]*U[bdu,id1,rdu]))*U[bd1,id1,rd1])
F3 = projalg((U[bd2,id2,rd2]\U[bd2,id1,rd2])*(U[bud,id2,rud]/U[b,id1,r]))
frc[b,id1,r] -= F1 - F3
frc[b,id2,r] += F1 - F2
end
Ush = @cuStaticSharedMem(T, (D,N))
for id in 1:N
@inbounds Ush[b,id] = U[b,id,r]
end
return nothing
end
function krnl_add_force_plns!(frc, fpl, lp::SpaceParm)
b, r = CUDA.threadIdx().x, CUDA.blockIdx().x
sync_threads()
@inbounds for ipl in 1:lp.npls
id1, id2 = lp.plidx[ipl]
frc[b,id1,r] += fpl[b,ipl,r,1,1] + fpl[b,ipl,r,1,2]
frc[b,id2,r] += fpl[b,ipl,r,2,1] + fpl[b,ipl,r,2,2]
bu1, ru1 = up((b, r), id1, lp)
bu2, ru2 = up((b, r), id2, lp)
if ru2 == r
gt2 = Ush[bu2,id1]
else
gt2 = U[bu2,id1,ru2]
end
if ru1 == r
gt1 = Ush[bu1,id2]
else
gt1 = U[bu1,id2,ru1]
end
g1 = gt1/gt2
g2 = Ush[b,id2]\Ush[b,id1]
fpl[b ,ipl, r ,1] = projalg(Ush[b,id1]*g1/Ush[b,id2])
fpl[bu1,ipl, ru1,2] = projalg(g1*g2)
fpl[bu2,ipl, ru2,3] = projalg(g2*g1)
end
return nothing
end
function force0_wilson_pln!(frc1, ftmp, U, lp::SpaceParm)
function krnl_add_force_plns!(frc::AbstractArray{T}, fpl, lp::SpaceParm{N,M,D}) where {T,N,M,D}
b, r = CUDA.threadIdx().x, CUDA.blockIdx().x
Fsh = @cuStaticSharedMem(T, (D,M))
@inbounds for j in 1:M
Fsh[b,j] = fpl[b,j,r,1]
end
sync_threads()
@inbounds for ipl in 1:lp.npls
id1, id2 = lp.plidx[ipl]
frc[b,id1,r] += -Fsh[b,ipl] + fpl[b,ipl,r,3]
frc[b,id2,r] += Fsh[b,ipl] - fpl[b,ipl,r,2]
end
return nothing
end
function force_wilson(ymws::YMworkspace, U, lp::SpaceParm)
if ymws.fpln == nothing
force_wilson_pln!(ymws.frc1, ymws.frc2, U, lp::SpaceParm)
else
force_wilson_nw!(ymws.frc1, ymws.fpln, U, lp::SpaceParm)
end
return nothing
end
function force_wilson_pln!(frc1, ftmp, U, lp::SpaceParm)
fill!(frc1, zero(eltype(frc1)))
fill!(ftmp, zero(eltype(ftmp)))
@ -140,20 +164,10 @@ function force0_wilson_pln!(frc1, ftmp, U, lp::SpaceParm)
return nothing
end
function force0_wilson!(frc1, U, lp::SpaceParm)
fill!(frc1, zero(eltype(frc1)))
CUDA.@sync begin
CUDA.@cuda threads=lp.bsz blocks=lp.rsz krnl_force_wilson!(frc1,U,lp)
end
return nothing
end
function force0_wilson_nw!(frc1, fpln, U, lp::SpaceParm)
function force_wilson_nw!(frc1, fpln, U, lp::SpaceParm)
CUDA.@sync begin
CUDA.@cuda threads=lp.bsz blocks=lp.rsz*lp.npls krnl_force_wilson_nw!(fpln,U,lp)
CUDA.@cuda threads=lp.bsz blocks=lp.rsz krnl_force_wilson_nw!(fpln,U,lp)
end
fill!(frc1, zero(eltype(frc1)))
CUDA.@sync begin