How to overdub static parameters in overdubbing

Internals & Design
1

Hello,

This question has stemmed from the work described in Does overdubbing in generated function inserts inlined code , where I would like to understand the code rewriting of CodeInstances. I decided to try to do everything manually without IRTools / Cassete to know, on which internals they are built on.

Currently, I am stucked on the following problem. Consider following function

function test(x::T) where T<:Union{Float64, Float32}
    x < T(pi)
end

which test, if some variable is smaller then pi (I have encountered this when overdubbing Main.sin). When I retrieve the code using

retrieve_code_info 
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)
    if isempty(_methods) 
        @info("method $(sigtypes) does not exist")
        return(nothing)
    end
    type_signature, raw_static_params, method = _methods[1]
    method_instance = Core.Compiler.specialize_method(method, type_signature, raw_static_params, false)
    code_info = Core.Compiler.retrieve_code_info(method_instance)
end

the code looks like this

CodeInfo(
    @ REPL[33]:2 within `test'
1 ─ %1 = ($(Expr(:static_parameter, 1)))(Main.pi)
│   %2 = x < %1
└──      return %2
)

from which I conclude that T(pi) is translated to ($(Expr(:static_parameter, 1)))(Main.pi) and the Expr(:static_parameter, 1) is an accessor to the first Type parameter T.
Now comes the problem.
When I am overdubbing, I do substitute function signature test(x) with a generated function overdub(f args...) which takes code info for test(x) and manipulates it (and it does the same for any other function) does not have any type parameters. Hence, I do not have access to T in a generated function overdub(typeof(test), args...). My question is, what is the general strategy to translate such line, for example in Zygote to make it compatible with a signature of overdub(f, args...)? The only strategy I can imagine at the moment is to remove the $(Expr(:static_parameter, 1))from the type definition, yet I am afraid that doing so will make the code type-unstable. I have search IRTools, Zygote and Cassette forstatic_parameters` but without luck.
Thanks for help and clarifications.

Tomas

2

Zygote and IRTools call Meta.partially_inline! to inline the static parameters from the method match directly into the IR.

I’ll highly recommend just using Cassette directly though, since all these tricky details can be quite annoying to handle correctly.

3

I see. Thanks.
I know that I should just use directly one of these libraries but as I said, I would like to get at least a partial understanding what is going on. It is for educational purposes for a course on Julia on Czech Technical University (https://github.com/JuliaTeachingCTU/Scientific-Programming-in-Julia).
I am now playing with IRTools and of course, it makes life much simpler, as expected…

I appreciate your help.
Best wishes,
Tomas

1 Like