diff --git a/src/findall.jl b/src/findall.jl index d81b0334..2f305a2d 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -257,9 +257,9 @@ function findall(values::AbstractArray; items::AbstractArray=keys(values), findall_run(identity, values, items, s, scratch_buffers(s.plan, workspace, values, items)) end -function findall(pred, v::AbstractArray; items::AbstractArray=keys(v), +function findall(pred::F, v::AbstractArray; items::AbstractArray=keys(v), backend::Union{Nothing, Backend}=nothing, alg::Algorithm=Auto(), - workspace=nothing) + workspace=nothing) where {F} check_bool_result(pred, v) s = findall_setup(v, items, true, backend, alg) findall_run(pred, v, items, s, scratch_buffers(s.plan, workspace, v, items)) @@ -268,8 +268,8 @@ end plan(::typeof(findall), values::AbstractArray; items::AbstractArray=keys(values), backend=nothing, alg::Algorithm=Auto()) = findall_setup(values, items, false, backend, alg).plan -plan(::typeof(findall), pred, v::AbstractArray; items::AbstractArray=keys(v), backend=nothing, - alg::Algorithm=Auto()) = +plan(::typeof(findall), pred::F, v::AbstractArray; items::AbstractArray=keys(v), + backend=nothing, alg::Algorithm=Auto()) where {F} = findall_setup(v, items, true, backend, alg).plan # The plan of `findall` over `v`: its scratch is a mask of `pred(v[i])` for the predicate form @@ -291,7 +291,7 @@ function findall_setup(v, items, predicate::Bool, backend, alg) return (; plan=Plan(backend, a, merge(mask, counts))) end -function findall_run(pred, v, items, s, bufs) +function findall_run(pred::F, v, items, s, bufs) where {F} backend, a = s.plan.backend, s.plan.alg bools = if haskey(bufs, :mask) findall_mask!(bufs.mask, pred, v, backend, a) diff --git a/src/predicates.jl b/src/predicates.jl index cebd6c6b..39ccabf8 100644 --- a/src/predicates.jl +++ b/src/predicates.jl @@ -162,19 +162,20 @@ end complex_any(CuArray(rand(Float32, 100)), CuArray(rand(Float32, 100))) ``` """ -function any(pred, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, - alg::Algorithm=Auto(), workspace=nothing) +function any(pred::F, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), workspace=nothing) where {F} p = predicate_setup(pred, |, false, v, backend, alg).plan - any_run(pred, v, p.backend, p.alg, scratch_buffers(p, workspace, v)) + # (asserted, as `Auto()` picks a GPU's algorithm from its tuning, which inference does not see) + any_run(pred, v, p.backend, p.alg, scratch_buffers(p, workspace, v))::Bool end -plan(::typeof(any), pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) = +plan(::typeof(any), pred::F, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) where {F} = predicate_setup(pred, |, false, v, backend, alg).plan # The predicate must return a `Bool`, as in Base. Where inference shows it cannot, the call fails # on the host, before launching, with an `ArgumentError` instead of an exception from the kernel # (which GPU backends report without Base's `TypeError`). -function check_bool_result(pred, v) +function check_bool_result(pred::F, v) where {F} isempty(v) && return nothing # (Base never calls the predicate then) R = Base.promote_op(pred, eltype(v)) R === Union{} || Bool <: R || throw(ArgumentError( @@ -184,7 +185,7 @@ end # The plan of `any`/`all` (whose `ViaReduce` reduces with `op` from `init`): the flag of # `ConcurrentWrite`, or the scratch of the reduction -function predicate_setup(pred, op, init, v, backend, alg) +function predicate_setup(pred::F, op::OP, init, v, backend, alg) where {F, OP} check_bool_result(pred, v) backend = resolve_backend(backend, v) a = resolve_predicate(alg, backend, eltype(v)) @@ -199,7 +200,7 @@ function predicate_setup(pred, op, init, v, backend, alg) return (; plan=Plan(backend, a, sizes)) end -function any_run(pred, v, backend, alg::ConcurrentWrite, bufs) +function any_run(pred::F, v, backend, alg::ConcurrentWrite, bufs) where {F} # `ndrange` must not be zero isempty(v) && return false out = bufs.flag @@ -208,10 +209,10 @@ function any_run(pred, v, backend, alg::ConcurrentWrite, bufs) return @allowscalar(out[1]) != 0 end -any_run(pred, v, backend, alg::ViaReduce, bufs) = +any_run(pred::F, v, backend, alg::ViaReduce, bufs) where {F} = mapreduce_nested(BoolResult(pred), |, v, bufs.reduce; backend, init=false, alg=alg.reduce) -function any_run(pred, v, backend, alg::CPUThreads.Partitioned, bufs) +function any_run(pred::F, v, backend, alg::CPUThreads.Partitioned, bufs) where {F} found = Threads.Atomic{Bool}(false) first_index = firstindex(v) task_partition(length(v), alg.max_tasks, alg.min_elems) do irange @@ -252,17 +253,18 @@ end complex_all(MtlArray(rand(Float32, 100)), MtlArray(rand(Float32, 100))) ``` """ -function all(pred, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, - alg::Algorithm=Auto(), workspace=nothing) +function all(pred::F, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), workspace=nothing) where {F} p = predicate_setup(pred, &, true, v, backend, alg).plan - all_run(pred, v, p.backend, p.alg, scratch_buffers(p, workspace, v)) + all_run(pred, v, p.backend, p.alg, scratch_buffers(p, workspace, v))::Bool # (as in `any`) end -plan(::typeof(all), pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) = +plan(::typeof(all), pred::F, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) where {F} = predicate_setup(pred, &, true, v, backend, alg).plan -all_run(pred, v, backend, alg::ConcurrentWrite, bufs) = !any_run(!pred, v, backend, alg, bufs) -all_run(pred, v, backend, alg::ViaReduce, bufs) = +all_run(pred::F, v, backend, alg::ConcurrentWrite, bufs) where {F} = + !any_run(!pred, v, backend, alg, bufs) +all_run(pred::F, v, backend, alg::ViaReduce, bufs) where {F} = mapreduce_nested(BoolResult(pred), &, v, bufs.reduce; backend, init=true, alg=alg.reduce) -all_run(pred, v, backend, alg::CPUThreads.Partitioned, bufs) = +all_run(pred::F, v, backend, alg::CPUThreads.Partitioned, bufs) where {F} = !any_run(!pred, v, backend, alg, bufs) diff --git a/src/reduce/shorthands.jl b/src/reduce/shorthands.jl index cdedae24..5f97c709 100644 --- a/src/reduce/shorthands.jl +++ b/src/reduce/shorthands.jl @@ -109,7 +109,7 @@ c = AK.count(CuArray(rand(Bool, 10, 100_000)); init=Int32(0), dims=2) ``` """ count(src::AbstractArray; kwargs...) = count(identity, src; kwargs...) -count(f, src::AbstractArray; init=0, kwargs...) = +count(f::F, src::AbstractArray; init=0, kwargs...) where {F} = mapreduce(BoolResult(f), Base.add_sum, src; init, kwargs...) plan(::typeof(sum), src::AbstractArray; kwargs...) = plan(reduce, Base.add_sum, src; kwargs...) @@ -117,5 +117,5 @@ plan(::typeof(prod), src::AbstractArray; kwargs...) = plan(reduce, Base.mul_prod plan(::typeof(maximum), src::AbstractArray; kwargs...) = plan(reduce, max, src; kwargs...) plan(::typeof(minimum), src::AbstractArray; kwargs...) = plan(reduce, min, src; kwargs...) plan(::typeof(count), src::AbstractArray; kwargs...) = plan(count, identity, src; kwargs...) -plan(::typeof(count), f, src::AbstractArray; init=0, kwargs...) = +plan(::typeof(count), f::F, src::AbstractArray; init=0, kwargs...) where {F} = plan(mapreduce, BoolResult(f), Base.add_sum, src; init, kwargs...) diff --git a/test/generic/predicates.jl b/test/generic/predicates.jl index 14561ec1..a97d9d88 100644 --- a/test/generic/predicates.jl +++ b/test/generic/predicates.jl @@ -22,6 +22,10 @@ HOST_KERNELS || pushfirst!(PRED_ALGS, AK.Auto()) @test AK.all(x->x>0, v; alg) === true @test AK.all(x->x<100, v; alg) === false + # The result infers, also where `Auto()` picks the algorithm from a GPU's tuning + @test @inferred(AK.any(x->x<0, v; alg)) === false + @test @inferred(AK.all(x->x>0, v; alg)) === true + for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) @@ -76,3 +80,30 @@ HOST_KERNELS || pushfirst!(PRED_ALGS, AK.Auto()) # Test that undefined kwargs are not accepted @test_throws MethodError AK.any(x->x<0, v; alg, bad=:kwarg) end + +# Julia does not specialize a method on a function argument that it only passes on, so the +# methods that take the predicate name its type. Otherwise they run code compiled for any +# `Function`, in which `check_bool_result` infers the predicate's result type on every call. +# (Called by dynamic dispatch, as an inferred caller may inline past the entry point.) +@testset "predicates: no inference at run time" begin + # The compilation time of the fastest of a few warm calls (other tasks may compile once) + function compile_time(f, args...; kwargs...) + invokelatest(f, args...; kwargs...) + minimum(1:3) do _ + Base.cumulative_compile_timing(true) + try + t = Base.cumulative_compile_time_ns()[1] + invokelatest(f, args...; kwargs...) + Base.cumulative_compile_time_ns()[1] - t + finally + Base.cumulative_compile_timing(false) + end + end + end + v = array_from_host(rand(Float32, 1000)) + for alg in PRED_ALGS + @test compile_time(AK.any, isnan, v; alg) == 0 + @test compile_time(AK.all, isfinite, v; alg) == 0 + end + @test compile_time(AK.findall, isnan, v; alg=FINDALL_ALG) == 0 +end