Permit n-argument operators

ext/DynamicExpressionsBumperExt.jl

@@ -5,8 +5,7 @@ using DynamicExpressions:
     OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array, EvalOptions
 using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
-import DynamicExpressions.ExtensionInterfaceModule:
-    bumper_eval_tree_array, bumper_kern1!, bumper_kern2!
+import DynamicExpressions.ExtensionInterfaceModule: bumper_eval_tree_array, bumper_kern!
 
 function bumper_eval_tree_array(
     tree::AbstractExpressionNode{T},
@@ -37,8 +36,7 @@ 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..., eval_options),
+            KernelDispatcher(operators, eval_options),
             tree;
             break_sharing=Val(true),
         )
@@ -49,63 +47,44 @@ function bumper_eval_tree_array(
     return (result, all_ok[])
 end
 
-function dispatch_kerns!(
-    operators, branch_node, cumulator, eval_options::EvalOptions{<:Any,true,early_exit}
-) where {early_exit}
-    cumulator.ok || return cumulator
-
-    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, eval_options)
-    return ResultOk(out, early_exit ? is_valid_array(out) : true)
-end
-function dispatch_kerns!(
-    operators,
-    branch_node,
-    cumulator1,
-    cumulator2,
-    eval_options::EvalOptions{<:Any,true,early_exit},
-) where {early_exit}
-    cumulator1.ok || return cumulator1
-    cumulator2.ok || return cumulator2
-
-    out = dispatch_kern2!(
-        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, eval_options
-    )
-    return ResultOk(out, early_exit ? is_valid_array(out) : true)
+struct KernelDispatcher{O<:OperatorEnum,E<:EvalOptions{<:Any,true,<:Any}} <: Function
+    operators::O
+    eval_options::E
 end
 
-@generated function dispatch_kern1!(unaops, op_idx, cumulator, eval_options::EvalOptions)
-    nuna = counttuple(unaops)
+@generated function (kd::KernelDispatcher{<:Any,<:EvalOptions{<:Any,true,early_exit}})(
+    branch_node, inputs::Vararg{Any,degree}
+) where {degree,early_exit}
     quote
-        Base.@nif(
-            $nuna,
-            i -> i == op_idx,
-            i -> let op = unaops[i]
-                return bumper_kern1!(op, cumulator, eval_options)
-            end,
-        )
+        Base.Cartesian.@nexprs($degree, i -> inputs[i].ok || return inputs[i])
+        cumulators = Base.Cartesian.@ntuple($degree, i -> inputs[i].x)
+        out = dispatch_kerns!(kd.operators, branch_node, cumulators, kd.eval_options)
+        return ResultOk(out, early_exit ? is_valid_array(out) : true)
     end
 end
-@generated function dispatch_kern2!(
-    binops, op_idx, cumulator1, cumulator2, eval_options::EvalOptions
-)
-    nbin = counttuple(binops)
+@generated function dispatch_kerns!(
+    operators::OperatorEnum{OPS},
+    branch_node,
+    cumulators::Tuple{Vararg{Any,degree}},
+    eval_options::EvalOptions,
+) where {OPS,degree}
+    nops = length(OPS.types[degree].types)
     quote
-        Base.@nif(
-            $nbin,
+        op_idx = branch_node.op
+        Base.Cartesian.@nif(
+            $nops,
             i -> i == op_idx,
-            i -> let op = binops[i]
-                return bumper_kern2!(op, cumulator1, cumulator2, eval_options)
-            end,
+            i -> bumper_kern!(operators[$degree][i], cumulators, eval_options)
         )
     end
 end
-function bumper_kern1!(op::F, cumulator, ::EvalOptions{false,true}) where {F}
-    @. cumulator = op(cumulator)
-    return cumulator
-end
-function bumper_kern2!(op::F, cumulator1, cumulator2, ::EvalOptions{false,true}) where {F}
-    @. cumulator1 = op(cumulator1, cumulator2)
-    return cumulator1
+
+function bumper_kern!(
+    op::F, cumulators::Tuple{Vararg{Any,degree}}, ::EvalOptions{false,true,early_exit}
+) where {F,degree,early_exit}
+    cumulator_1 = first(cumulators)
+    @. cumulator_1 = op(cumulators...)
+    return cumulator_1
 end
 
 end

ext/DynamicExpressionsLoopVectorizationExt.jl

@@ -8,37 +8,28 @@ using DynamicExpressions.EvaluateModule:
 import DynamicExpressions.EvaluateModule:
     deg1_eval,
     deg2_eval,
+    degn_eval,
     deg1_l2_ll0_lr0_eval,
     deg1_l1_ll0_eval,
     deg2_l0_r0_eval,
     deg2_l0_eval,
     deg2_r0_eval
 import DynamicExpressions.ExtensionInterfaceModule:
-    _is_loopvectorization_loaded, bumper_kern1!, bumper_kern2!
+    _is_loopvectorization_loaded, bumper_kern!
 
 _is_loopvectorization_loaded(::Int) = true
 
-function deg2_eval(
-    cumulator_l::AbstractVector{T},
-    cumulator_r::AbstractVector{T},
-    op::F,
-    ::EvalOptions{true},
-)::ResultOk where {T<:Number,F}
-    @turbo for j in eachindex(cumulator_l)
-        x = op(cumulator_l[j], cumulator_r[j])
-        cumulator_l[j] = x
-    end
-    return ResultOk(cumulator_l, true)
-end
-
-function deg1_eval(
-    cumulator::AbstractVector{T}, op::F, ::EvalOptions{true}
-)::ResultOk where {T<:Number,F}
-    @turbo for j in eachindex(cumulator)
-        x = op(cumulator[j])
-        cumulator[j] = x
+@generated function degn_eval(
+    cumulators::NTuple{N,<:AbstractVector{T}}, op::F, ::EvalOptions{true}
+)::ResultOk where {N,T,F}
+    # Fast general implementation of `cumulators[1] .= op.(cumulators[1], cumulators[2], ...)`
+    quote
+        Base.Cartesian.@nexprs($N, i -> cumulator_i = cumulators[i])
+        @turbo for j in eachindex(cumulator_1)
+            cumulator_1[j] = Base.Cartesian.@ncall($N, op, i -> cumulator_i[j])
+        end
+        return ResultOk(cumulator_1, true)
     end
-    return ResultOk(cumulator, true)
 end
 
 function deg1_l2_ll0_lr0_eval(
@@ -217,18 +208,13 @@ function deg2_r0_eval(
     end
 end
 
-## Interface with Bumper.jl
-function bumper_kern1!(
-    op::F, cumulator, ::EvalOptions{true,true,early_exit}
-) where {F,early_exit}
-    @turbo @. cumulator = op(cumulator)
-    return cumulator
-end
-function bumper_kern2!(
-    op::F, cumulator1, cumulator2, ::EvalOptions{true,true,early_exit}
-) where {F,early_exit}
-    @turbo @. cumulator1 = op(cumulator1, cumulator2)
-    return cumulator1
+# Interface with Bumper.jl
+function bumper_kern!(
+    op::F, cumulators::Tuple{Vararg{Any,degree}}, ::EvalOptions{true,true,early_exit}
+) where {F,degree,early_exit}
+    cumulator_1 = first(cumulators)
+    @turbo @. cumulator_1 = op(cumulators...)
+    return cumulator_1
 end
 
 end

ext/DynamicExpressionsSymbolicUtilsExt.jl

@@ -175,9 +175,13 @@ function Base.convert(
         findoperation(op, operators.unaops)
     end
 
-    return constructorof(N)(;
-        op=ind, children=map(x -> convert(N, x, operators; variable_names), args)
-    )
+    if length(args) == 2
+        children = map(x -> convert(N, x, operators; variable_names), (args[1], args[2]))
+        return constructorof(N)(; op=ind, children)
+    else
+        children = map(x -> convert(N, x, operators; variable_names), (only(args),))
+        return constructorof(N)(; op=ind, children)
+    end
 end
 
 _node_type(::Type{<:AbstractExpression{T,N}}) where {T,N<:AbstractExpressionNode} = N

ext/DynamicExpressionsZygoteExt.jl

@@ -3,30 +3,14 @@ module DynamicExpressionsZygoteExt
 using Zygote: gradient
 import DynamicExpressions.ExtensionInterfaceModule: _zygote_gradient, ZygoteGradient
 
-function _zygote_gradient(op::F, ::Val{1}) where {F}
-    return ZygoteGradient{F,1,1}(op)
-end
-function _zygote_gradient(op::F, ::Val{2}, ::Val{side}=Val(nothing)) where {F,side}
-    # side should be either nothing (for both), 1, or 2
-    @assert side === nothing || side in (1, 2)
-    return ZygoteGradient{F,2,side}(op)
+function _zygote_gradient(op::F, ::Val{degree}) where {F,degree}
+    return ZygoteGradient{F,degree}(op)
 end
 
-function (g::ZygoteGradient{F,1,1})(x) where {F}
-    out = only(gradient(g.op, x))
-    return out === nothing ? zero(x) : out
-end
-function (g::ZygoteGradient{F,2,nothing})(x, y) where {F}
-    (∂x, ∂y) = gradient(g.op, x, y)
-    return (∂x === nothing ? zero(x) : ∂x, ∂y === nothing ? zero(y) : ∂y)
-end
-function (g::ZygoteGradient{F,2,1})(x, y) where {F}
-    ∂x = only(gradient(Base.Fix2(g.op, y), x))
-    return ∂x === nothing ? zero(x) : ∂x
-end
-function (g::ZygoteGradient{F,2,2})(x, y) where {F}
-    ∂y = only(gradient(Base.Fix1(g.op, x), y))
-    return ∂y === nothing ? zero(y) : ∂y
+# All this does is remove `nothing`, so that we get type stability
+function (g::ZygoteGradient{F,degree})(args::Vararg{Any,degree}) where {F,degree}
+    partials = gradient(g.op, args...)
+    return ntuple(i -> @something(partials[i], zero(args[i])), Val(degree))
 end
 
 end

src/DynamicExpressions.jl

@@ -51,6 +51,8 @@ import .NodeModule:
     constructorof,
     with_type_parameters,
     preserve_sharing,
+    max_degree,
+    with_max_degree,
     leaf_copy,
     branch_copy,
     leaf_hash,