Cassette, IRTools, Diffractor and others use generated functions to transform a method’s CodeInfo / IRCode. Although I used these tools dozen times, I never managed to understand or reproduce it. Is there a minimal self-containing example of such a transformation?
For the context, I’m currently fighting extremely slow tracing of deeply nested function calls using Umlaut. Umlaut uses CompilerPluginTools to extract IRCode and execute it instruction by instruction, but this way most time is spent in abstract interpreter / type inference. I guess messing up with @generared functions directly may give me more control. If not, I hope to at least learn something new about Julia internals
As always, writing down a question brings a lot of new ideas. Such as watching the Jarrett’s talk on Cassette design, especially this moment where he explains that generated functions can return either an expressions, or CodeInfo instance. So I tried this:
@generated function transform(f, args...)
ci = @code_lowered f(args...)
return ci
end
inc(x) = x + 1
transform(inc, 2.0)
which didn’t work and just returned an empty Core.CodeInfo[]. The problem of course was that generated functions work not on values, but on types, so f, args are actually (typeof(inc), (Float64,)). Thus we need a slightly different way of getting the CodeInfo:
@generated function transform(f, args...)
ci = code_lowered(f.instance, args)[1]
return ci
end
transform(inc, 2.0)
This version fails with segfault, but in this case it’s a good thing - it means that Julia accepted the new CodeInfo instance, just couldn’t validate it. I know that generating a valid instance of CI can be done using tools from the JuliaCompilerPlugins org, but if somebody wants to jump in with a quick example, you are welcome!
Unfortunately I’m not aware of any remotely accessible and self-contained example in the wild for this, but here are some possible leads. Does overdubbing in generated function inserts inlined code - #7 by Tomas_Pevny has a self-contained snippet, but didn’t seem to work? Cassette’s docs shows how to create a valid CodeInfo output using its API but doesn’t go into details. The underlying code may thus be worth a read. Lastly, the genfuncs in Zygote’s compiler show most of the steps required for ingesting a CodeInfo for a function argument and manipulating it to be compatible with the generated function signature, but you’ll have to read through some IRTools methods to know exactly what it’s doing under the hood and a lot of the AD-related stuff may be a distraction.
With some examples in the directory, but not sure, if it helps. I might improve it in next weeks, as my teaching of that subject is approaching. Bug reports and suggestions for improvement is welcomed.
