Create CUDA extension

Project.toml

@@ -28,7 +28,11 @@ StatsBase = "0.32, 0.33, 0.34"
 Tables = "1.9"
 julia = "1.6"
 
+[extensions]
+EvoTreesCUDAExt = "CUDA"
+
 [extras]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
@@ -37,4 +41,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 docs = ["Documenter"]
-test = ["DataFrames", "Test", "MLJBase", "MLJTestInterface"]
+test = ["CUDA", "DataFrames", "Test", "MLJBase", "MLJTestInterface"]
+
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

ext/EvoTreesCUDAExt/EvoTreesCUDAExt.jl

@@ -0,0 +1,22 @@
+module EvoTreesCUDAExt
+
+using EvoTrees
+using CUDA
+
+# This should be different on CPUs and GPUs
+EvoTrees.device_ones(::Type{<:GPU}, ::Type{T}, n::Int) where {T} = CUDA.ones(T, n)
+EvoTrees.device_array_type(::Type{<:GPU}) = CuArray
+function EvoTrees.post_fit_gc(::Type{<:GPU})
+    GC.gc(true)
+    CUDA.reclaim()
+end
+
+include("loss.jl")
+include("eval.jl")
+include("predict.jl")
+include("init.jl")
+include("subsample.jl")
+include("fit-utils.jl")
+include("fit.jl")
+
+end # module

src/gpu/eval.jl → ext/EvoTreesCUDAExt/eval.jl

@@ -8,7 +8,7 @@ function eval_mse_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y::CuDe
     end
     return nothing
 end
-function mse(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.mse(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_mse_kernel!(eval, p, y, w)
@@ -19,8 +19,8 @@ end
 ########################
 # RMSE
 ########################
-rmse(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat} =
-    sqrt(rmse(p, y, w; MAX_THREADS, kwargs...))
+EvoTrees.rmse(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat} =
+    sqrt(EvoTrees.rmse(p, y, w; MAX_THREADS, kwargs...))
 
 ########################
 # MAE
@@ -32,7 +32,7 @@ function eval_mae_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y::CuDe
     end
     return nothing
 end
-function mae(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.mae(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_mae_kernel!(eval, p, y, w)
@@ -51,7 +51,7 @@ function eval_logloss_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y::
     end
     return nothing
 end
-function logloss(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.logloss(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_logloss_kernel!(eval, p, y, w)
@@ -70,7 +70,7 @@ function eval_gaussian_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y:
     end
     return nothing
 end
-function gaussian_mle(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.gaussian_mle(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_gaussian_kernel!(eval, p, y, w)
@@ -91,7 +91,7 @@ function eval_poisson_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y::
     return nothing
 end
 
-function poisson(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.poisson(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_poisson_kernel!(eval, p, y, w)
@@ -111,7 +111,7 @@ function eval_gamma_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y::Cu
     return nothing
 end
 
-function gamma(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.gamma(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_gamma_kernel!(eval, p, y, w)
@@ -133,7 +133,7 @@ function eval_tweedie_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y::
     return nothing
 end
 
-function tweedie(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.tweedie(p::CuMatrix{T}, y::CuVector{T}, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_tweedie_kernel!(eval, p, y, w)
@@ -158,10 +158,10 @@ function eval_mlogloss_kernel!(eval::CuDeviceVector{T}, p::CuDeviceMatrix{T}, y:
     return nothing
 end
 
-function mlogloss(p::CuMatrix{T}, y::CuVector, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
+function EvoTrees.mlogloss(p::CuMatrix{T}, y::CuVector, w::CuVector{T}, eval::CuVector{T}; MAX_THREADS=1024, kwargs...) where {T<:AbstractFloat}
     threads = min(MAX_THREADS, length(y))
     blocks = cld(length(y), threads)
     @cuda blocks = blocks threads = threads eval_mlogloss_kernel!(eval, p, y, w)
     CUDA.synchronize()
     return sum(eval) / sum(w)
-end
+end

src/gpu/fit.jl → ext/EvoTreesCUDAExt/fit.jl

@@ -1,15 +1,15 @@
-function grow_evotree!(evotree::EvoTree{L,K}, cache, params::EvoTypes{L}, ::Type{GPU}) where {L,K}
+function EvoTrees.grow_evotree!(evotree::EvoTree{L,K}, cache, params::EvoTrees.EvoTypes{L}, ::Type{GPU}) where {L,K}
 
     # compute gradients
-    update_grads!(cache.∇, cache.pred, cache.y, params)
+    EvoTrees.update_grads!(cache.∇, cache.pred, cache.y, params)
     # subsample rows
     cache.nodes[1].is =
-        subsample(cache.is_in, cache.is_out, cache.mask, params.rowsample, params.rng)
+        EvoTrees.subsample(cache.is_in, cache.is_out, cache.mask, params.rowsample, params.rng)
     # subsample cols
-    sample!(params.rng, cache.js_, cache.js, replace=false, ordered=true)
+    EvoTrees.sample!(params.rng, cache.js_, cache.js, replace=false, ordered=true)
 
     # assign a root and grow tree
-    tree = Tree{L,K}(params.max_depth)
+    tree = EvoTrees.Tree{L,K}(params.max_depth)
     grow! = params.tree_type == "oblivious" ? grow_otree! : grow_tree!
     grow!(
         tree,
@@ -34,9 +34,9 @@ end
 
 # grow a single binary tree - grow through all depth
 function grow_tree!(
-    tree::Tree{L,K},
+    tree::EvoTrees.Tree{L,K},
     nodes::Vector{N},
-    params::EvoTypes{L},
+    params::EvoTrees.EvoTypes{L},
     ∇::CuMatrix,
     edges,
     js,
@@ -66,7 +66,7 @@ function grow_tree!(
 
     # initialize summary stats
     nodes[1].∑ .= Vector(vec(sum(∇[:, nodes[1].is], dims=2)))
-    nodes[1].gain = get_gain(params, nodes[1].∑) # should use a GPU version?
+    nodes[1].gain = EvoTrees.get_gain(params, nodes[1].∑) # should use a GPU version?
 
     # grow while there are remaining active nodes
     while length(n_current) > 0 && depth <= params.max_depth
@@ -91,13 +91,13 @@ function grow_tree!(
                 end
             end
             @threads for n ∈ sort(n_current)
-                update_gains!(nodes[n], js, params, feattypes, monotone_constraints)
+                EvoTrees.update_gains!(nodes[n], js, params, feattypes, monotone_constraints)
             end
         end
 
         for n ∈ sort(n_current)
             if depth == params.max_depth || nodes[n].∑[end] <= params.min_weight
-                pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
+                EvoTrees.pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
             else
                 best = findmax(findmax.(nodes[n].gains))
                 best_gain = best[1][1]
@@ -126,8 +126,8 @@ function grow_tree!(
                     nodes[n<<1].is, nodes[n<<1+1].is = _left, _right
                     nodes[n<<1].∑ .= nodes[n].hL[best_feat][:, best_bin]
                     nodes[n<<1+1].∑ .= nodes[n].hR[best_feat][:, best_bin]
-                    nodes[n<<1].gain = get_gain(params, nodes[n<<1].∑)
-                    nodes[n<<1+1].gain = get_gain(params, nodes[n<<1+1].∑)
+                    nodes[n<<1].gain = EvoTrees.get_gain(params, nodes[n<<1].∑)
+                    nodes[n<<1+1].gain = EvoTrees.get_gain(params, nodes[n<<1+1].∑)
 
                     if length(_right) >= length(_left)
                         push!(n_next, n << 1)
@@ -137,7 +137,7 @@ function grow_tree!(
                         push!(n_next, n << 1)
                     end
                 else
-                    pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
+                    EvoTrees.pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
                 end
             end
         end
@@ -151,9 +151,9 @@ end
 
 # grow a single oblivious tree - grow through all depth
 function grow_otree!(
-    tree::Tree{L,K},
+    tree::EvoTrees.Tree{L,K},
     nodes::Vector{N},
-    params::EvoTypes{L},
+    params::EvoTrees.EvoTypes{L},
     ∇::CuMatrix,
     edges,
     js,
@@ -183,7 +183,7 @@ function grow_otree!(
 
     # initialize summary stats
     nodes[1].∑ .= Vector(vec(sum(∇[:, nodes[1].is], dims=2)))
-    nodes[1].gain = get_gain(params, nodes[1].∑) # should use a GPU version?
+    nodes[1].gain = EvoTrees.get_gain(params, nodes[1].∑) # should use a GPU version?
 
     # grow while there are remaining active nodes
     while length(n_current) > 0 && depth <= params.max_depth
@@ -197,7 +197,7 @@ function grow_otree!(
         if depth == params.max_depth || min_weight_flag
             for n in n_current
                 # @info "length(nodes[n].is)" length(nodes[n].is) depth n
-                pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
+                EvoTrees.pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
             end
         else
             # update histograms
@@ -218,7 +218,7 @@ function grow_otree!(
                 end
             end
             @threads for n ∈ n_current
-                update_gains!(nodes[n], js, params, feattypes, monotone_constraints)
+                EvoTrees.update_gains!(nodes[n], js, params, feattypes, monotone_constraints)
             end
 
             # initialize gains for node 1 in which all gains of a given depth will be accumulated
@@ -273,8 +273,8 @@ function grow_otree!(
                     nodes[n<<1].is, nodes[n<<1+1].is = _left, _right
                     nodes[n<<1].∑ .= nodes[n].hL[best_feat][:, best_bin]
                     nodes[n<<1+1].∑ .= nodes[n].hR[best_feat][:, best_bin]
-                    nodes[n<<1].gain = get_gain(params, nodes[n<<1].∑)
-                    nodes[n<<1+1].gain = get_gain(params, nodes[n<<1+1].∑)
+                    nodes[n<<1].gain = EvoTrees.get_gain(params, nodes[n<<1].∑)
+                    nodes[n<<1+1].gain = EvoTrees.get_gain(params, nodes[n<<1+1].∑)
 
                     if length(_right) >= length(_left)
                         push!(n_next, n << 1)
@@ -286,7 +286,7 @@ function grow_otree!(
                 end
             else
                 for n in n_current
-                    pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
+                    EvoTrees.pred_leaf_cpu!(tree.pred, n, nodes[n].∑, params, ∇, nodes[n].is)
                 end
             end
         end
@@ -295,4 +295,4 @@ function grow_otree!(
     end # end of loop over current nodes for a given depth
 
     return nothing
-end
+end

src/gpu/init.jl → ext/EvoTreesCUDAExt/init.jl

@@ -1,56 +1,56 @@
-function init_core(params::EvoTypes{L}, ::Type{GPU}, data, fnames, y_train, w, offset) where {L}
+function EvoTrees.init_core(params::EvoTrees.EvoTypes{L}, ::Type{EvoTrees.GPU}, data, fnames, y_train, w, offset) where {L}
 
     # binarize data into quantiles
-    edges, featbins, feattypes = get_edges(data; fnames, nbins=params.nbins, rng=params.rng)
-    x_bin = CuArray(binarize(data; fnames, edges))
+    edges, featbins, feattypes = EvoTrees.get_edges(data; fnames, nbins=params.nbins, rng=params.rng)
+    x_bin = CuArray(EvoTrees.binarize(data; fnames, edges))
     nobs, nfeats = size(x_bin)
     T = Float32
 
     target_levels = nothing
-    if L == Logistic
+    if L == EvoTrees.Logistic
         @assert eltype(y_train) <: Real && minimum(y_train) >= 0 && maximum(y_train) <= 1
         K = 1
         y = T.(y_train)
-        μ = [logit(mean(y))]
-        !isnothing(offset) && (offset .= logit.(offset))
-    elseif L in [Poisson, Gamma, Tweedie]
+        μ = [EvoTrees.logit(EvoTrees.mean(y))]
+        !isnothing(offset) && (offset .= EvoTrees.logit.(offset))
+    elseif L in [EvoTrees.Poisson, EvoTrees.Gamma, EvoTrees.Tweedie]
         @assert eltype(y_train) <: Real
         K = 1
         y = T.(y_train)
-        μ = fill(log(mean(y)), 1)
+        μ = fill(log(EvoTrees.mean(y)), 1)
         !isnothing(offset) && (offset .= log.(offset))
-    elseif L == MLogLoss
-        if eltype(y_train) <: CategoricalValue
-            target_levels = CategoricalArrays.levels(y_train)
-            y = UInt32.(CategoricalArrays.levelcode.(y_train))
+    elseif L == EvoTrees.MLogLoss
+        if eltype(y_train) <: EvoTrees.CategoricalValue
+            target_levels = EvoTrees.CategoricalArrays.levels(y_train)
+            y = UInt32.(EvoTrees.CategoricalArrays.levelcode.(y_train))
         elseif eltype(y_train) <: Integer || eltype(y_train) <: Bool || eltype(y_train) <: String || eltype(y_train) <: Char
             target_levels = sort(unique(y_train))
-            yc = CategoricalVector(y_train, levels=target_levels)
-            y = UInt32.(CategoricalArrays.levelcode.(yc))
+            yc = EvoTrees.CategoricalVector(y_train, levels=target_levels)
+            y = UInt32.(EvoTrees.CategoricalArrays.levelcode.(yc))
         else
             @error "Invalid target eltype: $(eltype(y_train))"
         end
         K = length(target_levels)
-        μ = T.(log.(proportions(y, UInt32(1):UInt32(K))))
+        μ = T.(log.(EvoTrees.proportions(y, UInt32(1):UInt32(K))))
         μ .-= maximum(μ)
         !isnothing(offset) && (offset .= log.(offset))
-    elseif L == GaussianMLE
+    elseif L == EvoTrees.GaussianMLE
         @assert eltype(y_train) <: Real
         K = 2
         y = T.(y_train)
-        μ = [mean(y), log(std(y))]
+        μ = [EvoTrees.mean(y), log(EvoTrees.std(y))]
         !isnothing(offset) && (offset[:, 2] .= log.(offset[:, 2]))
-    elseif L == LogisticMLE
+    elseif L == EvoTrees.LogisticMLE
         @assert eltype(y_train) <: Real
         K = 2
         y = T.(y_train)
-        μ = [mean(y), log(std(y) * sqrt(3) / π)]
+        μ = [EvoTrees.mean(y), log(EvoTrees.std(y) * sqrt(3) / π)]
         !isnothing(offset) && (offset[:, 2] .= log.(offset[:, 2]))
     else
         @assert eltype(y_train) <: Real
         K = 1
         y = T.(y_train)
-        μ = [mean(y)]
+        μ = [EvoTrees.mean(y)]
     end
     y = CuArray(y)
     μ = T.(μ)
@@ -94,8 +94,8 @@ function init_core(params::EvoTypes{L}, ::Type{GPU}, data, fnames, y_train, w, o
     )
 
     # initialize model
-    nodes = [TrainNode(featbins, K, view(is_in, 1:0)) for n = 1:2^params.max_depth-1]
-    bias = [Tree{L,K}(μ)]
+    nodes = [EvoTrees.TrainNode(featbins, K, view(is_in, 1:0)) for n = 1:2^params.max_depth-1]
+    bias = [EvoTrees.Tree{L,K}(μ)]
     m = EvoTree{L,K}(bias, info)
 
     # build cache
@@ -125,4 +125,4 @@ function init_core(params::EvoTypes{L}, ::Type{GPU}, data, fnames, y_train, w, o
         monotone_constraints=monotone_constraints,
     )
     return m, cache
-end
+end