WIP

Project.toml

@@ -11,6 +11,7 @@ BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Continuables = "79afa230-ca09-11e8-120b-5decf7bf5e25"
 Crayons = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
+Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

misc/dev_float32.jl

@@ -1,9 +1,95 @@
 using Dolo
 
 root_dir = pkgdir(Dolo)
-model32 = include("$(root_dir)/misc/rbc_float32.jl")
+model_32 = include("$(root_dir)/misc/rbc_float32.jl")
 
 
-dm32 = Dolo.discretize(model, Dict(:endo=>[10000000]) )
+model32 = Dolo.convert_precision(Float32, model_32)
 
-Dolo.convert
+dm32 = Dolo.discretize(model32, Dict(:endo=>[100]) )
+
+
+typeof(dm32)
+
+using Adapt
+
+
+wk0 = Dolo.time_iteration_workspace(dm32)
+
+import oneAPI: oneArray
+import Adapt: adapt_structure
+
+wk = adapt(oneArray, wk0)
+using KernelAbstractions: get_backend
+
+if typeof(dm32.grid)<:Dolo.CGrid
+    p = dm32.grid.ranges[1][3]
+    q = dm32.grid.ranges[2][3]
+else
+    p = length(dm32.grid.grids[1])
+    q = length(dm32.grid.grids[2])
+end
+
+t_e = get_backend(wk.x0)
+
+using KernelAbstractions
+using StaticArrays
+
+@kernel function ggg(r, dm, x, φ)
+# @kernel function fun(dm, r, x)
+
+
+    n = @index(Global, Linear)
+    ind = @index(Global, Cartesian)
+    (i,j) = ind.I
+
+
+    # k = length(dm.grid.grids[1])
+    # i_ = (n-1)÷k
+    # j_ = (n-1) - (i)*κ
+    # i = i_+1
+    # j = j_+1
+
+    # TODO: special function here
+    s_ = dm.grid[n]
+    s = Dolo.QP((i,j), s_)
+
+    xx = x[n]
+
+    # (i,j) = @inline Dolo.from_linear(model.grid, n)    
+
+    rr = Dolo.F(dm, s, xx, φ)
+
+    r[n] = rr
+
+end
+
+
+
+fun_ = ggg(t_e)
+@time fun_(wk.r0, dm32, wk.x0, wk.φ; ndrange=(p,q))
+
+
+
+@time Dolo.F!(wk0.r0, dm32, wk0.x0, wk0.φ)
+
+# try
+
+    # res = fun_(wk.r0, dm32, wk.x0, wk.φ; ndrange=(p,q))
+# catch err
+    # nothing
+# end
+
+
+
+# Dolo.time_iteration(dm32, wk; engine=:gpu)
+
+
+s_0 = dm32.grid[10]
+s0 = Dolo.QP((2,2), s_0)
+xx0 = wk0.x0[10]
+
+
+fon(dm32, s0, xx0) = sum(w*S.val for (w,S) in Dolo.τ(dm32, s0, xx0))
+
+@code_warntype fon(dm32, s0, xx0)

misc/test_gpu.jl

@@ -107,12 +107,6 @@ spl = Dolo.SplineInterpolator{typeof(grid), typeof(C), 3}(grid, C)
 # end
 
 
-adapt_structure(to, s::Dolo.SplineInterpolator{G,C,k}) where G where C where k = let
-    θθ = adapt(to, s.θ)
-    CC = typeof(θθ)
-    Dolo.SplineInterpolator{G,CC,k}(s.grid, θθ)
-end
-
 
 gspl = adapt(CuArray, spl)
 
@@ -170,6 +164,14 @@ Dolo.calibrated(model, :controls)
 
 import Adapt: adapt_storage
 
+
+
+adapt_structure(to, s::Dolo.SplineInterpolator{G,C,k}) where G where C where k = let
+    θθ = adapt(to, s.θ)
+    CC = typeof(θθ)
+    Dolo.SplineInterpolator{G,CC,k}(s.grid, θθ)
+end
+
 function adapt_structure(to, v::Dolo.GVector{G,V}) where G where V
     data = adapt(to, v.data)
     Dolo.GVector{G, typeof(data)}(

src/Dolo.jl

@@ -83,6 +83,7 @@ module Dolo
     include("algos/vfi.jl")
 
 
+    include("adapt.jl")
     include("utils.jl")
 
     # function yaml_import(filename)

src/adapt.jl

@@ -0,0 +1,36 @@
+import Adapt: adapt, adapt_structure
+
+adapt_structure(to, s::SplineInterpolator{G,C,k}) where G where C where k = let
+    θθ = adapt(to, s.θ)
+    CC = typeof(θθ)
+    SplineInterpolator{G,CC,k}(s.grid, θθ)
+end
+
+function adapt_structure(to, v::GVector{G,V}) where G where V
+    data = adapt(to, v.data)
+    GVector{G, typeof(data)}(
+        v.grid,
+        data
+    )
+end
+
+function adapt_structure(to, v::GArray{G,V}) where G where V
+    data = adapt(to, v.data)
+    GArray(
+        v.grid,
+        data
+    )
+end
+function adapt_structure(to, f::DFun{Dom, Gar, Itp, vars}) where Dom where Gar where Itp where vars
+    itp = adapt(to, f.itp)
+    values = adapt(to, f.values)
+    DFun{Dom, typeof(values), typeof(itp), vars}(
+        f.domain,
+        values,
+        itp
+    )
+end
+
+import KernelAbstractions: get_backend
+
+get_backend(g::GArray) = get_backend(g.data)

src/algos/time_iteration.jl

@@ -126,9 +126,9 @@ end
 
 
 using LinearMaps
+using Adapt
 
-
-function time_iteration_workspace(dmodel; interp_mode=:linear)
+function time_iteration_workspace(dmodel; interp_mode=:linear, dest=Array)
 
     T = eltype(dmodel)
 
@@ -145,7 +145,10 @@ function time_iteration_workspace(dmodel; interp_mode=:linear)
     )
     vars = variables(dmodel.model.controls)
     φ = DFun(dmodel.model.states, x0, vars; interp_mode=interp_mode)
-    return (;x0, x1, x2, r0, dx, J, φ)
+
+    tt = (;x0, x1, x2, r0, dx, J, φ)
+
+    return adapt(dest, tt)
 
 end
 
@@ -199,16 +202,14 @@ function time_iteration(model::DYModel,
     convergence = false
     iterations = T
 
-    if engine==:cpu
-        t_engine = CPU()
-    elseif engine==:gpu
-        t_engine = GPU()
+    (;x0, x1, x2, dx, r0, J, φ) = workspace
+
+    if engine in (:cpu, :gpu)
+        t_engine = get_backend(workspace.x0)
     else
         t_engine = nothing
     end
 
-    (;x0, x1, x2, dx, r0, J, φ) = workspace
-
     ti_trace = trace ? IterationTrace(typeof(φ)[]) : nothing
 
     if improve
@@ -224,7 +225,6 @@ function time_iteration(model::DYModel,
         trace && push!(ti_trace.data, deepcopy(φ))
 
         F!(r0, model, x0, φ, t_engine)
-
         # r0 = F(model, x0, φ)
 
         ε = norm(r0)

src/algos/time_iteration_accelerated.jl

@@ -37,14 +37,12 @@ using KernelAbstractions
 # end
 
 # This is for a PGrid model only 
-function F!(r, model, x, φ, engine::Union{CPU, GPU})
+function F!(r, model, x, φ, engine)
 
     @kernel function FF_(r, @Const(model), @Const(x), @Const(φ))
 
-        # c = @index(Global, Cartesian)
-        # c = @index(Global, Cartesian)
+
         n = @index(Global, Linear)
-        # i,j = c.I
         (i,j) = Dolo.from_linear(model.grid, n)
 
         s_ = model.grid[n]

src/grids.jl

@@ -15,7 +15,7 @@ size(cg::CGrid{d}) where d = tuple((e[3] for e in cg.ranges)... )
 
 getindex(g::CGrid{d}, inds::Vararg{Int64,d}) where d = SVector{d}(
     ( 
-        ( g.ranges[i][1] + (g.ranges[i][2]-g.ranges[i][1])*( (inds[i]-1)/(g.ranges[i][3]-1)) )
+        ( g.ranges[i][1] + (g.ranges[i][2]-g.ranges[i][1])*( (inds[i]-1)/(g.ranges[i][3]-convert(eltype(eltype(g)),1))) )
         for i=1:d
     )
 )
@@ -85,13 +85,17 @@ cross(g1::SGrid{d1}, g2::CGrid{d2}) where d1 where d2 = PGrid(g1,g2)
 
 import Base: getindex
 
-from_linear(g::PGrid{G1, G2, d}, n) where G1 where G2 where d = let x=divrem(n-1, length(g.grids[1])); (x[2]+1, x[1]+1) end
+@inline from_linear(g::PGrid{G1, G2, d}, n::Int) where G1 where G2 where d = let x=divrem(n-1, length(g.grids[1])); (x[2]+1, x[1]+1) end
 
-getindex(g::PGrid{G1, G2, d}, n::Int) where G1 where G2 where d = getindex(g, from_linear(g, n)...)
+@inline getindex(g::PGrid{G1, G2, d}, n::Int) where G1 where G2 where d = getindex(g, from_linear(g, n)...)
 
-function getindex(g::PGrid{G1, G2, d}, i::Int64, j::Int64) where G1<:SGrid{d1} where G2<:CGrid{d2} where d where d1 where d2
-    Tf = eltype(g)
-    SVector{d,Tf}(g.grids[1][i]..., g.grids[2][j]...)
+# eltype(cg::CGrid{d,Tf}) where d where Tf = SVector{d, Tf}
+
+
+@inline function getindex(g::PGrid{G1, G2, d}, i::Int64, j::Int64) where G1<:SGrid{d1} where G2<:CGrid{d2} where d where d1 where d2
+    # Tf = eltype(g)
+    # SVector{d,Tf}(g.grids[1][i]..., g.grids[2][j]...)
+    SVector(g.grids[1][i]..., g.grids[2][j]...)
 end
 
 

src/splines/csplines.jl

@@ -4,29 +4,29 @@ using StaticArrays
 
 
 # old style call
-function eval_UC_spline(a, b, orders, C::AbstractArray{Float64}, S::Matrix{Float64})
+function eval_UC_spline(a, b, orders, C::AbstractArray{Tf}, S::Matrix{Tf}) where Tf
     d,N = size(S)
 
     n_x = size(C,1)
     dims = tuple(size(C)[2:end]...)
     out = zeros(n_x,N)
-    CC = reshape(reinterpret(SVector{n_x,Float64},vec(C)),dims)
-    SS = reshape(reinterpret(SVector{d,Float64},vec(S)),(N,))
-    V  = reshape(reinterpret(SVector{n_x,Float64},vec(out)),(N,))
+    CC = reshape(reinterpret(SVector{n_x,Tf},vec(C)),dims)
+    SS = reshape(reinterpret(SVector{d,Tf},vec(S)),(N,))
+    V  = reshape(reinterpret(SVector{n_x,Tf},vec(out)),(N,))
     eval_UC_spline!(a, b, orders, CC, SS, V)
     return out
 end
 
 
-function eval_UC_spline(a, b, orders, C::AbstractArray{T, d}, S::Vector{SVector{d,Float64}}) where T where d
+function eval_UC_spline(a, b, orders, C::AbstractArray{T, d}, S::Vector{SVector{d,Tf}}) where T where d where Tf
     N = length(S)
     vals = zeros(T,N)
     eval_UC_spline!(a, b, orders, C, S, vals)
     return vals
 end
 
 
-@generated function eval_UC_spline!(a, b, orders, C::AbstractArray{T, d}, S::AbstractVector{SVector{d, Float64}}, V::AbstractVector{T}) where d where T
+@generated function eval_UC_spline!(a, b, orders, C::AbstractArray{T, d}, S::AbstractVector{SVector{d,Tf}}, V::AbstractVector{T}) where d where T where Tf
     # d = C.parameters[2]-1 # first dimension of C indexes the splines
     # the speed penalty of extrapolating points when iterating over point
     # seems very small so this is the default

src/splines/interp.jl

@@ -58,7 +58,8 @@ function interp(ranges::NTuple{d, Tuple{Tf, Tf, Int64}}, values::AbstractArray{T
 
     λ = (x.-(a .+ δ.*i))./δ
 
-    i_ = floor.(Int, i) .+ 1
+    # i_ = floor.(Int, i) .+ 1
+    i_ = unsafe_trunc.(Int, i) .+ 1
 
     M = matextract(values, i_...)