RFC: @jet_ignore to let JET.jl ignore specific functions

General Usage
1

Unleashing JET.jl on my codebase tends to flood the prompt with errors due to the use of IO, DataFrames, etc…
A lot of these runtime-dispatches are okay, but some are not.

As an example, assume that I’m okay with the println statement below, I can tag g with a @jet_ignore:

f(x) = 1 + g(x) 
@jet_ignore g(x) = (println(x); x)

Regular JET.jl will fail, but a modified @report_opt could ignore any runtime-dispatches in g, i.e.:

JET.@report_opt f(2)  # reports 1 error
JET.@test_opt f(2)  # test failed

@xjet f(2)  # no errors detected
@test_xjet f(2)  # test passed

A simple implementation of @jet_ignore/@xjet:

using MacroTools
using JET

const JET_IGNORE = gensym("_jet_ignore_")

macro jet_ignore(ex)
    f = splitdef(ex)
    jet_ignore_f = splitdef(ex)
    jet_ignore_nm = Symbol("$JET_IGNORE$(f[:name])")
    f[:body] = :($jet_ignore_nm($(f[:args]...); $(f[:kwargs]...)))
    jet_ignore_f[:name] = jet_ignore_nm
    return quote
        $(esc(combinedef(f)))
        $(esc(combinedef(jet_ignore_f)))
    end
end

is_jet_ignore(frame) = startswith(string(frame.linfo.def.name), string(JET_IGNORE))

# ~JET.report_opt
function xjet(args...; kwargs...)
    result = JET.report_opt(args...; kwargs...)
    filter!(JET.get_reports(result)) do report
        return !any(is_jet_ignore, report.vst)
    end
    result
end

# ~JET.@report_opt
macro xjet(ex0...)
    return JET.InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :xjet, ex0)
end

# ~JET.@test_opt
macro test_xjet(ex0...)
    return JET.call_test_ex(:xjet, Symbol("@test_opt"), ex0, __module__, __source__)
end

The ability to annotate specific functions was inspired by DispatchDoctor.jl.
Like DispatchDoctor.jl, @jet_ignore works by inserting an extra function layer. In this case a function with a specific name that can be picked up and filtered out.

Does this have broader appeal?
Could/should it be part of JET.jl?
Any other comments?

ping @aviatesk

1 Like