Expose run_jit to allow for custom behaviour in JIT flow

[tinco]

Hi, I'm making this PR primarily to let you know I'm working on a project that uses the JIT in a slightly weird way, I'm making a language that is JITted with Rust as an intermediate language, or at least for its parts that are linked to Rust crate dependencies.

To understand the JIT better, and to give myself a point to start thinking on what else I need I started by splitting out the JIT driver into smaller functions. I'm not sure if this is desirable, it just helped me understand it better.

My next step I've described in this thoughts:

I need some way to make the lazy jit state publicly accessible. It should only be accessible from the rustc thread. This is achieved by using the message channel to send requests to the rustc thread for the regular JIT flow. I need it to be publicly accessible so I can add new functions to the jit_module, and it has to be in the rustc_thread because I assume it
is thread_local for good reason.

I want to enable arbitrary code to invoke the JIT runtime to compile new functions, so I need to expose a way to send messages to the rustc thread.

So I should add a new message type that requests a function to be JIT to the UnsafeMessage enum. And then I should add a new message handler argument to the run_jit function that will be called when a new message is received.

I would love your thoughts on if I'm going in a right direction, and if you think this has a chance of getting merged or if I should just continue in my fork without worrying about upstreaming.

Thanks!

[bjorn3]

Hi, I'm making this PR primarily to let you know I'm working on a project that uses the JIT in a slightly weird way, I'm making a language that is JITted with Rust as an intermediate language, or at least for its parts that are linked to Rust crate dependencies.

Sounds cool! It may be hard to implement though as rustc isn't really designed for this.

I want to enable arbitrary code to invoke the JIT runtime to compile new functions, so I need to expose a way to send messages to the rustc thread.

If you want to define entirely new functions, you will have to create a new compiler session every time you add a new function. You could still reuse the same JitModule however to avoid having to codegen all functions again every time. https://github.com/rust-lang/rustc_codegen_cranelift/tree/wip_hot_code_swapping is an old branch with some experiments regarding hot code swapping, which is almost exactly what you need to compile new functions as opposed to lazily codegening on the first call as implemented in the lazy-jit mode.

To understand the JIT better, and to give myself a point to start thinking on what else I need I started by splitting out the JIT driver into smaller functions.

I would prefer if the entire JIT driver stays in src/driver/jit.rs (or if necessary submodules thereof) rather than scattering it around. That said, moving predefine_mono_items to src/base.rs makes sense to me, though please move it above CodegenedFunction. Feel free to make a separate PR for that change.

As for making more things public, I'm a bit split. While I do like the idea of your project, I fear that it will make refactorings to cg_clif that are beneficial to the core objective of cg_clif harder.

[tinco]

I want to enable arbitrary code to invoke the JIT runtime to compile new functions, so I need to expose a way to send messages to the rustc thread.

If you want to define entirely new functions, you will have to create a new compiler session every time you add a new function. You could still reuse the same JitModule however to avoid having to codegen all functions again every time. https://github.com/rust-lang/rustc_codegen_cranelift/tree/wip_hot_code_swapping is an old branch with some experiments regarding hot code swapping, which is almost exactly what you need to compile new functions as opposed to lazily codegening on the first call as implemented in the lazy-jit mode.

Thanks a lot, I didn't know this and would've spent a bunch of time figuring it out without this advice. I think it might mean I have to rethink the strategy a bit too, recreating the compiler session means running the full analysis pass again right? That sounds kind of expensive for something that's supposed to happen as a JIT/optimizing behavior. If it's really going to be expensive then I might take a more traditional approach and rely more on cranelift itself and less on rustc, at least for the parts that don't directly interact with Rust crates.

As for making more things public, I'm a bit split. While I do like the idea of your project, I fear that it will make refactorings to cg_clif that are beneficial to the core objective of cg_clif harder.

Totally understandable. My main objective would be to expose the JitModule and use it to define new functions. But if doing that isn't compatible with the way the rustc backend works, then probably I'm better off keeping my own module next to it. Maybe I could even rely more on some other dynamic language implementation built on cranelift.

[bjorn3]

recreating the compiler session means running the full analysis pass again right?

Assuming that you finished the previous compilation session cleanly (which the current jit mode doesn't do as it exits with the exit code of the jitted program before reaching that point), you will get a fair bit of reused work from incr comp, but it is still likely not cheap enough for your purposes.

If it's really going to be expensive then I might take a more traditional approach and rely more on cranelift itself and less on rustc, at least for the parts that don't directly interact with Rust crates.

That is what I would recommend doing.