Does overdubbing in generated function inserts inlined code

Tomas_Pevny · November 21, 2021, 8:55pm

Dear All,

I am trying to understand the role of generated functions in Zygote / Cassette, as want to properly understand “overdubbing design pattern”. I decided to test the idea by writing a profiler, which will log the start and end of each function called. This is kind of an alternative to the sampling profiler available in Julia. Note that the implementation below is not a package (and likely never going to be), it is something I decided to implement to educate myself. The full code is

"Full code

@generated function overdub(ctx::Context, f::F, args...) where {F}
    ci = retrieve_code_info((F, args...))
    slot_vars = Dict(enumerate(ci.slotnames))
    # ssa_vars = Dict(i => gensym(:left) for i in 1:length(ci.code))
    ssa_vars = Dict(i => Symbol(:L, i) for i in 1:length(ci.code))
    used = assigned_vars(ci.code) |> distinct
    exprs = []
    for i in 1:length(args)
        push!(exprs, Expr(:(=), ci.slotnames[i+1], :(args[$(i)])))
    end
    for (i, ex) in enumerate(ci.code)
        ex = rename_args(ex, slot_vars, ssa_vars)
        if ex isa Core.ReturnNode 
            push!(exprs, Expr(:return, ex.val))
            continue
        end
        if timable(ex)
            fname = exportname(ex)
            fname = :(Symbol($(fname)))
            push!(exprs, Expr(:call, :push!, :to, :(:start), fname))
            ex = overdubbable(ex) ? Expr(:call, :overdub, :ctx, ex.args...) : ex
            ex = i ∈ used ? Expr(:(=) , ssa_vars[i], ex) : ex
            push!(exprs, ex)
            push!(exprs, Expr(:call, :push!, :to, :(:stop), fname))
        else
            ex = i ∈ used ? Expr(:(=) , ssa_vars[i], ex) : ex
            push!(exprs, ex)
        end
    end
    r = Expr(:block, exprs...)
    @show r 
    # println("  ")
    r
end

For debugging purposes I have not hygienized inserted symbols and also I print the generated code, which is against the generated functions not to have side effects, but it is for debugging and final code will not have it. The above code does not recursively descend into the function calls, and it produces the result I expect, as

julia> function foo(x, y)
          z = x * y
          z + sin(y)
       end
foo (generic function with 1 method)

julia> reset!(to)
0

julia> overdub(Context(), foo, 1.0, 1.0)
r = quote
    x = args[1]
    y = args[2]
    z = x * y
    L2 = z
    push!(to, :start, Symbol(Main.sin))
    L3 = Main.sin(y)
    push!(to, :stop, Symbol(Main.sin))
    push!(to, :start, Symbol(Main.:+))
    L4 = L2 + L3
    push!(to, :stop, Symbol(Main.:+))
    return L4
end
1.8414709848078965

julia> to
0.0  start  sin
1.9073486328125e-6  stop  sin
3.814697265625e-6  start  +
3.814697265625e-6  stop  +

If I know enable the recursion by redefining overdubbable as

function overdubbable(ex::Expr) 
    ex.head != :call && return(false)
    length(ex.args) < 2 && return(false)
    (ex.args[1] isa Core.IntrinsicFunction) && return(false)
    return(true)
end

the overdubbing goes nuts and I am apparently overdubbing a function that I have already overdubbed.

A recursive overdubbing that went nuts

julia> overdub(Context(), foo, 1.0, 1.0)
r = quote
x = args[1]
y = args[2]
z = x * y
L2 = z
push!(to, :start, Symbol(Main.sin))
L3 = overdub(ctx, Main.sin, y)
push!(to, :stop, Symbol(Main.sin))
push!(to, :start, Symbol(Main.:+))
L4 = overdub(ctx, Main.:+, L2, L3)
push!(to, :stop, Symbol(Main.:+))
return L4
end

r = quote
x = args[1]
Core.NewvarNode(:(_3))
Core.NewvarNode(:(_4))
Core.NewvarNode(:(_5))
absx = Base.Math.abs(x)
L5 = absx
push!(to, :start, Symbol((Expr(:static_parameter, 1)))) L6 = overdub(ctx, (Expr(:static_parameter, 1)), Base.Math.pi)
push!(to, :stop, Symbol((Expr(:static_parameter, 1)))) push!(to, :start, Symbol(Base.Math.:/)) L7 = overdub(ctx, Base.Math.:/, L6, 4) push!(to, :stop, Symbol(Base.Math.:/)) push!(to, :start, Symbol(Base.Math.:<)) overdub(ctx, Base.Math.:<, L5, L7) push!(to, :stop, Symbol(Base.Math.:<)) goto %18 if not %8 L10 = absx push!(to, :start, Symbol(Base.Math.eps)) L11 = overdub(ctx, Base.Math.eps, (Expr(:static_parameter, 1)))
push!(to, :stop, Symbol(Base.Math.eps))
push!(to, :start, Symbol(Base.Math.sqrt))
L12 = overdub(ctx, Base.Math.sqrt, L11)
push!(to, :stop, Symbol(Base.Math.sqrt))
push!(to, :start, Symbol(Base.Math.:<))
overdub(ctx, Base.Math.:<, L10, L12)
push!(to, :stop, Symbol(Base.Math.:<))
goto %16 if not %13
return x
push!(to, :start, Symbol(Base.Math.sin_kernel))
L16 = overdub(ctx, Base.Math.sin_kernel, x)
push!(to, :stop, Symbol(Base.Math.sin_kernel))
return L16
push!(to, :start, Symbol(Base.Math.isnan))
overdub(ctx, Base.Math.isnan, x)
push!(to, :stop, Symbol(Base.Math.isnan))
goto %22 if not %18
push!(to, :start, Symbol((Expr(:static_parameter, 1)))) L20 = overdub(ctx, (Expr(:static_parameter, 1)), Base.Math.NaN)
push!(to, :stop, Symbol($(Expr(:static_parameter, 1))))
return L20
push!(to, :start, Symbol(Base.Math.isinf))
overdub(ctx, Base.Math.isinf, x)
push!(to, :stop, Symbol(Base.Math.isinf))
goto %25 if not %22
push!(to, :start, Symbol(Base.Math.sin_domain_error))
overdub(ctx, Base.Math.sin_domain_error, x)
push!(to, :stop, Symbol(Base.Math.sin_domain_error))
push!(to, :start, Symbol(Base.Math.rem_pio2_kernel))
L25 = overdub(ctx, Base.Math.rem_pio2_kernel, x)
push!(to, :stop, Symbol(Base.Math.rem_pio2_kernel))
push!(to, :start, Symbol(Base.indexed_iterate))
L26 = overdub(ctx, Base.indexed_iterate, L25, 1)
push!(to, :stop, Symbol(Base.indexed_iterate))
n = Core.getfield(L26, 1)
var"" = Core.getfield(L26, 2)
push!(to, :start, Symbol(Base.indexed_iterate))
L29 = overdub(ctx, Base.indexed_iterate, L25, 2, var"")
push!(to, :stop, Symbol(Base.indexed_iterate))
y = Core.getfield(L29, 1)
n = n & 3
push!(to, :start, Symbol(Base.Math.:(==)))
overdub(ctx, Base.Math.:(==), n, 0)
push!(to, :stop, Symbol(Base.Math.:(==)))
goto %36 if not %32
push!(to, :start, Symbol(Base.Math.sin_kernel))
L34 = overdub(ctx, Base.Math.sin_kernel, y)
push!(to, :stop, Symbol(Base.Math.sin_kernel))
return L34
push!(to, :start, Symbol(Base.Math.:(==)))
overdub(ctx, Base.Math.:(==), n, 1)
push!(to, :stop, Symbol(Base.Math.:(==)))
goto %40 if not %36
push!(to, :start, Symbol(Base.Math.cos_kernel))
L38 = overdub(ctx, Base.Math.cos_kernel, y)
push!(to, :stop, Symbol(Base.Math.cos_kernel))
return L38
push!(to, :start, Symbol(Base.Math.:(==)))
overdub(ctx, Base.Math.:(==), n, 2)
push!(to, :stop, Symbol(Base.Math.:(==)))
goto %45 if not %40
push!(to, :start, Symbol(Base.Math.sin_kernel))
L42 = overdub(ctx, Base.Math.sin_kernel, y)
push!(to, :stop, Symbol(Base.Math.sin_kernel))
push!(to, :start, Symbol(Base.Math.:-))
L43 = overdub(ctx, Base.Math.:-, L42)
push!(to, :stop, Symbol(Base.Math.:-))
return L43
push!(to, :start, Symbol(Base.Math.cos_kernel))
L45 = overdub(ctx, Base.Math.cos_kernel, y)
push!(to, :stop, Symbol(Base.Math.cos_kernel))
push!(to, :start, Symbol(Base.Math.:-))
L46 = overdub(ctx, Base.Math.:-, L45)
push!(to, :stop, Symbol(Base.Math.:-))
return L46
end

Can someone help me by explaining what I doing wrong or which part went nuts? I would like to be able to recursively overdubb functions that are called within overdubbed functions (for example here the function sin. Also, I am also totally surprised where the nodes Core.NewvarNode(:(_3)) came from.

Note that I know about IRTools and Cassette, but as I have said, I would like to understand the mechanics of those packages.

Thank you very much in advance.

Tomas Pevny

goerch · November 21, 2021, 9:06pm

Regarding your MWE: where does the definition of Context comes from?

On Julia 1.7.0-rc3 after

using IRTools
using Cassette
using Zygote

I still see

ERROR: LoadError: UndefVarError: Context not defined

simeonschaub · November 21, 2021, 9:15pm

You are making your life unnecessarily hard by trying to convert the linearized IR back into Julia’s surface AST. While that’s theoretically possible, there are a lot of special forms only valid inside IR you are not handling here. Instead, you can just return CodeInfo objects from generated functions, which is what Cassette and IRTools do.

Tomas_Pevny · November 22, 2021, 4:44am

Thanks Simon,

and the generated function can return CodeInfo object?

Tomas_Pevny · November 22, 2021, 10:41am

@simeonschaub Are there any docs, how I can create CodeInfor objects manually?

Tomas_Pevny · November 22, 2021, 10:48am

Sorry,

I have defined that by myself as a dummy struct

struct Context{T<:Union{Nothing, Vector{Symbol}}}
    functions::T
end
Context() = Context(nothing)

ctx = Context()

Tomas_Pevny · November 22, 2021, 8:11pm

As @simeonschaub has recommended, I have rewritten my code to return CodeInfo object, thought the documentation is scarce. My current solution looks like

# Generated functions
using Dictionaries
function retrieve_code_info(sigtypes, world = Base.get_world_counter())
  S = Tuple{map(s -> Core.Compiler.has_free_typevars(s) ? typeof(s.parameters[1]) : s, sigtypes)...}
  _methods = Base._methods_by_ftype(S, -1, world)
  isempty(_methods) && @error("method $(sigtypes) does not exist, may-be run it once")
  type_signature, raw_static_params, method = _methods[1] # method is the same as we would get by invoking methods(+, (Int, Int)).ms[1]  

  # this provides us with the CodeInfo
  method_instance = Core.Compiler.specialize_method(method, type_signature, raw_static_params, false)
  code_info = Core.Compiler.retrieve_code_info(method_instance)
end

struct Calls
    stamps::Vector{Float64} # contains the time stamps
    event::Vector{Symbol}  # name of the function that is being recorded
    startstop::Vector{Symbol} # if the time stamp corresponds to start or to stop
    i::Ref{Int}
end

function Calls(n::Int)
    Calls(Vector{Float64}(undef, n+1), Vector{Symbol}(undef, n+1), Vector{Symbol}(undef, n+1), Ref{Int}(0))
end

function Base.show(io::IO, calls::Calls)
    for i in 1:calls.i[]
        println(io, calls.stamps[i] - calls.stamps[1],"  ", calls.startstop[i],"  ",calls.event[i])
    end
end

function Base.push!(calls::Calls, s::Symbol, ev::Symbol)
    n = calls.i[] = calls.i[] + 1
    n > length(calls.stamps) && return 
    calls.event[n] = ev
    calls.startstop[n] = s
    calls.stamps[n] = time()
end

reset!(calls::Calls) = calls.i[] = 0


function overdubbable(ex::Expr) 
    ex.head != :call && return(false)
    length(ex.args) < 2 && return(false)
    (ex.args[1] isa Core.IntrinsicFunction) && return(false)
    return(true)
end

overdubbable(ex) = false
timable(ex::Expr) = ex.head == :call
timable(ex) = false

rename_args(ex, ssamap) = ex
rename_args(ex::Expr, ssamap) = Expr(ex.head, rename_args(ex.args, ssamap)...)
rename_args(args::AbstractArray, ssamap) = map(a -> rename_args(a, ssamap), args)
rename_args(r::Core.ReturnNode, ssamap) = Core.ReturnNode(rename_args(r.val, ssamap))
rename_args(a::Core.SSAValue, ssamap) = Core.SSAValue(ssamap[a.id])

exportname(ex::GlobalRef) = ex.name
exportname(ex::Expr) = ex.args[1]

using Base: invokelatest
dummy() = return

overdub(f::Core.IntrinsicFunction, args...) = f(args...)


@generated function overdub(f::F, args...) where {F}
    F = typeof(foo)
    args = (Float64, Float64)
    # ci = code_lowered((F, args))[1]
    ci = retrieve_code_info((F, args...))

    # this is to initialize a new CodeInfo and fill it with values from the 
    # overdubbed function
    new_ci = code_lowered(dummy, Tuple{})[1]
    empty!(new_ci.code)
    empty!(new_ci.slotnames)
    foreach(s -> push!(new_ci.slotnames, s), ci.slotnames)
    new_ci.slotnames = vcat([Symbol("#self#"), :f, :args], ci.slotnames[length(args)+2:end])
    empty!(new_ci.linetable)
    foreach(s -> push!(new_ci.linetable, s), ci.linetable)
    empty!(new_ci.codelocs)

    ssamap = Dict{Int, Int}()
    slotvar = Dict{Int, Any}()
    for i in 1:length(args)
        push!(new_ci.code, :(Base.getindex(args, $(i))))
        slotvar[i+1] = Core.SSAValue(i)
        push!(new_ci.codelocs, ci.codelocs[1])
    end
    slotvar[1] = Core.SlotNumber(1)
    foreach(i -> slotvar[i[2]] = Core.SlotNumber(i[1]+3), enumerate(length(args)+2:length(ci.slotnames)))

    j = length(args)
    for (i, ex) in enumerate(ci.code)
        ex = rename_args(ex, slotvar, ssamap)
        if timable(ex)
            fname = exportname(ex)
            fname = :(Symbol($(fname)))
            push!(new_ci.code, Expr(:call, :push!, :to, :(:start), fname))
            j += 1
            push!(new_ci.codelocs, ci.codelocs[i])
            ex = overdubbable(ex) ? Expr(:call, :overdub, ex.args...) : ex
            # ex = i ∈ used ? Expr(:(=) , ssa_vars[i], ex) : ex
            push!(new_ci.code, ex)
            push!(new_ci.codelocs, ci.codelocs[i])
            j += 1
            ssamap[i] = j
            push!(new_ci.code, Expr(:call, :push!, :to, :(:stop), fname))
            push!(new_ci.codelocs, ci.codelocs[i])
            j += 1
        else
            # ex = i ∈ used ? Expr(:(=) , ssa_vars[i], ex) : ex
            push!(new_ci.code, ex)
            push!(new_ci.codelocs, ci.codelocs[i])
            j += 1
            ssamap[i] = j
        end
    end
    new_ci

    # Core.Compiler.replace_code_newstyle!(ci, ir, length(ir.argtypes)-1)
    new_ci.inferred = false
    new_ci.ssavaluetypes = length(new_ci.code)
    new_ci
end

But when I call it as follows

function foo(x, y)
   z = x * y
   z + sin(y)
end

global const to = Calls(100)
reset!(to)
overdub(foo, 1.0, 1.0)

, the overdub does not executed foo but the Core.CodeInfo. This quite puzzling, as based on Simon’s post I would expect it is done automatically.

Thanks for any help.

Best wishes,
Tomas

simeonschaub · November 22, 2021, 8:33pm

Try adding an explicit return new_ci at the end. See https://github.com/JuliaLang/julia/issues/25678.

Tomas_Pevny · November 23, 2021, 4:26am

Thanks Simon,

it seems to be that. Julia has instantly crashed, which seems like I am having some bug in my codegen. I expect something like this.