diff --git a/Project.toml b/Project.toml index 05a7ce8..0453009 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FunctionProperties" uuid = "f62d2435-5019-4c03-9749-2d4c77af0cbc" -version = "0.1.4" +version = "0.1.5" authors = ["SciML"] [deps] diff --git a/src/FunctionProperties.jl b/src/FunctionProperties.jl index 761e7fa..bd17be9 100644 --- a/src/FunctionProperties.jl +++ b/src/FunctionProperties.jl @@ -2,11 +2,15 @@ module FunctionProperties using Core: GotoIfNot +# Backstop against pathological recursion depth; real call trees that matter here are shallow. +const RECURSION_LIMIT = 256 + """ is_leaf(f, args...) -> Bool Override this to exempt a function from `hasbranching` analysis. -Return `true` to treat `f` as branch-free regardless of its implementation. +Return `true` to treat `f` as branch-free regardless of its implementation, which also +stops [`hasbranching`](@ref) from recursing into its callees. ## Example @@ -31,12 +35,19 @@ like ReverseDiff.jl. ## Outputs -Boolean for whether the function's immediate IR contains a conditional branch (`GotoIfNot`). +Boolean for whether the function contains a value-dependent conditional branch +(`GotoIfNot`). The type-inferred IR (`code_typed(...; optimize = false)`) of `f` is scanned, +and the scan recurses through statically resolved calls into other user-defined methods so +that branches living behind a non-inlined call boundary are still detected. Calls into +`Base`, `Core`, and the standard libraries are treated as leaves: their internal branches +are structural/compile-time rather than value-dependent user logic, and recursing into them +(e.g. matrix multiply, broadcasting, `getindex` bounds checks) produces false positives. ## Customizing and Removing Functions from the Checks Some functions may produce false positives because their internal branches are compile-time -constants. Override `FunctionProperties.is_leaf` to opt them out: +constants. Override [`is_leaf`](@ref) to opt them out (this also prevents recursion into +them): ```julia FunctionProperties.is_leaf(::typeof(my_fn)) = true @@ -44,17 +55,121 @@ FunctionProperties.is_leaf(::typeof(my_fn)) = true """ function hasbranching(f, x...) is_leaf(f, x...) && return false - argtypes = Tuple{Core.Typeof.(x)...} + sig = Tuple{Core.Typeof(f), Core.Typeof.(x)...} + return _hasbranching(sig, Set{Any}(), 0) +end + +function _hasbranching(@nospecialize(sig), seen, depth) + depth > RECURSION_LIMIT && return false + sig in seen && return false + push!(seen, sig) + results = try - code_typed(f, argtypes; optimize = false) + Base.code_typed_by_type(sig; optimize = false) + catch + return false + end + + for pair in results + ci = first(pair) + # Generated functions that were not expanded come back as `Method`, not `CodeInfo`; + # there is no body to scan, so treat them as leaves. + ci isa Core.CodeInfo || continue + any(stmt -> isa(stmt, GotoIfNot), ci.code) && return true + for stmt in ci.code + _recurse_call(stmt, ci, seen, depth) && return true + end + end + return false +end + +# Inspect a single IR statement: if it is a statically resolvable call into a non-library +# method, recurse into that method's IR. Returns `true` if a branch is found downstream. +function _recurse_call(@nospecialize(stmt), ci, seen, depth) + call = Meta.isexpr(stmt, :(=)) ? stmt.args[2] : stmt + + if Meta.isexpr(call, :invoke) + mi = call.args[1] + callsig = mi isa Core.MethodInstance ? mi.specTypes : + ( + isdefined(mi, :def) && getfield(mi, :def) isa Core.MethodInstance ? + getfield(mi, :def).specTypes : nothing + ) + callsig === nothing && return false + return _recurse_sig(callsig, nothing, seen, depth) + end + + Meta.isexpr(call, :call) || return false + ftype, fval = _resolve_callee(call.args[1], ci) + ftype === nothing && return false + argtypes = Any[_value_type(a, ci) for a in @view call.args[2:end]] + return _recurse_sig(Tuple{ftype, argtypes...}, fval, seen, depth) +end + +function _recurse_sig(@nospecialize(callsig), @nospecialize(fval), seen, depth) + # Honor user `is_leaf` overrides when the concrete function value is recoverable. + fval !== nothing && is_leaf(fval) && return false + m = try + Base.which(callsig) catch return false end - isempty(results) && return false - ci = first(results)[1] - return any(isa(s, GotoIfNot) for s in ci.code) + _is_library_method(m) && return false + return _hasbranching(callsig, seen, depth + 1) +end + +# Library code (`Base`, `Core`, stdlibs) is treated as a leaf: its branches are structural or +# compile-time, not the value-dependent user logic `hasbranching` is meant to surface. +function _is_library_method(m::Method) + root = Base.moduleroot(m.module) + (root === Base || root === Core) && return true + pkgdir = Base.pkgdir(root) + pkgdir === nothing && return false + return startswith(pkgdir, Sys.STDLIB) +end + +function _resolve_callee(@nospecialize(fexpr), ci) + if fexpr isa GlobalRef + if isdefined(fexpr.mod, fexpr.name) && isconst(fexpr.mod, fexpr.name) + v = getglobal(fexpr.mod, fexpr.name) + return (Core.Typeof(v), v) + end + return (nothing, nothing) + elseif fexpr isa QuoteNode + return (Core.Typeof(fexpr.value), fexpr.value) + elseif fexpr isa Core.SSAValue + t = ci.ssavaluetypes[fexpr.id] + t isa Core.Const && return (Core.Typeof(t.val), t.val) + return (_widen(t), nothing) + else + return (_value_type(fexpr, ci), nothing) + end +end + +function _value_type(@nospecialize(a), ci) + if a isa Core.SSAValue + return _widen(ci.ssavaluetypes[a.id]) + elseif a isa Core.Argument + st = ci.slottypes + return st === nothing ? Any : _widen(st[a.n]) + elseif a isa Core.SlotNumber + st = ci.slottypes + return st === nothing ? Any : _widen(st[a.id]) + elseif a isa GlobalRef + return (isdefined(a.mod, a.name) && isconst(a.mod, a.name)) ? + Core.Typeof(getglobal(a.mod, a.name)) : Any + elseif a isa QuoteNode + return Core.Typeof(a.value) + else + return Core.Typeof(a) + end end +_widen(@nospecialize t) = + t isa Core.Const ? Core.Typeof(t.val) : + t isa Core.PartialStruct ? t.typ : + isa(t, Type) ? t : Any + export hasbranching, is_leaf end diff --git a/test/core_tests.jl b/test/core_tests.jl index e46c81b..e166bd7 100644 --- a/test/core_tests.jl +++ b/test/core_tests.jl @@ -17,6 +17,23 @@ f_branch() = true ? 1 : 0 FunctionProperties.is_leaf(::typeof(f_branch)) = true @test !FunctionProperties.hasbranching(f_branch) +# Branches behind a non-inlined call boundary must still be detected: the value-dependent +# `if` lives in `branchy_helper`, not in the immediate IR of the entry function. +@noinline branchy_helper(x) = x < 0 ? -x : x +nested_branch_rhs(u, p, t) = branchy_helper(u) + p +@test FunctionProperties.hasbranching(nested_branch_rhs, 1.0, 2.0, 0.0) + +@noinline branchfree_helper(x) = x * x + one(x) +nested_branchfree_rhs(u, p, t) = branchfree_helper(u) + p +@test !FunctionProperties.hasbranching(nested_branchfree_rhs, 1.0, 2.0, 0.0) + +# An `is_leaf` override stops recursion into the marked callee. +@noinline opted_out_helper(x) = x < 0 ? -x : x +opted_out_rhs(u, p, t) = opted_out_helper(u) + p +@test FunctionProperties.hasbranching(opted_out_rhs, 1.0, 2.0, 0.0) +FunctionProperties.is_leaf(::typeof(opted_out_helper)) = true +@test !FunctionProperties.hasbranching(opted_out_rhs, 1.0, 2.0, 0.0) + # Test simple mutating functions function f(dx, x) return @inbounds dx[1] = x[1]