wip: introduce special assignment operator

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -12,6 +12,7 @@ using DispatchDoctor: @stable, @unstable
     include("NodePreallocation.jl")
     include("Strings.jl")
     include("Evaluate.jl")
+    include("SpecialOperators.jl")
     include("EvaluateDerivative.jl")
     include("ChainRules.jl")
     include("EvaluationHelpers.jl")
@@ -76,6 +77,7 @@ import .StringsModule: get_op_name
 @reexport import .EvaluateModule:
     eval_tree_array, differentiable_eval_tree_array, EvalOptions
 import .EvaluateModule: ArrayBuffer
+@reexport import .SpecialOperatorsModule: AssignOperator
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
 @reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!

src/Evaluate.jl

@@ -218,6 +218,10 @@ function eval_tree_array(
             "Bumper and LoopVectorization features are only compatible with numeric element types",
         )
     end
+    if any_special_operators(typeof(operators))
+        cX = copy(cX)
+        # TODO: This is dangerous if the element type is mutable
+    end
     if _eval_options.bumper isa Val{true}
         return bumper_eval_tree_array(tree, cX, operators, _eval_options)
     end
@@ -329,22 +333,26 @@ end
     long_compilation_time = nbin > OPERATOR_LIMIT_BEFORE_SLOWDOWN
     if long_compilation_time
         return quote
+            op = operators.binops[op_idx]
+            special_operator(op) && return deg2_eval_special(tree, cX, op, eval_options)
             result_l = _eval_tree_array(tree.l, cX, operators, eval_options)
             !result_l.ok && return result_l
             @return_on_nonfinite_array(eval_options, result_l.x)
             result_r = _eval_tree_array(tree.r, cX, operators, eval_options)
             !result_r.ok && return result_r
             @return_on_nonfinite_array(eval_options, result_r.x)
             # op(x, y), for any x or y
-            deg2_eval(result_l.x, result_r.x, operators.binops[op_idx], eval_options)
+            deg2_eval(result_l.x, result_r.x, op, eval_options)
         end
     end
     return quote
         return Base.Cartesian.@nif(
             $nbin,
             i -> i == op_idx,
             i -> let op = operators.binops[i]
-                if tree.l.degree == 0 && tree.r.degree == 0
+                if special_operator(op)
+                    deg2_eval_special(tree, cX, op, eval_options)
+                elseif tree.l.degree == 0 && tree.r.degree == 0
                     deg2_l0_r0_eval(tree, cX, op, eval_options)
                 elseif tree.r.degree == 0
                     result_l = _eval_tree_array(tree.l, cX, operators, eval_options)
@@ -380,13 +388,16 @@ end
     eval_options::EvalOptions,
 ) where {T}
     nuna = get_nuna(operators)
+    special_operators = any_special_operators(operators)
     long_compilation_time = nuna > OPERATOR_LIMIT_BEFORE_SLOWDOWN
     if long_compilation_time
         return quote
+            op = operators.unaops[op_idx]
+            special_operator(op) && return deg1_eval_special(tree, cX, op, eval_options)
             result = _eval_tree_array(tree.l, cX, operators, eval_options)
             !result.ok && return result
             @return_on_nonfinite_array(eval_options, result.x)
-            deg1_eval(result.x, operators.unaops[op_idx], eval_options)
+            deg1_eval(result.x, op, eval_options)
         end
     end
     # This @nif lets us generate an if statement over choice of operator,
@@ -396,13 +407,18 @@ end
             $nuna,
             i -> i == op_idx,
             i -> let op = operators.unaops[i]
-                if tree.l.degree == 2 && tree.l.l.degree == 0 && tree.l.r.degree == 0
+                if special_operator(op)
+                    deg1_eval_special(tree, cX, op, eval_options)
+                elseif !special_operators &&
+                    tree.l.degree == 2 &&
+                    tree.l.l.degree == 0 &&
+                    tree.l.r.degree == 0
                     # op(op2(x, y)), where x, y, z are constants or variables.
                     l_op_idx = tree.l.op
                     dispatch_deg1_l2_ll0_lr0_eval(
                         tree, cX, op, l_op_idx, operators.binops, eval_options
                     )
-                elseif tree.l.degree == 1 && tree.l.l.degree == 0
+                elseif !special_operators && tree.l.degree == 1 && tree.l.l.degree == 0
                     # op(op2(x)), where x is a constant or variable.
                     l_op_idx = tree.l.op
                     dispatch_deg1_l1_ll0_eval(
@@ -925,4 +941,10 @@ end
     end
 end
 
+# Overloaded by SpecialOperators.jl:
+function any_special_operators end
+function special_operator end
+function deg2_eval_special end
+function deg1_eval_special end
+
 end

src/SpecialOperators.jl

@@ -0,0 +1,57 @@
+module SpecialOperatorsModule
+
+using ..OperatorEnumModule: OperatorEnum
+using ..EvaluateModule: _eval_tree_array, @return_on_nonfinite_array, deg2_eval
+
+import ..EvaluateModule:
+    special_operator, deg2_eval_special, deg1_eval_special, any_special_operators
+
+function any_special_operators(::Type{OperatorEnum{B,U}}) where {B,U}
+    return any(special_operator, B.types) || any(special_operator, U.types)
+end
+
+# Use this to customize evaluation behavior for operators:
+@inline special_operator(::Type) = false
+@inline special_operator(f) = special_operator(typeof(f))
+
+# Base.@kwdef struct WhileOperator <: Function
+#     max_iters::Int = 100
+# end
+Base.@kwdef struct AssignOperator <: Function
+    target_register::Int
+end
+
+# @inline special_operator(::Type{WhileOperator}) = true
+@inline special_operator(::Type{AssignOperator}) = true
+
+# function deg2_eval_special(tree, cX, op::WhileOperator, eval_options)
+#     cond = tree.l
+#     body = tree.r
+#     for _ in 1:(op.max_iters)
+#         let cond_result = _eval_tree_array(cond, cX, operators, eval_options)
+#             !cond_result.ok && return cond_result
+#             @return_on_nonfinite_array(eval_options, cond_result.x)
+#         end
+#         let body_result = _eval_tree_array(body, cX, operators, eval_options)
+#             !body_result.ok && return body_result
+#             @return_on_nonfinite_array(eval_options, body_result.x)
+#             # TODO: Need to somehow mask instances 
+#         end
+#     end
+
+#     return get_filled_array(eval_options.buffer, zero(eltype(cX)), cX, axes(cX, 2))
+# end
+# TODO: Need to void any instance of buffer when using while loop.
+
+function deg1_eval_special(tree, cX, op::AssignOperator, eval_options)
+    result = _eval_tree_array(tree.l, cX, operators, eval_options)
+    !result.ok && return result
+    @return_on_nonfinite_array(eval_options, result.x)
+    target_register = op.target_register
+    @inbounds @simd for i in eachindex(axes(cX, 2))
+        cX[target_register, i] = result.x[i]
+    end
+    return result
+end
+
+end