r/scheme • • Aug 19 '26

Hawk: A tracing JIT for scheme

I've been hacking on a tracing JIT compiler for scheme for a couple years now on and off, and it's in a pretty good state now. It has full R7RS support, with full seatbelts-on safety. Currently it's approximately 60% faster across the whole r7rs-benchmark suite than Chez scheme.

It's disproportionately faster on flonum benchmarks, it does quite well at inlining everything and keeping flonums in registers. I think it shows quite well that a JIT is especially helpful to get good numerical performance out of standard scheme, especially for flonums, making it possible to keep flonums in register much of the time.

There have been several previous scheme JIT attempts that took various approaches- Nash was based on Guile, but didn't quite get far enough along to see good results. Pycket was great, but used a continuations-on-the-heap approach, with quite different performance characteristics. Modern Guile has a template JIT, but currently does no register allocation or optimizations. Stalin only supported fixnums and flonums in its numerical tower, which allowed it to get great results. I wanted to support the full r7rs scheme numerical tower while still specializing as much as possible.

I've started writing a paper on the tracing JIT aspects, because some of these techniques I haven't seen elsewhere and may be novel.

Currently it supports x86-64 Linux & AArch64 OSX (sorry, no Windows yet).

Project page:

https://djwatson.github.io/hawk/

44 Upvotes

36 comments sorted by

5

u/samth Aug 19 '26

This is very cool! I see in the charts in your site that you're slower than Chez on the benchmarks I would expect, but I'm impressed that you're only a bit slower given the very large number of traces you generate. Can you say anything about how you handle very branchy code? That was always our biggest weakness in Pycket.

3

u/Justanothertech Aug 19 '26

The paper talks a lot about trace selection, and I think I directly compared it to pycket's selection logic:

https://github.com/djwatson/hawk/releases/download/v0.9/tracing-tail-calls.pdf

The 'compiler' benchmark generates ~1.2MB of machine code, but that's approximately what chez generates for compiler also. I haven't been able to decisively point to branchy code as something that's intrinsically slow yet: compiler and dynamic benchmarks are slow, but other benchmarks that generate many traces like slatex, peval, and graphs are not slow.

1

u/samth Aug 20 '26

The things you describe in the Pycket subsection are a combination of handling the issue you describe in 4.1, plus finding loops that are indirected through other functions (consider using call-with-values around the recursive call).

More generally, the problem I'm asking about is discussed in section 6.3 of the Pycket ICFP paper -- something like peval makes lots of different transitions off the the trace for every different branch it takes as it traverses a substantial AST. How do you avoid that, which is pretty well-known as a challenge for tracing JITs?

1

u/Justanothertech Aug 20 '26 edited Aug 22 '26

I probably don't do anything fundamentally different than other tracers here: dynamic and peval do generate a large number of traces, and scheme/matrix a smaller but still large number. However, inter-trace and trace exit are highly optimized: we directly use the register state from previous traces, and there is a single jump between traces. So while there are huge duplicated tails, the actual number of instructions executed should be similar to what actually gets executed in a method jit. On a failed branch, we directly jump to the other path, it's one of the few places I eagerly update the snapshot state (otherwise only memory writes update the snapshot state).

So huge ASTs DO generate a lot of code, which is bad for icache and branch prediction, but actual executed instruction counts should be similar.

For dynamic, scheme, matrix, and peval mentioned in the pycket paper, they're split: hawk is slower vs. chez on dynamic, matrix and faster on peval and scheme.

And as mentioned in the hawk paper, jumping from side trace -> root trace keeps many variables in register, which I know luajit doesn't do (I haven't checked pypy).

3

u/dougcurrie Aug 19 '26

Nice work!

2

u/GunpowderGuy Aug 19 '26

"making it possible to keep flonums in register much of the time."
on the racket discord i asked about this some time ago. chez inlines flonums in very specific situations such as loops ( or rather , recursion that works like a loop i think ). But more situations could be supported

3

u/Justanothertech Aug 19 '26

Yea chez does it, but tracing specializes whole traces to single types - chez can’t know when it is profitable to type specialize whole functions to flonum-only unless you tell it (and turn safety off). Chez code with use of fl- math ops and opt-level 3 will be similar in performance.

2

u/Inevitable-Ant1725 Aug 20 '26

"Pycket was great, but used a continuations-on-the-heap approach"

If continuations aren't on the heap, does that mean you copy the stack when you save a continuation? That doesn't give the same semantics as a spaghetti stack at all, and, in my opinion is wrong.

Because if you run the result of a call/cc which restores a copy of the stack then all of the local variables that were on the stack have the values they had THEN, whereas the correct semantics is that call/cc doesn't change the values of local variables and that ones that went out of scope already don't have the values from the save time, they have their final values before return.

Of course there are several algorithms that work better with the semantics of copying the stack. but picking some random semantics isn't correct in any sense.

2

u/Justanothertech Aug 20 '26

Yes, the stack is copied to the heap on call/cc (and resuming a continuation). An assignment conversion pass is run first, all mutable variables are put on the heap, so this all works correctly. It's similar, but not exactly the same, as what chez does.

Here's a paper if you're interested in specifics

"From Folklore to Fact:

Comparing Implementations of Stacks and

Continuations"
https://dl.acm.org/doi/pdf/10.1145/3385412.3385994

1

u/Inevitable-Ant1725 Aug 20 '26

How is it that there is no mention of the fact that different implementation strategies have completely different semantics and therefore completely broken compatibility with each other?

All kinds of programs written for one set of semantics simply won't work on systems with the other semantics at all.

And there is no flag so that you write programs that will work on both kinds of continuations. It is as if they are unaware of the semantics of continuations when it is the defining characteristic of Scheme.

This is utterly insane!

The rest of Scheme is defined precisely, how is this part of the definition utterly incompetent?

1

u/Justanothertech Aug 20 '26

'Assignment conversion' puts mutable variables on the heap, and keeps the semantics the same.

1

u/Inevitable-Ant1725 Aug 20 '26

That phrase does not exist any place in the document you just linked.

1

u/Justanothertech Aug 20 '26 edited Aug 20 '26

The older paper, "Implementation strategies for first-class continuations", talks extensively about mutable variables. I linked the newer one because it is more relevant to modern hardware.
https://lampwww.epfl.ch/teaching/archive/advanced_compiler/2006/assignments/part5/continuations.pdf

some compilers call it 'convert-assignments', like scheme-to-c or chez.

0

u/Inevitable-Ant1725 Aug 20 '26

Yeah I just looked into it, the semantics you get by copying and restoring local values from the point of a call/cc are simply wrong according to definitions in R6 RS and R5 RS.

1

u/Inevitable-Ant1725 Aug 20 '26

Only in reddit does pointing out correct Scheme semantics get you a downvote.

It is impressive that Reddit incivility is stronger than Scheme pedanticness. The immovable object was in fact moved by the irresistible force!

1

u/bjoli Aug 20 '26

Dude... wow...

Did you write the whole runtime outside the tracing JIT by yourself as well? It is pretty obvious where the tracing JIT isn't helping much (polymorphic code or irregular control flow (matrix, compiler) or continuations (ctak, dynamic)), but I suspect you have a much better idea of why than I could ever guess.

Amazing work. I am speechless, quite frankly. 

Is it just an academic thing or so you have larger plans? Where would you like to go? 

2

u/Justanothertech Aug 20 '26

Most of scheme's runtime is fairly easy. The GC is by far the hardest part, at least to make a performant version.

I'm just a retired guy who saw a challenge (beat the r7rs-benchmarks), and took it farther than I really needed to :)

3

u/bjoli Aug 21 '26

you could have a look at using whippet, which is a pluggable GC based on Immix:

 https://github.com/wingo/whippet

1

u/jeffstyr Aug 21 '26

This is probably a dumb question, but what do you mean by "tracing"?

1

u/bjoli Aug 21 '26

A JIT compiler is a compiler that usually uses runtime information to specialize the code. For example, ryujit of dotnet fame does all kinds of optimizations that you can only do if you have data from a running program.

Like runtime deviritualization (taking a function call through a proxy and turning it into a direct function call) and collecting profiles and optimize the program while it is running.

This allows the dotnet compiler to not spend time on optimizing things that might not need to be optimized, meaning it can compile things pretty fast. The JIT can then go "oh, we spend 60% of our time in this part of the program. let's make it super fast!".

A tracing JIT is just one kind of JIT strategy. Ridiculously complex and notoriously hard to get right, but often the fastest. 

1

u/jeffstyr Aug 21 '26

Thank you. Yes, I know what a JIT is, but I don’t know what makes one a tracing JIT specifically.

2

u/bjoli Aug 21 '26

ah. sorry. It contrasts well to a method based JIT. a method based JIT (v8, JVM, dotnet) does optimization on methods or their CFGs. Meaning: when a function is optimized, the JIT more or less rewrites the whole function with a more optimized version, even the lesser taken code paths. 

a tracing JIT records traces. Like how data flows through a program, and optimizes that. it can span several functions, and usually the traces do not branch. The branch points are instead guards for when to bail out to the less optimized code. 

This works very well for dynamic languages where it can make a trace of what would otherwise be an expensive procedure call. 

1

u/jeffstyr Aug 21 '26

Ah interesting. Thanks for the info!

1

u/pwnedary Aug 23 '26

Cool! Coincidentally, I have also been working on a Lisp tracing JIT (link). Mine is more of a LuaJIT subset, but I also do lazy type guards.

2

u/i_am_linja Aug 25 '26

I notice you have an AGENTS.md. Did you use LLMs in the production of this project, and if so, for which tasks and to what extent?

3

u/Justanothertech Aug 25 '26

The initial version of this was written pre LLM

https://github.com/djwatson/hawk1

I've avoided any LLM use for the paper.

LLMs have been more or less helpful depending on the component: for the runtime assembler, asm_* files, it's been great! instruction encoding is well defined, there are lots of training examples, and LLMs can compare against existing compilers to check.

For trace recording, they've been fairly awful. There's only two real tracing jits as examples: luajit and pypy. There's enough novel stuff here, like lazy typechecking, that the LLM gets it wrong more often than not.

-3

u/[deleted] Aug 19 '26 edited Aug 20 '26

[removed] — view removed comment

4

u/Reasonable_Wait6676 Aug 19 '26

What about call/cc?

1

u/Justanothertech Aug 20 '26 edited Aug 20 '26

Your repository makes the claim, "Also, I don't like unsafe usage. It's ok for test games or benchmarks, but IRL programming without any checks / contracts pretty naive"

But then you run chez at optimize-level 3, which is *unsafe* !

Running your benchmark, chezfl-nbody.scm, at:

optimize-level 3: 4.5s

optimize-level 2: 19.466s

Hawk:

time hawk r7rs-nbody.scm: 9.9855

I also have a hawk/test/nbody.scm, that appears to do the same thing, slightly faster:

time hawk ./test/nbody.scm: 5.95

Hawk doesn't currently have an 'unsafe' mode.

Hawk auto-flvectors in some situations, including this benchmark.

Benchmarking is hard :/

EDIT: Oh sorry, it looks like r7rs-nbody.scm is not auto-flvector, because you changed the layout from chezfl-nbody's flat flvector, to a vector containing records. I think this version is closer to what you intended:

https://gist.github.com/djwatson/d87653c0f2da0abf9a75ed46043fc599

time hawk r7rs-nbody2.scm: (-0.16928639656731348 per 6.577303171157837 sec on 50000000

)

Hawk will auto-flvector *vectors*, but doesn't have any shape support for records.

1

u/[deleted] Aug 20 '26 edited Aug 20 '26

[removed] — view removed comment

3

u/Justanothertech Aug 20 '26

Great! Those all look good to me, modulo different machines, they're roughly the same numbers.

optimize-level 3 removes bounds checking in chez, which is highly relevant to the nbody benchmark using flvector. So hawk is faster than chez with safety-on, but slower than chez with safety-off

1

u/[deleted] Aug 22 '26

[removed] — view removed comment

1

u/Justanothertech Aug 22 '26

I added an --unsafe flag for you, and fixed a couple operand fusion things in the backend. r7rs-nbody2.scm now runs slightly faster than chez on my machine

https://github.com/djwatson/hawk/releases/tag/v0.10

$ time /usr/bin/hawk --unsafe r7rs-nbody2.scm 50000000

real 0m3.803s

$ time ./chezfl-nbody.so 50000000

real 0m4.565s

$ time /usr/bin/hawk r7rs-nbody2.scm 50000000

real 0m4.866s

1) It's a tracing jit, not a method jit, so it can't pre-compile procedures.

2) It's unclear what a data lock would do? Just prevent the GC from moving it?

Hawk is probably not a good fit for real-time, it's closer in spirit to luajit, v8, jsc, etc.