/
diffusion_2D_kp.jl
100 lines (94 loc) · 4.02 KB
/
diffusion_2D_kp.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
using AMDGPU, ImplicitGlobalGrid, Plots
macro d_xa(A,ix,iy) esc(:( ($A[$ix+1,$iy+0] - $A[$ix,$iy+0]) )) end
macro d_ya(A,ix,iy) esc(:( ($A[$ix+0,$iy+1] - $A[$ix+0,$iy]) )) end
macro d_xi(A,ix,iy) esc(:( ($A[$ix+1,$iy+1] - $A[$ix,$iy+1]) )) end
macro d_yi(A,ix,iy) esc(:( ($A[$ix+1,$iy+1] - $A[$ix+1,$iy]) )) end
macro inn(A,ix,iy) esc(:( ($A[$ix+1,$iy+1]) )) end
macro all(A,ix,iy) esc(:( ($A[$ix ,$iy ]) )) end
"""
# Fourier's law of heat conduction:
q_x = -λ ∂T/∂x
qx .= -lam .* d_xi(T)./dx
qy .= -lam .* d_yi(T)./dy
"""
function Flux!(qx, qy, T, lam, _dx, _dy)
ix = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x # only workgroupIdx starts at 1
iy = (workgroupIdx().y - 1) * workgroupDim().y + workitemIdx().y # only workgroupIdx starts at 1
if (ix<=size(qx,1) && iy<=size(qx,2))
@all(qx,ix,iy) = -lam * @d_xi(T,ix,iy) * _dx
end
if (ix<=size(qy,1) && iy<=size(qy,2))
@all(qy,ix,iy) = -lam * @d_yi(T,ix,iy) * _dy
end
return
end
"""
# Conservation of energy:
∂T/∂t = 1/cₚ (-∂q_x/∂x - ∂q_y/∂y)
dTdt .= 1.0./inn(Cp) .* (-d_xa(qx)./dx .- d_ya(qy)./dy)
"""
function Residual!(dTdt, qx, qy, Cp, _dx, _dy)
ix = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x # only workgroupIdx starts at 1
iy = (workgroupIdx().y - 1) * workgroupDim().y + workitemIdx().y # only workgroupIdx starts at 1
if (ix<=size(dTdt,1) && iy<=size(dTdt,2))
@all(dTdt,ix,iy) = 1.0/@inn(Cp,ix,iy) * - (@d_xa(qx,ix,iy) * _dx + @d_ya(qy,ix,iy) * _dy)
end
return
end
"""
# Update of temperature
T[2:end-1,2:end-1] .= inn(T) .+ dt .* dTdt
T_new = T_old + ∂T/∂t
"""
function Update!(T, dt, dTdt)
ix = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x # only workgroupIdx starts at 1
iy = (workgroupIdx().y - 1) * workgroupDim().y + workitemIdx().y # only workgroupIdx starts at 1
if (ix<=size(T,1)-2 && iy<=size(T,2)-2)
@inn(T,ix,iy) += dt * @all(dTdt,ix,iy)
end
return
end
@views function diffusion2D()
# Physics
lx, ly = 10.0, 10.0 # Length of domain in dimensions x, y and z
lam = 1.0 # Thermal conductivity
Cp0 = 1.0
# Numerics
nx, ny = 128, 128
threads = (32, 8)
grid = (nx, ny)
nt = 1e3 # Number of time steps
me, dims, nprocs, coords, comm_cart = init_global_grid(nx, ny, 1) # Initialize the implicit global grid
println("Process $me selecting device $(AMDGPU.default_device_id())")
dx, dy = lx/nx_g(), ly/ny_g() # Space step in dimension x
_dx,_dy = 1.0/dx, 1.0/dy
dt = min(dx*dx,dy*dy)*Cp0/lam/4.1 # Time step for the 3D Heat diffusion
# Array initializations
dTdt = AMDGPU.zeros(Float64, nx-2, ny-2)
qx = AMDGPU.zeros(Float64, nx-1, ny-2)
qy = AMDGPU.zeros(Float64, nx-2, ny-1)
Cp = Cp0.*AMDGPU.ones(Float64, nx, ny)
T = zeros(Float64, nx, ny)
# Initial conditions
T = ROCArray([exp(-(x_g(ix,dx,T)+dx/2 -lx/2)^2 -(y_g(iy,dy,T)+dy/2 -ly/2)^2) for ix=1:size(T,1), iy=1:size(T,2)])
# visu
if (me==0) gr(); ENV["GKSwstype"]="nul"; !ispath("../output") && mkdir("../output"); end
nx_v = (nx-2)*dims[1]
ny_v = (ny-2)*dims[2]
T_v = zeros(nx_v, ny_v)
T_nh = zeros(nx-2, ny-2)
# Time loop
me==0 && print("Starting the time loop 🚀...")
for it = 1:nt
wait( @roc groupsize=threads gridsize=grid Flux!(qx, qy, T, lam, _dx, _dy) )
wait( @roc groupsize=threads gridsize=grid Residual!(dTdt, qx, qy, Cp, _dx, _dy) )
wait( @roc groupsize=threads gridsize=grid Update!(T, dt, dTdt) )
update_halo!(T)
end
me==0 && println("done")
T_nh .= Array(T[2:end-1,2:end-1])
gather!(T_nh, T_v)
if (me==0) heatmap(transpose(T_v)); png("../output/Temp_kp_$(nprocs)_$(nx_g())_$(ny_g()).png"); end
finalize_global_grid() # Finalize the implicit global grid
end
diffusion2D()