search

tshort / StaticCompiler.jl

Compiles Julia code to a standalone library (experimental)
Other
489 stars 31 forks source link

Cassette-type functionality with Mixtape #69

Closed tshort closed 1 year ago

tshort commented 2 years ago

It'd be nice to be able to change functionality when statically compiling. This PR includes code from Mixtape to do that.

In the bigger picture, there are different options to support this feature:

  1. Wait for the AbstractInterpreter and the compiler-plugin interface to settle. I haven't seen updates to https://github.com/JuliaLang/julia/pull/41632 in a while.
  2. Add functionality like this to GPUCompiler.
  3. [This PR] Add functionality here, and switch it later based on advancements in option 1 or 2.

Pinging @femtomc for opinions...

Notes:

There are lot's of possibilities to use this with StaticTools. This PR includes an example where Strings are replaced by StaticStrings.

femtomc commented 2 years ago

@tshort I really like this idea - this was something I tried early on, but I didn't understand the static compilation process very well.

Re -- @vchuravy: is it possible that we can chat about fixing the OrcV2 linking questions in the next few weeks?

femtomc commented 2 years ago

Re -- @tshort: if StaticCompiler is using GPUCompiler anyways (and doesn't intercept the native execution engine somehow), then the plugin interface won't necessarily offer anything over Mixtape.

femtomc commented 2 years ago

The biggest issue with Mixtape is that it uses AbstractInterpreter -- which has the problem that composition is not easy. So e.g. running Mixtape passes and then doing an AD pass is something which there's not really an idiom for. I'd really like to alleviate this issue in general (but I've been busy doing other research).

MasonProtter commented 2 years ago

Nice, this is a great idea.

tshort commented 2 years ago

The main holdup this Mixtape support is getting it to work with Julia v1.8. I haven't gotten far on that.

femtomc commented 2 years ago

@vchuravy the hold up is likely an ABI thing, as well as maybe a "move to the new OrcJIT" thing

tshort commented 2 years ago

@MasonProtter, with this one tweak from your Orc v2 PR, the StaticCompiler tests progress much better. Only nightly fails immediately.

MasonProtter commented 2 years ago

Okay, that’s promising!

brenhinkeller commented 2 years ago

I wonder if this could work on Nightly now after the update to GPUCompiler 0.16? Or is GPUCompiler working merely a necessary but not sufficient condition for Mixtape.jl working?

tshort commented 2 years ago

I wonder if this could work on Nightly now after the update to GPUCompiler 0.16? Or is GPUCompiler working merely a necessary but not sufficient condition for Mixtape.jl working?

I'm hopeful, too. I gave it a quick try, but it needed more work than I had time for.

tshort commented 1 year ago

Here are some updates on Mixtape support. This branch is caught up to master, and tests pass. I used the AbstractInterpreter straight out of GPUCompiler with minor tweaks. That should make it easier to keep up with GPUCompiler, and if it works, we may want to upstream it at some point.

Some things that still need work:

tshort commented 1 year ago

I think the Mixtape support is ready to merge. But, note that it is experimental and hasn't been tested much. Any other comments prior to merging? @brenhinkeller or @MasonProtter?

brenhinkeller commented 1 year ago

Looks like tests are all passing!

MasonProtter commented 1 year ago

Great that you got this working. I think most of the functionality here is already covered by the new @device_override stuff, but I think this is probably useful to have if anyone needs to do more complex IR transformations than just method replacement.

tshort commented 1 year ago

Great that you got this working. I think most of the functionality here is already covered by the new @device_override stuff, but I think this is probably useful to have if anyone needs to do more complex IR transformations than just method replacement.

Yeah, overriding methods is sure easier.