Merge d35d0ce into 6e87e04

Manifest.toml

@@ -16,10 +16,10 @@ version = "1.0.0"
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[BinaryProvider]]
-deps = ["Libdl", "SHA"]
+deps = ["Libdl", "Logging", "SHA"]
 git-tree-sha1 = "c7361ce8a2129f20b0e05a89f7070820cfed6648"
 uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
-version = "0.5.4"
+version = "0.5.6"
 
 [[CodecZlib]]
 deps = ["BinaryProvider", "Libdl", "Test", "TranscodingStreams"]
@@ -46,10 +46,10 @@ uuid = "a2441757-f6aa-5fb2-8edb-039e3f45d037"
 version = "0.9.2"
 
 [[DataStructures]]
-deps = ["InteractiveUtils", "OrderedCollections", "Random", "Serialization", "Test"]
-git-tree-sha1 = "ca971f03e146cf144a9e2f2ce59674f5bf0e8038"
+deps = ["InteractiveUtils", "OrderedCollections"]
+git-tree-sha1 = "0809951a1774dc724da22d26e4289bbaab77809a"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.15.0"
+version = "0.17.0"
 
 [[Dates]]
 deps = ["Printf"]
@@ -77,9 +77,9 @@ version = "1.0.7"
 
 [[HTTP]]
 deps = ["Base64", "Dates", "IniFile", "MbedTLS", "Sockets"]
-git-tree-sha1 = "854fad2a2b9cc6678f70ed15a65fc7e572056d39"
+git-tree-sha1 = "03ddc88af7f2d963fac5aa9f3ac8e11914d68a78"
 uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
-version = "0.8.2"
+version = "0.8.4"
 
 [[IniFile]]
 deps = ["Test"]
@@ -174,6 +174,12 @@ git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
 uuid = "189a3867-3050-52da-a836-e630ba90ab69"
 version = "0.2.0"
 
+[[Requires]]
+deps = ["Test"]
+git-tree-sha1 = "f6fbf4ba64d295e146e49e021207993b6b48c7d1"
+uuid = "ae029012-a4dd-5104-9daa-d747884805df"
+version = "0.5.2"
+
 [[SHA]]
 uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 

Project.toml

@@ -18,5 +18,18 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[compat]
+FFTW = "≥ 0.2.4"
+Interpolations = "≥ 0.8.0"
+JLD2 = "≥ 0.1.2"
+julia = "≥ 0.7.0"
+
+[extras]
+CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test", "CuArrays"]

src/CuFourierFlows.jl

@@ -0,0 +1,34 @@
+using .CuArrays
+
+# Discard `effort` argument for CuArrays
+plan_flows_fft(a::CuArray, effort) = plan_fft(a)
+plan_flows_rfft(a::CuArray, effort) = plan_rfft(a)
+
+OneDGrid(dev::GPU, args...; kwargs...) = OneDGrid(args...; ArrayType=CuArray, kwargs...)
+TwoDGrid(dev::GPU, args...; kwargs...) = TwoDGrid(args...; ArrayType=CuArray, kwargs...)
+
+function Base.zeros(::GPU, T, dims)
+    a = CuArray{T}(undef, dims...)
+    a .= 0
+    return a
+end
+
+ArrayType(::GPU, T, dim) = CuArray{T, dim}
+ArrayType(::GPU) = CuArray
+
+supersize(a::CuArray) = size(a)
+
+getetdcoeffs(dt, L::CuArray; kwargs...) = 
+    (CuArray(ζ) for ζ in getetdcoeffs(dt, Array(L); kwargs...))
+
+makefilter(K::CuArray; kwargs...) = CuArray(makefilter(Array(K); kwargs...))
+
+function makefilter(g::AbstractGrid{Tg, <:CuArray}, T, sz; kwargs...) where Tg
+    CuArray(ones(T, sz)) .* makefilter(g; realvars=sz[1]==g.nkr, kwargs...)
+end
+
+function gridpoints(g::AbstractGrid{Tg, <:CuArray}) where Tg
+  X = [ g.x[i] for i=1:g.nx, j=1:g.ny]
+  Y = [ g.y[j] for i=1:g.nx, j=1:g.ny]
+  CuArray(X), CuArray(Y)
+end

src/FourierFlows.jl

@@ -1,6 +1,15 @@
 module FourierFlows
 
 export
+  # Helper variables and macros for determining if machine is CUDA-enabled.
+  HAVE_CUDA,
+  @hascuda,
+
+  Device,
+  CPU,
+  GPU,
+  ArrayType,
+
   cxtype,
   fltype,
   innereltype,
@@ -27,8 +36,10 @@ export
   savediagnostic,
 
   @zeros,
+  @devzeros,
   @createarrays,
   @superzeros,
+  devzeros,
   superzeros,
   supersize,
 
@@ -47,22 +58,25 @@ using
   FFTW,
   JLD2,
   Statistics,
-  Interpolations
+  Interpolations,
+  Requires
 
 import Base: resize!, getindex, setindex!, lastindex, push!, append!
 
 using Base: fieldnames
 
 using LinearAlgebra: mul!, ldiv!
 
-abstract type AbstractGrid{T} end
-abstract type AbstractTwoDGrid{T} <: AbstractGrid{T} end
-abstract type AbstractOneDGrid{T} <: AbstractGrid{T} end
+abstract type AbstractGrid{T, Ta} end
 abstract type AbstractTimeStepper{T} end
 abstract type AbstractParams end
 abstract type AbstractVars end
 abstract type AbstractDiagnostic end
 
+abstract type Device end
+struct CPU <: Device end
+struct GPU <: Device end
+
 # The main show
 include("problem.jl")
 include("domains.jl")
@@ -74,4 +88,21 @@ include("timesteppers.jl")
 # Physics
 include("diffusion.jl")
 
+# Import CUDA utilities if cuda is detected.
+const HAVE_CUDA = try
+    using CuArrays
+    true
+catch
+    false
+end
+
+macro hascuda(ex)
+    return HAVE_CUDA ? :($(esc(ex))) : :(nothing)
+end
+
+
+function __init__()
+    @require CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae" include("CuFourierFlows.jl")
+end
+
 end # module

src/diffusion.jl

@@ -25,48 +25,55 @@ function Problem(;
          kappa = 0,
             dt = 0.01,
        stepper = "RK4",
-             T = Float64
+             T = Float64,
+           dev = CPU()
   )
 
-    grid = OneDGrid(nx, Lx; T=T)
-  params = Params(kappa)
-    vars = Vars(grid)
-     eqn = Equation(kappa, grid)
+    grid = OneDGrid(dev, nx, Lx; T=T)
+  params = Params(dev, kappa)
+    vars = Vars(dev, grid)
+     eqn = DiffusionEquation(dev, kappa, grid)
 
-  FourierFlows.Problem(eqn, stepper, dt, grid, vars, params)
+  FourierFlows.Problem(dev, eqn, stepper, dt, grid, vars, params)
 end
 
 struct Params{T} <: AbstractParams
   kappa::T
 end
 
+Params(dev, kappa::Number) = Params(kappa)
+Params(dev, kappa::AbstractArray) = Params(ArrayType(dev)(kappa))
+
 """
     Equation(p, g)
 
 Returns the equation for constant diffusivity problem with params p and grid g.
 """
-function Equation(kappa::T, g) where T<:Number
-  FourierFlows.Equation(-kappa*g.kr.^2, calcN!, g)
+function DiffusionEquation(dev::Device, kappa::T, grid) where T<:Number
+  L = zeros(dev, T, grid.nkr)
+  @. L = -kappa * grid.kr^2
+  FourierFlows.Equation(L, calcN!, grid)
 end
 
-function Equation(kappa::T, g::AbstractGrid{Tg}) where {T<:AbstractArray,Tg}
-  FourierFlows.Equation(0, calcN!, g; dims=(g.nkr,), T=cxtype(Tg))
+function DiffusionEquation(dev::Device, kappa::T, grid::AbstractGrid{Tg}) where {T<:AbstractArray, Tg}
+  FourierFlows.Equation(0, calcN!, grid; dims=(grid.nkr,), T=cxtype(Tg))
 end
 
-# Construct Vars types
-const physicalvars = [:c, :cx]
-const  fouriervars = [:ch, :cxh]
-
-eval(varsexpression(:Vars, physicalvars, fouriervars))
+struct Vars{Aphys, Atrans} <: AbstractVars
+    c :: Aphys
+   cx :: Aphys
+   ch :: Atrans
+  cxh :: Atrans
+end
 
 """
-    Vars(g)
+    Vars(dev, grid)
 
 Returns the vars for constant diffusivity problem on grid g.
 """
-function Vars(g::AbstractGrid{T}) where T
-  @zeros T g.nx c cx
-  @zeros Complex{T} g.nkr ch cxh
+function Vars(::Dev, grid::AbstractGrid{T}) where {Dev, T}
+  @devzeros Dev T grid.nx c cx
+  @devzeros Dev Complex{T} grid.nkr ch cxh
   Vars(c, cx, ch, cxh)
 end