diff --git a/src/Evaluate.jl b/src/Evaluate.jl index 37c0f19a..f8e550b9 100644 --- a/src/Evaluate.jl +++ b/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 diff --git a/src/EvaluateDerivative.jl b/src/EvaluateDerivative.jl index 6841557a..20e196f0 100644 --- a/src/EvaluateDerivative.jl +++ b/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, +, diff --git a/src/NodeUtils.jl b/src/NodeUtils.jl index 5b75f7b1..2a3f3551 100644 --- a/src/NodeUtils.jl +++ b/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 diff --git a/src/ParametricExpression.jl b/src/ParametricExpression.jl index e17a3cf7..5b027273 100644 --- a/src/ParametricExpression.jl +++ b/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,7 +241,7 @@ 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 @@ -249,7 +249,7 @@ 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) ] diff --git a/src/Strings.jl b/src/Strings.jl index 55ff9e5b..655203e5 100644 --- a/src/Strings.jl +++ b/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, ' ') diff --git a/src/Utils.jl b/src/Utils.jl index 2522a9b1..f8bfcaaa 100644 --- a/src/Utils.jl +++ b/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