Refactor compile and test buffer donation

[mofeing]

This PR refactors compile in a series of smaller functions that makes it easier to test different parts of it. In particular, we use it to test buffer donation but could be used to test other parts as well.

The new functions are:

• compile_mlir This is the old compile_to_module! but added a new method such that it creates and activates the MLIR context by itself and returns the MLIR module.
  • This patter was used in several ocasions (i.e. code_hlo and compile) with slight variations.
  • We don't really need to store the MLIR context (so no need for CompiledModule).
  • Also the MLIR context is only needed in the beginning but it's activated through all compile... so this just activates the context where is needed.
• compile_xla Calls to compile_mlir and passes the result to XLA for compilation.
• codegen_flatten Generates the Julia code that unwraps the input arguments into the XLA buffers (i.e. the linearization).
  • The _flatten suffix is a reference to the flatten function passed to register_pytree_node in JAX, which has a similar functionality.
• codegen_unflatten Generates the Julia code that rewraps the processed XLA buffers into the result datatypes (i.e. the delinearization).
  • The _unflatten suffix is a reference to the unflatten function passed to register_pytree_node in JAX, which has a similar functionality.
• codegen_xla_call Generates the Julia code that calls the XLA-generated executable.

This way, buffer donation can be checked just using compile_mlir and then checking manually against the returned linear_args and preserved_args.

The refactor of compile is due to the fact that I had to debug it many times, and it's hard to follow without more structure.

I'm trying to rebase #123 since it changed some stuff on compile.

[github-actions]

Benchmark Results

	main	5707adf...	main/5707adfcfcbcf0...
comptime/basics/2D sum	26.2 ± 0.98 ms	26.3 ± 0.91 ms	0.995
comptime/basics/Basic cos	24 ± 0.42 ms	23.9 ± 0.43 ms	1
comptime/basics/Basic grad cos	0.0436 ± 0.00068 s	0.0435 ± 0.00054 s	1
comptime/lux neural networks/vgg11	0.365 ± 0.0023 s	0.375 ± 0.025 s	0.973
comptime/lux neural networks/vgg13	0.417 ± 0.01 s	0.422 ± 0.026 s	0.99
comptime/lux neural networks/vgg16	0.499 ± 0.014 s	0.492 ± 0.0045 s	1.01
comptime/lux neural networks/vgg19	0.584 ± 0.021 s	0.573 ± 0.0076 s	1.02
time_to_load	1.49 ± 0.01 s	1.49 ± 0.0089 s	1

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

[wsmoses]

LGTM just ome minor nits