feat: assume :terminates_locally for most loops

src/Evaluate.jl

@@ -6,7 +6,7 @@ import ..NodeModule:
     AbstractExpressionNode, constructorof, get_children, get_child, with_type_parameters
 import ..StringsModule: string_tree
 import ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum, GenericOperatorEnum
-import ..UtilsModule: fill_similar, counttuple, ResultOk
+import ..UtilsModule: fill_similar, counttuple, ResultOk, @finite
 import ..NodeUtilsModule: is_constant
 import ..ExtensionInterfaceModule: bumper_eval_tree_array, _is_loopvectorization_loaded
 import ..ValueInterfaceModule: is_valid, is_valid_array
@@ -295,7 +295,7 @@ end
     # Fast general implementation of `cumulators[1] .= op.(cumulators[1], cumulators[2], ...)`
     quote
         Base.Cartesian.@nexprs($N, i -> cumulator_i = cumulators[i])
-        @inbounds @simd for j in eachindex(cumulator_1)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator_1)
             cumulator_1[j] = Base.Cartesian.@ncall($N, op, i -> cumulator_i[j])::T
         end  # COV_EXCL_LINE
         return ResultOk(cumulator_1, true)
@@ -581,7 +581,7 @@ function deg1_l2_ll0_lr0_eval(
         @return_on_nonfinite_val(eval_options, val_ll, cX)
         feature_lr = get_child(get_child(tree, 1), 2).feature
         cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in axes(cX, 2)
             x_l = op_l(val_ll, cX[feature_lr, j])::T
             x = is_valid(x_l) ? op(x_l)::T : T(Inf)
             cumulator[j] = x
@@ -592,7 +592,7 @@ function deg1_l2_ll0_lr0_eval(
         val_lr = get_child(get_child(tree, 1), 2).val
         @return_on_nonfinite_val(eval_options, val_lr, cX)
         cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j], val_lr)::T
             x = is_valid(x_l) ? op(x_l)::T : T(Inf)
             cumulator[j] = x
@@ -602,7 +602,7 @@ function deg1_l2_ll0_lr0_eval(
         feature_ll = get_child(get_child(tree, 1), 1).feature
         feature_lr = get_child(get_child(tree, 1), 2).feature
         cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j], cX[feature_lr, j])::T
             x = is_valid(x_l) ? op(x_l)::T : T(Inf)
             cumulator[j] = x
@@ -630,7 +630,7 @@ function deg1_l1_ll0_eval(
     else
         feature_ll = get_child(get_child(tree, 1), 1).feature
         cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x_l = op_l(cX[feature_ll, j])::T
             x = is_valid(x_l) ? op(x_l)::T : T(Inf)
             cumulator[j] = x
@@ -659,7 +659,7 @@ function deg2_l0_r0_eval(
         val_l = get_child(tree, 1).val
         @return_on_nonfinite_val(eval_options, val_l, cX)
         feature_r = get_child(tree, 2).feature
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(val_l, cX[feature_r, j])::T
             cumulator[j] = x
         end  # COV_EXCL_LINE
@@ -669,7 +669,7 @@ function deg2_l0_r0_eval(
         feature_l = get_child(tree, 1).feature
         val_r = get_child(tree, 2).val
         @return_on_nonfinite_val(eval_options, val_r, cX)
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(cX[feature_l, j], val_r)::T
             cumulator[j] = x
         end  # COV_EXCL_LINE
@@ -678,7 +678,7 @@ function deg2_l0_r0_eval(
         cumulator = get_array(eval_options.buffer, cX, axes(cX, 2))
         feature_l = get_child(tree, 1).feature
         feature_r = get_child(tree, 2).feature
-        @inbounds @simd for j in axes(cX, 2)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(cX[feature_l, j], cX[feature_r, j])::T
             cumulator[j] = x
         end  # COV_EXCL_LINE
@@ -697,14 +697,14 @@ function deg2_l0_eval(
     if get_child(tree, 1).constant
         val = get_child(tree, 1).val
         @return_on_nonfinite_val(eval_options, val, cX)
-        @inbounds @simd for j in eachindex(cumulator)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(val, cumulator[j])::T
             cumulator[j] = x
         end  # COV_EXCL_LINE
         return ResultOk(cumulator, true)
     else
         feature = get_child(tree, 1).feature
-        @inbounds @simd for j in eachindex(cumulator)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(cX[feature, j], cumulator[j])::T
             cumulator[j] = x
         end  # COV_EXCL_LINE
@@ -723,14 +723,14 @@ function deg2_r0_eval(
     if get_child(tree, 2).constant
         val = get_child(tree, 2).val
         @return_on_nonfinite_val(eval_options, val, cX)
-        @inbounds @simd for j in eachindex(cumulator)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(cumulator[j], val)::T
             cumulator[j] = x
         end  # COV_EXCL_LINE
         return ResultOk(cumulator, true)
     else
         feature = get_child(tree, 2).feature
-        @inbounds @simd for j in eachindex(cumulator)
+        @finite @inbounds @simd ivdep for j in eachindex(cumulator)
             x = op(cumulator[j], cX[feature, j])::T
             cumulator[j] = x
         end  # COV_EXCL_LINE

src/EvaluateDerivative.jl

@@ -2,7 +2,7 @@ module EvaluateDerivativeModule
 
 import ..NodeModule: AbstractExpressionNode, constructorof, get_children
 import ..OperatorEnumModule: OperatorEnum
-import ..UtilsModule: fill_similar, ResultOk2
+import ..UtilsModule: fill_similar, ResultOk2, @finite
 import ..ValueInterfaceModule: is_valid_array
 import ..NodeUtilsModule: count_constant_nodes, index_constant_nodes, NodeIndex
 import ..EvaluateModule:
@@ -105,7 +105,7 @@ end
             dx_cumulator_i = dx_cumulators[i]
         end)
         diff_op = _zygote_gradient(op, Val(N))
-        @inbounds @simd for j in eachindex(x_cumulator_1)
+        @finite @inbounds @simd ivdep for j in eachindex(x_cumulator_1)
             x = Base.Cartesian.@ncall($N, op, i -> x_cumulator_i[j])
             Base.Cartesian.@ntuple($N, i -> grad_i) = Base.Cartesian.@ncall(
                 $N, diff_op, i -> x_cumulator_i[j]
@@ -346,13 +346,13 @@ end
             d_cumulator_i = d_cumulators[i]
         end)
         diff_op = _zygote_gradient(op, Val($N))
-        @inbounds @simd for j in eachindex(x_cumulator_1)
+        @finite @inbounds for j in eachindex(x_cumulator_1)
             x = Base.Cartesian.@ncall($N, op, i -> x_cumulator_i[j])
             Base.Cartesian.@ntuple($N, i -> grad_i) = Base.Cartesian.@ncall(
                 $N, diff_op, i -> x_cumulator_i[j]
             )
             x_cumulator_1[j] = x
-            for k in axes(d_cumulator_1, 1)
+            @simd ivdep for k in axes(d_cumulator_1, 1)
                 d_cumulator_1[k, j] = Base.Cartesian.@ncall(
                     $N,
                     +,

src/NodeUtils.jl

@@ -1,6 +1,6 @@
 module NodeUtilsModule
 
-using ..UtilsModule: Nullable
+using ..UtilsModule: Nullable, @finite
 
 import ..NodeModule:
     AbstractNode,
@@ -108,7 +108,7 @@ function get_scalar_constants(
     else
         vals = Vector{BT}(undef, count_scalar_constants(tree))
         i = firstindex(vals)
-        for ref in refs
+        @finite for ref in refs
             i = pack_scalar_constants!(vals, i, ref[].val::T)
         end
         return vals, refs
@@ -123,13 +123,13 @@ Set the constants in a tree, in depth-first order. The function
 """
 function set_scalar_constants!(tree::AbstractExpressionNode{T}, constants, refs) where {T}
     if T <: Number
-        @inbounds for i in eachindex(refs, constants)
+        @finite @inbounds for i in eachindex(refs, constants)
             refs[i][].val = constants[i]
         end
     else
         nums_i = 1
         refs_i = 1
-        while nums_i <= length(constants) && refs_i <= length(refs)
+        @finite while nums_i <= length(constants) && refs_i <= length(refs)
             ix, v = unpack_scalar_constants(constants, nums_i, refs[refs_i][].val::T)
             refs[refs_i][].val = v
             nums_i = ix

src/ParametricExpression.jl

@@ -8,7 +8,7 @@ using ..NodeModule: AbstractExpressionNode, Node, tree_mapreduce
 using ..ExpressionModule:
     AbstractExpression, Metadata, with_contents, with_metadata, unpack_metadata
 using ..ChainRulesModule: NodeTangent
-using ..UtilsModule: Nullable, set_nan!
+using ..UtilsModule: Nullable, set_nan!, @finite
 
 import ..NodeModule:
     constructorof,
@@ -241,15 +241,15 @@ function _get_constants_array(parameter_refs, ::Type{BT}) where {BT}
     size = sum(count_scalar_constants, parameter_refs)
     flat = Vector{BT}(undef, size)
     ix = 1
-    for p in parameter_refs
+    @finite for p in parameter_refs
         ix = pack_scalar_constants!(flat, ix, p)
     end
     return flat
 end
 
 function _set_constants_array!(parameter_refs, flat)
     ix, i = 1, 1
-    while ix <= length(flat) && i <= length(parameter_refs)
+    @finite while ix <= length(flat) && i <= length(parameter_refs)
         ix, parameter_refs[i] = unpack_scalar_constants(flat, ix, parameter_refs[i])
         i += 1
     end
@@ -378,7 +378,7 @@ function eval_tree_array(
     @assert length(classes) == size(X, 2)
     @assert maximum(classes) <= size(ex.metadata.parameters, 2)  # TODO: Remove when comfortable
     parameters = ex.metadata.parameters
-    indexed_parameters = [
+    indexed_parameters = @finite [
         parameters[i_parameter, classes[i_row]] for
         i_parameter in eachindex(axes(parameters, 1)), i_row in eachindex(classes)
     ]

src/Strings.jl

@@ -1,6 +1,6 @@
 module StringsModule
 
-using ..UtilsModule: deprecate_varmap
+using ..UtilsModule: deprecate_varmap, @finite
 using ..OperatorEnumModule: AbstractOperatorEnum
 using ..NodeModule: AbstractExpressionNode, tree_mapreduce, max_degree
 
@@ -122,7 +122,7 @@ function combine_op_with_inputs(op, args::Vararg{Any,D})::Vector{Char} where {D}
         # "op(l, r)"
         out = copy(op)
         push!(out, '(')
-        for i in 1:(D - 1)
+        @finite for i in 1:(D - 1)
             append!(out, strip_brackets(args[i]))
             push!(out, ',')
             push!(out, ' ')

src/Utils.jl

@@ -75,4 +75,13 @@ function set_nan!(out)
     return nothing
 end
 
+"""
+    @finite ex
+
+Wraps `ex` in a `Base.@assume_effects :terminates_locally` block.
+"""
+macro finite(ex)
+    return esc(:(Base.@assume_effects :terminates_locally $ex))
+end
+
 end