Test buffered-evaluation of CUDA kernels

Author: MilesCranmer

ext/DynamicExpressionsCUDAExt.jl

@@ -15,7 +15,6 @@ Base.similar(x::FakeCuArray, dims::Integer...) = FakeCuArray(similar(x.a, dims..
 Base.getindex(x::FakeCuArray, i::Int...) = getindex(x.a, i...)
 Base.setindex!(x::FakeCuArray, v, i::Int...) = setindex!(x.a, v, i...)
 Base.size(x::FakeCuArray) = size(x.a)
-Base.Array(x::FakeCuArray) = Array(x.a)
 
 const MaybeCuArray{T,N} = Union{CuArray{T,N},FakeCuArray{T,N}}
 

test/test_cuda.jl

@@ -16,21 +16,34 @@ let
     )
     x1, x2, x3 = (i -> Node(Float64; feature=i)).(1:3)
 
-    for T in (Float32, Float64, ComplexF64), num_trees in (1, 2, 3), seed in 0:10
+    for T in (Float32, Float64, ComplexF64), ntrees in (1, 2, 3), seed in 0:10
         Random.seed!(seed)
-        num_rows = rand(10:30)
-        nodes_per = rand(10:25, num_trees)
-        trees = ntuple(
-            i -> gen_random_tree_fixed_size(nodes_per[i], operators, 3, T), num_trees
-        )
-        X = randn(T, 3, num_rows)
-        y, completed = eval_tree_array(trees, X, operators)
-        gpu_y, gpu_completed = eval_tree_array(trees, FakeCuArray(X), operators)
-        gpu_y = Array.(gpu_y)
-
-        for i in eachindex(completed, gpu_completed)
-            if completed[i]
-                @test y[i] ≈ gpu_y[i]
+
+        nrow = rand(10:30)
+        nnodes = rand(10:25, ntrees)
+
+        buffer = rand(Bool) ? ones(Int32, 8, sum(nnodes)) : nothing
+        gpu_buffer = rand(Bool) ? FakeCuArray(ones(Int32, 8, sum(nnodes))) : nothing
+        gpu_workspace = rand(Bool) ? FakeCuArray(ones(T, nrow + 1, sum(nnodes))) : nothing
+
+        trees = ntuple(i -> gen_random_tree_fixed_size(nnodes[i], operators, 3, T), ntrees)
+        X = randn(T, 3, nrow)
+        if ntrees > 1
+            y, completed = eval_tree_array(trees, X, operators)
+            gpu_y, gpu_completed = eval_tree_array(
+                trees, FakeCuArray(X), operators; buffer, gpu_workspace, gpu_buffer
+            )
+
+            for i in eachindex(completed, gpu_completed)
+                if completed[i]
+                    @test y[i] ≈ gpu_y[i]
+                end
+            end
+        else
+            y, completed = eval_tree_array(only(trees), X, operators)
+            gpu_y, gpu_completed = eval_tree_array(only(trees), FakeCuArray(X), operators)
+            if completed
+                @test y ≈ gpu_y
             end
         end
     end