Allow both Bumper and LoopVectorization

Author: MilesCranmer

benchmark/benchmarks.jl

@@ -36,7 +36,6 @@ function benchmark_evaluation()
         for turbo in (false, true), bumper in (false, true)
 
             (turbo || bumper) && !(T in (Float32, Float64)) && continue
-            turbo && bumper && continue
             if bumper
                 try
                     eval_tree_array(Node{T}(val=1.0), ones(T, 5, n), operators; turbo, bumper)
@@ -47,7 +46,15 @@ function benchmark_evaluation()
                 end
             end
 
-            extra_key = turbo ? "_turbo" : (bumper ? "_bumper" : "")
+            extra_key = if turbo && bumper
+                "_turbo_bumper"
+            elseif turbo
+                "_turbo"
+            elseif bumper
+                "_bumper"
+            else
+                ""
+            end
             extra_kws = bumper ? (; bumper=Val(true)) : ()
             eval_tree_array(
                 gen_random_tree_fixed_size(20, operators, 5, T),

ext/DynamicExpressionsBumperExt.jl

@@ -4,11 +4,15 @@ using Bumper: @no_escape, @alloc
 using DynamicExpressions: OperatorEnum, AbstractExpressionNode, tree_mapreduce
 using DynamicExpressions.UtilsModule: ResultOk, counttuple, is_bad_array
 
-import DynamicExpressions.ExtensionInterfaceModule: bumper_eval_tree_array
+import DynamicExpressions.ExtensionInterfaceModule:
+    bumper_eval_tree_array, bumper_kern1!, bumper_kern2!
 
 function bumper_eval_tree_array(
-    tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, operators::OperatorEnum
-) where {T}
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    operators::OperatorEnum,
+    ::Val{turbo},
+) where {T,turbo}
     result = similar(cX, axes(cX, 2))
     n = size(cX, 2)
     all_ok = Ref(false)
@@ -32,7 +36,8 @@ function bumper_eval_tree_array(
             branch_node -> branch_node,
             # In the evaluation kernel, we combine the branch nodes
             # with the arrays created by the leaf nodes:
-            ((args::Vararg{Any,M}) where {M}) -> dispatch_kerns!(operators, args...),
+            ((args::Vararg{Any,M}) where {M}) ->
+                dispatch_kerns!(operators, args..., Val(turbo)),
             tree;
             break_sharing=Val(true),
         )
@@ -43,43 +48,55 @@ function bumper_eval_tree_array(
     return (result, all_ok[])
 end
 
-function dispatch_kerns!(operators, branch_node, cumulator)
+function dispatch_kerns!(operators, branch_node, cumulator, ::Val{turbo}) where {turbo}
     cumulator.ok || return cumulator
 
-    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x)
+    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, Val(turbo))
     return ResultOk(out, !is_bad_array(out))
 end
-function dispatch_kerns!(operators, branch_node, cumulator1, cumulator2)
+function dispatch_kerns!(
+    operators, branch_node, cumulator1, cumulator2, ::Val{turbo}
+) where {turbo}
     cumulator1.ok || return cumulator1
     cumulator2.ok || return cumulator2
 
-    out = dispatch_kern2!(operators.binops, branch_node.op, cumulator1.x, cumulator2.x)
+    out = dispatch_kern2!(
+        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, Val(turbo)
+    )
     return ResultOk(out, !is_bad_array(out))
 end
 
-@generated function dispatch_kern1!(unaops, op_idx, cumulator)
+@generated function dispatch_kern1!(unaops, op_idx, cumulator, ::Val{turbo}) where {turbo}
     nuna = counttuple(unaops)
     quote
-        Base.@nif($nuna, i -> i == op_idx, i -> let op = unaops[i]
-            return kern1!(op, cumulator)
-        end,)
+        Base.@nif(
+            $nuna,
+            i -> i == op_idx,
+            i -> let op = unaops[i]
+                return bumper_kern1!(op, cumulator, Val(turbo))
+            end,
+        )
     end
 end
-@generated function dispatch_kern2!(binops, op_idx, cumulator1, cumulator2)
+@generated function dispatch_kern2!(
+    binops, op_idx, cumulator1, cumulator2, ::Val{turbo}
+) where {turbo}
     nbin = counttuple(binops)
     quote
         Base.@nif(
-            $nbin, i -> i == op_idx, i -> let op = binops[i]
-                return kern2!(op, cumulator1, cumulator2)
+            $nbin,
+            i -> i == op_idx,
+            i -> let op = binops[i]
+                return bumper_kern2!(op, cumulator1, cumulator2, Val(turbo))
             end,
         )
     end
 end
-function kern1!(op::F, cumulator) where {F}
+function bumper_kern1!(op::F, cumulator, ::Val{false}) where {F}
     @. cumulator = op(cumulator)
     return cumulator
 end
-function kern2!(op::F, cumulator1, cumulator2) where {F}
+function bumper_kern2!(op::F, cumulator1, cumulator2, ::Val{false}) where {F}
     @. cumulator1 = op(cumulator1, cumulator2)
     return cumulator1
 end

ext/DynamicExpressionsLoopVectorizationExt.jl

@@ -12,7 +12,8 @@ import DynamicExpressions.EvaluateEquationModule:
     deg2_l0_r0_eval,
     deg2_l0_eval,
     deg2_r0_eval
-import DynamicExpressions.ExtensionInterfaceModule: _is_loopvectorization_loaded
+import DynamicExpressions.ExtensionInterfaceModule:
+    _is_loopvectorization_loaded, bumper_kern1!, bumper_kern2!
 
 _is_loopvectorization_loaded(::Int) = true
 
@@ -201,4 +202,14 @@ function deg2_r0_eval(
     end
 end
 
+## Interface with Bumper.jl
+function bumper_kern1!(op::F, cumulator, ::Val{true}) where {F}
+    @turbo @. cumulator = op(cumulator)
+    return cumulator
+end
+function bumper_kern2!(op::F, cumulator1, cumulator2, ::Val{true}) where {F}
+    @turbo @. cumulator1 = op(cumulator1, cumulator2)
+    return cumulator1
+end
+
 end

src/EvaluateEquation.jl

@@ -72,14 +72,13 @@ function eval_tree_array(
     if v_turbo isa Val{true} || v_bumper isa Val{true}
         @assert T in (Float32, Float64)
     end
-    @assert !(v_turbo isa Val{true} && v_bumper isa Val{true})
-    if v_bumper isa Val{true}
-        return bumper_eval_tree_array(tree, cX, operators)
-    end
     if v_turbo isa Val{true}
         _is_loopvectorization_loaded(0) ||
             error("Please load the LoopVectorization.jl package to use this feature.")
     end
+    if v_bumper isa Val{true}
+        return bumper_eval_tree_array(tree, cX, operators, v_turbo)
+    end
 
     result = _eval_tree_array(tree, cX, operators, v_turbo)
     return (result.x, result.ok && !is_bad_array(result.x))

src/ExtensionInterface.jl

@@ -14,6 +14,8 @@ end
 function bumper_eval_tree_array(args...)
     return error("Please load the Bumper.jl package to use this feature.")
 end
+function bumper_kern1! end
+function bumper_kern2! end
 
 _is_loopvectorization_loaded(_) = false
 

test/test_evaluation.jl

@@ -43,7 +43,6 @@ for turbo in [Val(false), Val(true)],
 
     # Float16 not implemented:
     (turbo isa Val{true} || bumper isa Val{true}) && !(T in (Float32, Float64)) && continue
-    turbo isa Val{true} && bumper isa Val{true} && continue
     @testset "Test evaluation of trees with turbo=$turbo, bumper=$bumper, T=$T" begin
         for (i_func, fnc) in enumerate(functions)