Snapshot from this documentation build. SHA-256: d0896b17ea51632bc1b225757f3ecd52683f422e8597e24da1659c539370e3f7
1 module MHDKernelsMetalExt
2
3 using MHDKernels
4 using Metal
5
6 const MG = Metal.MPSGraphs
7 using Metal.MPSGraphs: MPSGraph, MPSGraphFFTDescriptor, MPSGraphTensor,
8 MPSGraphTensorData, placeholderTensor, id, NSDictionary, NSArray, NSNumber,
9 NSString
10
11 struct _RadiationForwardPlan
12 graph::Any
13 input::Any
14 output::Any
15 end
16
17 struct _RadiationInversePlan
18 graph::Any
19 input::Any
20 output::Any
21 end
22
23 const _RADIATION_FORWARD_CACHE=Dict{NTuple{3,Int},_RadiationForwardPlan}()
24 const _RADIATION_INVERSE_CACHE=Dict{NTuple{3,Int},_RadiationInversePlan}()
25 const _RADIATION_FORWARD_BATCH_CACHE=Dict{NTuple{4,Int},_RadiationForwardPlan}()
26 const _RADIATION_INVERSE_BATCH_CACHE=Dict{NTuple{4,Int},_RadiationInversePlan}()
27
28 function _radiation_fft_descriptor(;inverse::Bool,
29 scaling=MG.MPSGraphFFTScalingModeNone)
30 d=MPSGraphFFTDescriptor(MG.@objc [MPSGraphFFTDescriptor descriptor]::id{MPSGraphFFTDescriptor})
31 d.inverse=inverse; d.scalingMode=scaling; d.roundToOddHermitean=false
32 d
33 end
34
35 function _radiation_forward_plan(N)
36 get!(_RADIATION_FORWARD_CACHE,N) do
37 g=MPSGraph(); x=placeholderTensor(g,N,Float32,"radiation_real")
38 axes=NSArray(NSNumber.([0,1,2]))
39 desc=_radiation_fft_descriptor(inverse=false)
40 out=MPSGraphTensor(MG.@objc [g::id{MPSGraph} realToHermiteanFFTWithTensor:x::id{MPSGraphTensor} axes:axes::id{NSArray} descriptor:desc::id{MPSGraphFFTDescriptor} name:"radiation_rfft"::id{NSString}]::id{MPSGraphTensor})
41 _RadiationForwardPlan(g,x,out)
42 end
43 end
44
45
46 function _radiation_inverse_plan(H)
47 get!(_RADIATION_INVERSE_CACHE,H) do
48 g=MPSGraph(); x=placeholderTensor(g,H,ComplexF32,"radiation_hat")
49 axes=NSArray(NSNumber.([0,1,2]))
50 desc=_radiation_fft_descriptor(inverse=true,scaling=MG.MPSGraphFFTScalingModeSize)
51 out=MPSGraphTensor(MG.@objc [g::id{MPSGraph} HermiteanToRealFFTWithTensor:x::id{MPSGraphTensor} axes:axes::id{NSArray} descriptor:desc::id{MPSGraphFFTDescriptor} name:"radiation_irfft"::id{NSString}]::id{MPSGraphTensor})
52 _RadiationInversePlan(g,x,out)
53 end
54 end
55
56 function _radiation_forward_batch_plan(N)
57 get!(_RADIATION_FORWARD_BATCH_CACHE,N) do
58 g=MPSGraph(); x=placeholderTensor(g,N,Float32,"radiation_real_batch")
59 # MPSGraph tensor dimensions are reversed relative to Julia. Axis 0 is
60 # the four-channel batch dimension; axes 1:3 are spatial.
61 axes=NSArray(NSNumber.([1,2,3]))
62 desc=_radiation_fft_descriptor(inverse=false)
63 out=MPSGraphTensor(MG.@objc [g::id{MPSGraph} realToHermiteanFFTWithTensor:x::id{MPSGraphTensor} axes:axes::id{NSArray} descriptor:desc::id{MPSGraphFFTDescriptor} name:"radiation_rfft_batch"::id{NSString}]::id{MPSGraphTensor})
64 _RadiationForwardPlan(g,x,out)
65 end
66 end
67
68 function _radiation_inverse_batch_plan(H)
69 get!(_RADIATION_INVERSE_BATCH_CACHE,H) do
70 g=MPSGraph(); x=placeholderTensor(g,H,ComplexF32,"radiation_hat_batch")
71 axes=NSArray(NSNumber.([1,2,3]))
72 desc=_radiation_fft_descriptor(inverse=true,scaling=MG.MPSGraphFFTScalingModeSize)
73 out=MPSGraphTensor(MG.@objc [g::id{MPSGraph} HermiteanToRealFFTWithTensor:x::id{MPSGraphTensor} axes:axes::id{NSArray} descriptor:desc::id{MPSGraphFFTDescriptor} name:"radiation_irfft_batch"::id{NSString}]::id{MPSGraphTensor})
74 _RadiationInversePlan(g,x,out)
75 end
76 end
77
78 function MHDKernels._radiation_forward_batch!(
79 chat::Metal.MtlArray{ComplexF32,4},x::Metal.MtlArray{Float32,4})
80 P=_radiation_forward_batch_plan(size(x))
81 MG.@autoreleasepool begin
82 feeds=Dict{MPSGraphTensor,MPSGraphTensorData}(P.input=>MPSGraphTensorData(x))
83 results=Dict{MPSGraphTensor,MPSGraphTensorData}(P.output=>MPSGraphTensorData(chat))
84 cb=MG.MPSCommandBuffer(Metal.global_queue(Metal.device()))
85 MG.encode!(cb,P.graph,NSDictionary(feeds),NSDictionary(results),MG.nil,
86 MG.MPSGraphExecutionDescriptor())
87 MG.commit!(cb); MG.wait_completed(cb)
88 end
89 chat
90 end
91
92 function MHDKernels._radiation_inverse_batch!(
93 x::Metal.MtlArray{Float32,4},chat::Metal.MtlArray{ComplexF32,4})
94 P=_radiation_inverse_batch_plan(size(chat))
95 MG.@autoreleasepool begin
96 feeds=Dict{MPSGraphTensor,MPSGraphTensorData}(P.input=>MPSGraphTensorData(chat))
97 results=Dict{MPSGraphTensor,MPSGraphTensorData}(P.output=>MPSGraphTensorData(x))
98 cb=MG.MPSCommandBuffer(Metal.global_queue(Metal.device()))
99 MG.encode!(cb,P.graph,NSDictionary(feeds),NSDictionary(results),MG.nil,
100 MG.MPSGraphExecutionDescriptor())
101 MG.commit!(cb); MG.wait_completed(cb)
102 end
103 x
104 end
105
106
107 function MHDKernels._radiation_forward!(chat::Metal.MtlArray{ComplexF32,3},
108 x::Metal.MtlArray{Float32,3})
109 P=_radiation_forward_plan(size(x))
110 MG.@autoreleasepool begin
111 feeds=Dict{MPSGraphTensor,MPSGraphTensorData}(P.input=>MPSGraphTensorData(x))
112 results=Dict{MPSGraphTensor,MPSGraphTensorData}(P.output=>MPSGraphTensorData(chat))
113 cb=MG.MPSCommandBuffer(Metal.global_queue(Metal.device()))
114 MG.encode!(cb,P.graph,NSDictionary(feeds),NSDictionary(results),MG.nil,
115 MG.MPSGraphExecutionDescriptor())
116 MG.commit!(cb); MG.wait_completed(cb)
117 end
118 chat
119 end
120
121
122 function MHDKernels._radiation_inverse!(x::Metal.MtlArray{Float32,3},
123 chat::Metal.MtlArray{ComplexF32,3})
124 P=_radiation_inverse_plan(size(chat))
125 MG.@autoreleasepool begin
126 feeds=Dict{MPSGraphTensor,MPSGraphTensorData}(P.input=>MPSGraphTensorData(chat))
127 results=Dict{MPSGraphTensor,MPSGraphTensorData}(P.output=>MPSGraphTensorData(x))
128 cb=MG.MPSCommandBuffer(Metal.global_queue(Metal.device()))
129 MG.encode!(cb,P.graph,NSDictionary(feeds),NSDictionary(results),MG.nil,
130 MG.MPSGraphExecutionDescriptor())
131 MG.commit!(cb); MG.wait_completed(cb)
132 end
133 x
134 end
135
136 function __init__()
137 if Metal.functional()
138 MHDKernels.register_backend!(:metal, MetalBackend())
139 end
140 end
141
142 MHDKernels.device_zeros(::MetalBackend, ::Type{T}, dims::Dims) where {T} = Metal.zeros(T, dims)
143
144 function MHDKernels._cube_launch!(::Val, be::MetalBackend, s::MHDKernels.MHDState{T}, dt::Real, ch::Real, decay::Real) where {T}
145 N = s.dims[1]; nb = N ÷ MHDKernels.CTB; dtdx = T(dt) / s.dx; cx, cy, cz = MHDKernels.bc_codes(s)
146 rec = Val(MHDKernels.recon_code_of(s)); per = Val(MHDKernels.all_periodic(s)); rsol = Val(MHDKernels.riemann_code_of(s))
147 physics = (dtdx, s.γ, T(ch), T(decay), s.smallr, s.pfl, s.llf_dmin, s.llf_pmin)
148 gs = MHDKernels.METAL_CUBE_GS
149 if N % MHDKernels.RCTBX == 0
150 nbx = N ÷ MHDKernels.RCTBX
151 nby = N ÷ MHDKernels.RCTBY
152 nbz = N ÷ MHDKernels.RCTBZ
153 MHDKernels.step_cube_rect_kernel!(be, gs)(
154 s.scratch..., s.U..., N, nbx, nby, cx, cy, cz,
155 Val(Float32), rec, per, rsol, physics, Val(false),
156 MHDKernels._no_drag_coefficients(T),Val(false),Val(false),Val(false),
157 Val(false),Val(false);
158 ndrange = nbx * nby * nbz * gs)
159 else
160 MHDKernels.step_cube_kernel!(be, gs)(
161 s.scratch..., s.U..., N, nb, cx, cy, cz, Val(Float32), rec, per, rsol,
162 dtdx, s.γ, T(ch), T(decay), s.smallr, s.pfl, s.llf_dmin, s.llf_pmin,
163 Val(false),MHDKernels._no_drag_coefficients(T),Val(false),Val(false),
164 Val(false),Val(false),Val(false);
165 ndrange = nb * nb * nb * gs)
166 end
167 MHDKernels.KA.synchronize(be)
168 end
169
170 function MHDKernels._cube_drag_launch!(be::MetalBackend,
171 s::MHDKernels.MHDState{T}, dt::Real, ch::Real, decay::Real,
172 q::T) where {T}
173 N = s.dims[1]; nb = N ÷ MHDKernels.CTB; dtdx = T(dt) / s.dx
174 cx, cy, cz = MHDKernels.bc_codes(s)
175 rec = Val(MHDKernels.recon_code_of(s)); per = Val(MHDKernels.all_periodic(s))
176 rsol = Val(MHDKernels.riemann_code_of(s))
177 dc = MHDKernels._drag_coefficients(q)
178 physics = (dtdx, s.γ, T(ch), T(decay), s.smallr, s.pfl, s.llf_dmin, s.llf_pmin)
179 gs = MHDKernels.METAL_CUBE_GS
180 if N % MHDKernels.RCTBX == 0
181 nbx = N ÷ MHDKernels.RCTBX
182 nby = N ÷ MHDKernels.RCTBY
183 nbz = N ÷ MHDKernels.RCTBZ
184 MHDKernels.step_cube_rect_kernel!(be, gs)(
185 s.scratch..., s.U..., N, nbx, nby, cx, cy, cz,
186 Val(Float32), rec, per, rsol, physics, Val(true),
187 dc,Val(false),Val(false),Val(false),Val(false),Val(false);
188 ndrange = nbx * nby * nbz * gs)
189 else
190 MHDKernels.step_cube_kernel!(be, gs)(
191 s.scratch..., s.U..., N, nb, cx, cy, cz, Val(Float32), rec, per, rsol,
192 dtdx, s.γ, T(ch), T(decay), s.smallr, s.pfl, s.llf_dmin, s.llf_pmin,
193 Val(true),dc,Val(false),Val(false),Val(false),Val(false),Val(false);
194 ndrange = nb * nb * nb * gs)
195 end
196 MHDKernels.KA.synchronize(be)
197 end
198
199
200 function MHDKernels._cube_radiation_launch!(be::MetalBackend,
201 s::MHDKernels.MHDState{T},dt::Real,ch::Real,decay::Real,q::T,
202 correction,track_density::Val,track_magnetic::Val,track_energy::Val,
203 check::Val,
204 fallback::Val{FALLBACK},global_fallback::Val{GLOBAL},
205 robust_fallback::Val{ROBUST}) where {T,FALLBACK,GLOBAL,ROBUST}
206 N=s.dims[1]; dtdx=T(dt)/s.dx; cx,cy,cz=MHDKernels.bc_codes(s)
207 rec=ROBUST ? Val(MHDKernels.RECON_PLM_MINMOD) :
208 Val(MHDKernels.recon_code_of(s))
209 per=Val(MHDKernels.all_periodic(s))
210 rsol=ROBUST ? Val(MHDKernels.RSOLVE_HLL) :
211 GLOBAL ? Val(MHDKernels.RSOLVE_LLF) :
212 Val(MHDKernels.riemann_code_of(s))
213 dc=MHDKernels._drag_coefficients(q)
214 physics=(dtdx,s.γ,T(ch),T(decay),s.smallr,s.pfl,s.llf_dmin,s.llf_pmin)
215 kernel_check=GLOBAL ? Val(3) : FALLBACK ? Val(2) : check
216 gs=MHDKernels.METAL_CUBE_GS
217 if N % MHDKernels.RCTBX == 0
218 nbx=N÷MHDKernels.RCTBX; nby=N÷MHDKernels.RCTBY; nbz=N÷MHDKernels.RCTBZ
219 MHDKernels.step_cube_rect_kernel!(be,gs)(s.scratch...,s.U[1],s.U[2],
220 s.U[3],s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],correction,
221 N,nbx,nby,cx,cy,cz,Val(Float32),rec,per,rsol,physics,
222 Val(true),dc,Val(true),track_density,
223 track_magnetic,track_energy,kernel_check;
224 ndrange=nbx*nby*nbz*gs)
225 else
226 nb=N÷MHDKernels.CTB
227 MHDKernels.step_cube_kernel!(be,gs)(s.scratch...,s.U[1],s.U[2],s.U[3],
228 s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],correction,N,nb,cx,cy,cz,
229 Val(Float32),rec,per,rsol,dtdx,s.γ,T(ch),T(decay),s.smallr,s.pfl,
230 s.llf_dmin,s.llf_pmin,Val(true),dc,Val(true),track_density,
231 track_magnetic,track_energy,kernel_check;
232 ndrange=nb*nb*nb*gs)
233 end
234 MHDKernels.KA.synchronize(be)
235 end
236
237 end # module