Use vectors for vector input; tuple for tuple

Author: MilesCranmer

ext/DynamicExpressionsCUDAExt.jl

@@ -29,15 +29,15 @@ function eval_tree_array(
 end
 
 function eval_tree_array(
-    trees::Tuple{N,Vararg{N,M}},
+    trees::Union{NTuple{M,N} where M,AbstractVector{N}},
     gcX::MaybeCuArray{T,2},
     operators::OperatorEnum;
     buffer=nothing,
     gpu_workspace=nothing,
     gpu_buffer=nothing,
     kws...,
-) where {T<:Number,N<:AbstractExpressionNode{T},M}
-    (; val, execution_order, roots, buffer, num_nodes) = as_array(Int32, trees...; buffer)
+) where {T<:Number,N<:AbstractExpressionNode{T}}
+    (; val, execution_order, roots, buffer, num_nodes) = as_array(Int32, trees; buffer)
     num_launches = maximum(execution_order)
     num_elem = size(gcX, 2)
 
@@ -82,11 +82,8 @@ function eval_tree_array(
     )
     #! format: on
 
-    out = ntuple(i -> @view(gworkspace[begin:(end - 1), roots[i]]), Val(M + 1))
-    is_good = ntuple(
-        i -> true,  # Up to user to find NaNs
-        Val(M + 1),
-    )
+    out = (r -> @view(gworkspace[begin:(end - 1), r])).(roots)
+    is_good = (_ -> true).(trees)
 
     return (out, is_good)
 end

src/AsArray.jl

@@ -4,15 +4,20 @@ using ..EquationModule: AbstractExpressionNode, tree_mapreduce, count_nodes
 
 function as_array(
     ::Type{I},
-    tree::N,
-    additional_trees::Vararg{N,M};
+    trees::Union{NTuple{M,N} where M,AbstractVector{N}};
     buffer::Union{AbstractArray,Nothing}=nothing,
-) where {T,N<:AbstractExpressionNode{T},I,M}
-    trees = (tree, additional_trees...)
+) where {T,N<:AbstractExpressionNode{T},I}
     each_num_nodes = (t -> count_nodes(t; break_sharing=Val(true))).(trees)
     num_nodes = sum(each_num_nodes)
 
-    roots = cumsum(tuple(one(I), each_num_nodes[1:(end - 1)]...))
+    # Want `roots` to be tuple if `trees` is tuple and similar for vector
+    roots = cumsum(
+        if each_num_nodes isa Tuple
+            tuple(one(I), each_num_nodes[1:(end - 1)]...)
+        else
+            vcat(one(I), each_num_nodes[1:(end - 1)])
+        end,
+    )
 
     val = Array{T}(undef, num_nodes)
 

src/EvaluateEquation.jl

@@ -98,10 +98,13 @@ function eval_tree_array(
     return eval_tree_array(tree, cX, operators; turbo, bumper)
 end
 function eval_tree_array(
-    trees::Tuple{N,Vararg{N,M}}, cX::AbstractMatrix{T}, operators::OperatorEnum; kws...
-) where {T<:Number,N<:AbstractExpressionNode{T},M}
-    outs = ntuple(i -> eval_tree_array(trees[i], cX, operators; kws...), Val(M + 1))
-    return ntuple(i -> first(outs[i]), Val(M + 1)), ntuple(i -> last(outs[i]), Val(M + 1))
+    trees::Union{NTuple{M,N} where M,AbstractArray{N}},
+    cX::AbstractMatrix{T},
+    operators::OperatorEnum;
+    kws...,
+) where {T<:Number,N<:AbstractExpressionNode{T}}
+    outs = (t -> eval_tree_array(t, cX, operators; kws...)).(trees)
+    return first.(outs), last.(outs)
 end
 
 get_nuna(::Type{<:OperatorEnum{B,U}}) where {B,U} = counttuple(U)

test/test_cuda.jl

@@ -21,27 +21,41 @@ let
 
         nrow = rand(10:30)
         nnodes = rand(10:25, ntrees)
+        use_tuple = rand(Bool)
 
         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)
+        trees = use_tuple ? trees : collect(trees)
         X = randn(T, 3, nrow)
         if ntrees > 1
-            y, completed = eval_tree_array(trees, X, operators)
-            gpu_y, gpu_completed = eval_tree_array(
+            y, completed = @inferred eval_tree_array(trees, X, operators)
+            gpu_y, gpu_completed = @inferred eval_tree_array(
                 trees, FakeCuArray(X), operators; buffer, gpu_workspace, gpu_buffer
             )
 
+            # Should give same result either way
             for i in eachindex(completed, gpu_completed)
                 if completed[i]
                     @test y[i] ≈ gpu_y[i]
                 end
             end
+
+            # Should return same type as input
+            if use_tuple
+                @test y isa Tuple
+                @test gpu_y isa Tuple
+            else
+                @test y isa Vector
+                @test gpu_y isa Vector
+            end
         else
-            y, completed = eval_tree_array(only(trees), X, operators)
-            gpu_y, gpu_completed = eval_tree_array(only(trees), FakeCuArray(X), operators)
+            y, completed = @inferred eval_tree_array(only(trees), X, operators)
+            gpu_y, gpu_completed = @inferred eval_tree_array(
+                only(trees), FakeCuArray(X), operators
+            )
             if completed
                 @test y ≈ gpu_y
             end