Recursive `@code_lowered`

General Usage
1

We can get IR for a single call using @code_lowered:

foo(x) = x + 1
bar(x) = 2 * foo(x)

@code_lowered bar(2.0)

# CodeInfo(
# 1 1 ─ %1 = (Main.foo)(x)                                                    │
#   │   %2 = 2 * %1                                                           │
#   └──      return %2                                                        │
# )

But how can I dive deeper and get IR for a called function , e.g. Main.foo(x)?

2

Unfortunately, Sugar still doesn’t support Julia 1.0, but perhaps following their approach + MacroTools / similar to walk the expressions?

You would need code_typed + type stable functions, though. Otherwise, if there are different method definitions for Main.foo, which do you expand?

3

Do you desire this for debugging or code transformation?

This kind of “recursive reflection” is exactly how Cassette works; you might be interesting in https://jrevels.github.io/Cassette.jl/latest/contextualpass.html.

Even if you’re only using it for debugging, you can add side effects to an otherwise “no-op” Cassette pass just to e.g. print out the CodeInfo.

If you’re fine with doing things manually and just want a quicker interactive workflow, I find GitHub - timholy/Rebugger.jl: An expression-level debugger for Julia with a provocative command-line (REPL) user interface is helpful for this sort of thing.

4

Ah, how could I forget about it! I used Sugar a number of times to find original source code, but totally forgot that it works with IR even better. Thanks for this reminder!

5

It’s for code transformation. In fact, I’m trying to estimate effort required to implement something like Zygote or Capstan, but simpler (I’m trying to make a huge focus on simplicity these days, even trading it for a number of features). My latest beast in AD zoo is Yota.jl which is based on old school method dispatching (very much like ReverseDiff), but maybe it’s time step forward in complexity again.

I tried Cassette like 3 or 4 times, mostly before Julia 1.0, but then it was too much of a moving target for me to rely on it. Now I still see a disclaimer it its docs:

For now, each individual version of Cassette can technically only support a single specific version of Julia at a time ; differing by even a patch version could (theoretically) break Cassette entirely.

Cassette enables interaction with many parts of the Julia compiler, a lot of which are undocumented or sparsely documented.

Is API and runtime of Cassette stable enough for somebody not involved into development of Julia compiler or Cassette itself to implement AD on top of it?

2 Likes
6

Ooh, cool, I’ll check it out!

Actually, that specific disclaimer is now out of date, since Base Julia has committed to reverse-dependency testing before new releases (which will catch compiler changes that affect Cassette, so we can patch things ahead of time). I’ll update that when I get the chance.

Is API and runtime of Cassette stable enough for somebody not involved into development of Julia compiler or Cassette itself to implement AD on top of it?

The documented API (basically, everything but the “contextual tagging” stuff) is stable. Performance is fairly good these days in a lot of cases, though we know of some edge cases. Stressing it on real code is the best way to find those edge cases, of course :slight_smile: Feel free to ping me if you have questions!

4 Likes
7

Sounds promising, thanks! Perhaps this is indeed the way to go.

I guess for AD-like stuff I need to tag functions arguments, but the corresponding part of manual isn’t available. What is the most recent write up on the topic besides source and tests? Even if not up to date, it might be useful to read it for better conceptual view.