Thank you.
So, everyone, I'm Masajardano from UCL.
I'm a shoe from Ganty University.
So, we were talking about how to accelerate
scientific code using AI hardware
and the drill-up package to react into jail.
So, you know that every nowadays,
the AI hardware dominates modern computing world
if you wait a few weeks, you can see some news
about some expanding on AI
from different vendors and different AI providers,
also not just GPUs, also GPUs and other hardware,
and also they affect the world markets
when these things happen.
So, yeah, there is lots of this compute available,
but it's mainly used for AI,
but we want to do also something like using
for regular old-school scientific computing.
So, how do you use these accelerators?
If you use Python, probably you will
rewrite it in Jax or PyTorch,
if you have already some code,
but we are Julia uses.
Julia is a compile language based on LLVM.
So, this will not be an approach for us,
so we need to do something else.
But taking a step back,
any important thing in modern scientific computing
is to do automatic information,
and why this is important.
The derivatives are everywhere in the physical laws
or natural laws.
So, if you can compute the derivatives very fast,
you can accelerate your code a lot,
and it's also useful for machine learning applications
because you can do the back propagation,
or based on inference.
You can see here some applications,
like a folding of proteins,
or reconstructing a Hello Kitty image,
or accelerating the black hole imaging
from the event horizon telescope.
Yeah, so for example, in this case,
they were able to go from analysis,
taking one week on a cluster,
but when they used automatic differentiation,
they were able to reduce the compute time
to one hour on a single core on a CPU.
That's quite a speed up.
How automatic differential differential works
in most pipelines?
So, you take your code,
the source code, which could be written in C++,
really aroused with or whatever language you like.
You apply the automatic differentiation on the source code,
and then that's lowered into some intermediary presentation
of the compiler that can be optimized,
and then you generate the native code.
But this can be inefficient.
A more efficient way to compute derivatives
is the enzyme approach.
And I mean, I am out of this engine based on LLVM,
which works after optimization.
So, you take your code,
you lower to the compiler intermediary presentation,
you optimize it,
and then here you apply the automatic differentiation,
and then possibly you can apply even further optimization
to eventually get to the native code.
And why this is important?
As a case study of how the enzyme approach
speeds up derivatives,
for example, take the normalization of a vector.
Here you have this function, which takes an input and output array,
and you loop over all the elements of the input and the output,
and you divide and you set the output array
to the input array divided by the normalization of the input array.
The problem of this is that you're basically computing
the normalization factor in each iteration of the loop.
The easiest thing to do is to just
place the normalization out of the loop.
But you need to remember to do this, right?
You will select the compiler to do this for you.
So, if you do automatic differentiation before optimization,
the problem is you forget to place the normalization out of the loop.
You do auto-diff,
and so you get the derivative of your function inside the loop.
And so the cost of the derivative is 0 of n squared,
like it was here.
And after optimization, it's not possible to hoist
the factor out of the loop anymore,
because it depends on the iteration.
Instead, if you do what the order of the automatic differentiation
and optimization, you get that you first do the optimization,
so your friendly compiler can hoist the normalization out of the loop.
And then the derivative, in this case,
becomes a 0 of n derivative.
So you get a much faster derivative automatic derivative
by moving the auto-diff stage after optimization.
And as a scientific motivating problem,
to be used for applying all of these concepts together,
was oceananic.
So oceananic is a ocean model written in Julia.
It's a very fast.
It has been used up to 768 GPUs and demonstrating
very high memory and energy efficiency,
compared to other competitors.
And you can also see some cool pictures like this,
very high resolution of the ocean evolution on the globe.
And about GPU programming via LLVM,
which is, again, one of the most common compiler framework,
the problem of doing GPU programming at the LLVM level
is that LLVM IR is very low level.
It doesn't have a good high level representation of the parallelism,
which makes some optimization very hard or just impossible to do.
And also specifically for GPU programming,
the problem is the device code and the host code
live in two different separate modules.
So you have the host code calling a basically black box
on the device code and on the GPU.
And there is no way to do optimization between the device and the host code.
So you cannot, for example, again,
host some functions out of the loop inside the device code
because there is no communication between these two.
And if we look more deeply at the scientific codes,
like in an ocean ion,
we can say about 270 kernels, which look like this,
where you are doing basically stencil code.
And this is like the most natural way to write the code
if you take like the paper formulas.
But if you squint very hard,
you realize that this stencil code are actually convolution.
And much in learning accelerator are actually very good
in convolution, like a tip used in particular,
where the time to do convolution very efficiently.
So one way will be you write your code to express
this stencil code into convolutions,
or you hope that your friendly compiler will do that for you.
Again, this can be difficult to recognize this pattern
and if there are optimizations, which require communication
between the host code and the device code,
like in the previous slide, this can be very hard to do.
Instead, MLAR, MLAR is a compiler framework,
which is based on MLVM, but it's more high level.
So it has a better understanding of structures in your code
and also parallelism and crucially in MLAR,
the host code and the device code are actually in the same module.
So if you take this,
MLVM can erase into MLAR,
then you can do the kind of optimization that you were saying before,
like hosting a function out of a loop,
when you know that this is constant.
And also, MLAR is aware of parallelism
and so can do optimization in this way.
And if you apply derivatives or automatic derivatives,
notative at this MLAR level,
you also get much faster derivatives,
which as you have seen before, can provide a big speed
to your application.
And now, Jules, we'll talk about the Java package,
which brings all of this together.
Yes, so MLAR is very cool,
but you need a way to actually generate it.
So reactant.jl is a Julia package that compiles
Julia code to MLAR using the,
than the MLAR code,
when it can be compiled using the XLA compiler.
And we also have these nice optimizations on top,
and we can do automatic differentiation.
There's actually two ways to get MLAR from Julia code.
The first is to take existing kernels
and to automatically erase them to MLAR.
So in this example,
we have this stencil kernel.
It's lower to LLVM using the normal compilation flow.
And then we have a pass that raises this to
vendor, agnostic MLAR operations.
And at this point, we try to recognize
scalar operations and group these together
in these TensorFlow level operations.
And then we do additional optimizations.
And so you can see throughout this compilation flow in the end,
we end up with a convolution operation.
And this can be executed very efficiently on,
for example, Google TPUs.
The other way to generate MLAR is to use
reactant-straising engine.
So you can write a regular Julia function.
And then at the call site,
you can pass it special array types.
So two RRA creates a reactant array.
And this will make it such that every function
invocation with reactant array
operant will actually generate an MLAR
operation in the compiled MLAR.
So this approach doesn't only work for array operants.
You can also have arbitrary structures.
For example, in this function,
there's a function that takes two points.
And each point has two fields.
And when you compile this, when you trace this to MLAR,
the MLAR code will have flattened each point
into its two constituent sensors.
It's also possible to return structs.
And in this case, the flattened representation will be returned
by the MLAR code.
And there's a bridge to Julia that will reconstruct the structure.
It's also possible to trace control flow.
So in this case, we've added this trace macro in front of a four loop.
And this will make it such that there's a loop in the generated MLAR code as well.
If you don't include this trace,
it will actually kind of unroll the loop by tracing through each iteration
and adding the operations.
There's an asterisk there,
because this is a bit of a flaky system,
that we're working on doing this automatically.
So generating these loops without the annotation.
Lastly, some people might know that stable each hello is a pure format.
Each operation doesn't have side effects.
So it's not possible for an operation to mutate its operands.
While in the Julia code, this is possible.
The first element of its operand 20,
and it doesn't return anything.
But in the traced MLAR code, this can't be expressed.
So there's actually an operation that creates a new array,
where this first element is mutated.
And then it returns it.
And again, in the binding from Julia to the compiler backend,
it will insert this result value into the original operand.
This was just a very quick overview of what reactant is capable of.
But so there's more.
There will be a lot of optimizations on top of the traced code.
It's also possible to automatically distribute the execution.
You can automatically differentiate using runtime.
Like, Mozilla already talked about.
And I focused on the Julia front end,
but it's actually possible to do this with C++ code as well.
So this is what we're looking at with ocean anagons.
So we take existing Julia code from ocean anagons.
And we compile it using reactant.
And then we can automatically distribute it on a very large number of Google TPUs.
So we're still working on performance,
especially for multi-note cases.
We want to get this line down,
but we're getting there.
So yeah, I maybe want to go into too much detail on the conclusion
to get time for questions,
but reactant is very cool and it's open source so you can try that.
Thank you.
I think it's reliable.
So ocean anagons has hundreds of kernels.
And these are successfully raised.
Other people are using it and it works.
And if there's a book, you can report it.
Or like, if there's something missing in the raising pipeline,
you can report it and we'll add it.
So we have hundreds of kernels.
And these are successfully raised.
Other people are using it and it works.
And if there's a book, you can report it.
Or like, if there's something missing in the raising pipeline,
you can report it and we'll add it.
So you're, did reactant lower down into an amlyard.
Like those, right?
Composer, is it an amlyard.
Yeah.
So reactant lowers into stable ratio.
HLO and the question is whether this is like,
you can combine this with other dialects.
The answer is not really stable.
HLO is like, it's an amlyard.
I like, but it contains everything you need.
So you can also target different dialects.
I think, like for example, tryton.
There's like some experimentation going on,
but like stable echelos is the main thing.
Yeah. So the question was whether we can also target the MPI,
amlyard dialects, yes.
We haven't tested on a large scale, but the support is there
to translate MPI operations on the Julia code,
translate like, like, rise them to the MPI amlyard dialects.
No more questions?
Let's get started.
Yeah.
So the question is whether,
if you do state operational arrays,
is that from the command of original?
Okay. So the question is whether it's recommended to do a mutation
on a race or not.
I think it's actually a good style, so especially.
So if you act on sidestep the Julia garbage collector,
because it does all the memory applications on the,
like a mellar level, so it doesn't use the Julia runtime.
So in a sense, like in normal Julia code,
you would like to reuse as much code as possible,
so use mutation a lot, but you can still do it
and still find it because I can, or you want to.
Yeah, maybe also, I think if you do mutation within a function,
the back end compiler will actually realize,
in many cases, that you can reuse the memory,
so it will actually not not like materialize too much.
Yeah, so it does the automatic donation, right?
So the lucky installation of this concept of donating an argument
to a function, to reuse it, and reactant is able to do this process automatically.
Yeah?
So you mentioned that you could get slide at the point
that when it, like there's a boundary where it goes into M.
I already unwrap and then back to Julia and just do that.
Like if you have multiple functions that are all using the same thing,
let me know.
So, yeah, the question is whether, if you have multiple functions that take a point,
whether it needs to generate the code multiple times,
so the code to unwrap and wrap will actually be cache type belief when you run it,
but if you compile a reactant function, you actually compile like the whole reactant program,
typically, so you only unwrap at the beginning and wrap at the end,
so during execution, you can call different functions without any unwraping happening.
So the question is whether it's all ahead of time,
whether you can use Julia's jet, so you can use Julia's jet.
If you read the final function and you recall it with the adjet macro to have it reactant,
it will actually recompile the function, which is kind of what Julia does as well
if you run a function, it's ahead of time, compiles it just in time.
Thank you very much.