So, next talk, I have a pleasure to introduce Ramon, who will be talking about closing the loop,
and self-living compilers.
I've heard today that there is an NVIDIA published AI agent generated AI generator for
to generate models, which is six times slower than the one written by people.
So, let's see if we can do a much better work with self-learning compilers for AI.
Let's see.
Hi.
So, my name is Ramon Rush.
I represent Daisy Tuner, where I work as a compiler engineer, and I'll be presenting,
closing the loop, a self-learning compiler for AI accelerators.
Okay.
So, I think it's pretty safe to say we live in the age of hardware accelerators.
Early a month goes by where there's not news about a new accelerator or even a new architecture.
So, what does that mean for the user?
First, we have to choose what kind of accelerator would we want to target if we determine
that we need acceleration.
And then there's always the balance between portability and specialization.
You can remain agnostic to the actual hardware accelerator, easily swap them out if you
remain on the right side, but you'll properly sacrifice a lot of performance, and you
can get in the best performance on the left side, but probably looking at vendor lock-in or
API lock-in.
So, let's look at like an example stack, a user that comes with a portable system that hasn't
used in accelerator before, the problem proves that it works on complex problems.
And we still can here see the balance between generality and specialization.
If you want the best performance, you would write directly for the instructions at architecture
of the accelerator, but the heart of this problem is the mapping.
Here we can get all the performance or lose all the performance, but I'll get back to that
later.
First off, let's look at what we want to do.
We want to build a toolbox that allows for optimized mapping between all sorts of different
accelerators from basically any kind of input language or format.
Yeah.
So, let's look how we do with that.
First off, the representation on which we think, so you get a better idea on which level
we're tackling that mapping.
So we call our representation structured data of a stateful data flow multi-graphs.
Let's start out with a simple data flow graph.
We have access notes, which represent variables in this case.
We have tasklets, which is the basic math operation, and we have the edges, which we call
memlets.
This would be a load.
Here would be a store operation.
We can make this more complex.
We can add dependent data flow on any very simple, very easy to understand this format.
We can add parallel execution in that graph.
We can add offsets and operate on arrays instead of scalar variables.
Still looks pretty simple.
Then we can wrap this in a loop.
In this case, it would be a data parallel loop.
And the memlets are now indirect accesses, and we can take this further, wrap it into
more loops.
Then we have the naive math implementation.
Here the memlets are multidimensional accesses with a types attached technically.
So we can take this one abstraction further and just represent a call.
In this case, our memlet edges become just address calculations, no memory access, and
the call could be anything to other graphs, to external functions, in some kind of environment
or interendicts, which are just well-known functions, like gem that we could replace with
the underlying code, or just call something else.
Then lastly, we have conditional execution.
It's also pretty simple, and you can chain all of these elements for larger programs.
To summarize, the format expresses data flow and control flow at the same time.
It's structured.
You can see the loops directly, so that makes it very convenient if you're talking about
optimizing, transforming, changing loops.
It's a lot more tors than other formats with us, and it's simpler and easier to analyze.
What it doesn't allow is a regular control flow.
So you cannot represent a program that contains just like branch help.
You need to simplify that, but we don't think it's a large problem, because most pseudo-code
is structured and doesn't have that problem.
Also, you might have noticed the address and condition calculations are in a different
format than the data flow, especially for parallel problems.
This is very convenient, and the more separate they are, the more parallel the problem
tends to be.
Okay, so what are the elements that we have in our toolbox so far?
So we have our SDFGs, APIs to build them, we have serialization so far it's just JSON,
more to come.
Then we have analysis passes similar to compilers that operate on loops, data dependencies
and parallelism, and we have optimization passes that can optimize the loops, transform
it to parallel execution, and so on, and right now our output is CC++ code generation.
So this way we are currently very agnostic of what comes behind it, any CC++ compiler
could compile it to an architecture, just like what I showed in the stack graph.
In terms of front ends, right now part of our library are Python bindings, so you can build
those graphs from Python, and we also have a decorator, and let's look at example.
If we have this simple code, we can annotate it with our annotation, and when it executes,
our library will pass that code, transform it into our graph representation, and if you let
it do a default work, it will output CC++ code compile it, load it into the runtime again,
and execute it, so it will just make simple Python code faster, so far, but it's a good
platform to experiment and look at what the representation does.
So this would, for example, be our graph, that then executes, okay, we also are working on,
I believe it's been pushed today into the git, it's not really doing much so far, but
more to come, a front end to the ingest PyTorch models directly without having to go to Python
code, and also working on ingesting MLIR graphs directly. On the target side, we have the generic
host target, which basically uses the normal CC++ code generation, so basically if you want
to execute it on the CPU, we have an open MP back end, which basically adds open MP annotations
to the output code, so it runs on multiple cores, we use Google Highway for vectorization, so we can
use that, and we also have a CUDA back end, so take a look at that, okay, so for CUDA, so far we had
normal code linear code runs on the CPU, now we want to offload that, so we need to cut out a portion
and we need to generate code that runs on the device, we need to run the, generate the host code
and we need to generate memory management and transfers to the device and back, as a graph,
this would look something like this in the middle, the same mutmal graph from before and just the
memory transfers above and below it, and what we also added on top is optimizations around the
data structures, so if you were to execute this twice, the second one on the results of the first,
it would be a ways to transfer that data back to the host and back to the accelerator again,
and we can just remove those transfers, similar with read-only data that has never changed,
it can just remain on the device, and this optimization is also generic, it could apply to any
accelerator that is modeled as a target for this library, okay, so let's look at how well this works
so far, this is NP-bench, popular Python benchmark, uses a few different back ends also with decorators,
just as us, so they just in time replace the code that executes and measure the performance,
and NumPy is the baseline on the bottom in milliseconds and everything else is relative to
NumPy green indicates speed up, red indicates a slowdown, so these are just the frameworks that come
with NP-bench, and we run it on paper size, okay, so this is baseline, we don't do any crazy
optimizations yet, it's basically just cleaning up what we get from Python, and the C code we compile
with a clang on O3, and you can see we are achieving some speedups, there are a few gaps where
we don't map all operations just yet, but we're actually pretty happy with how good that works,
if we enable OpenMP, it still works, right now we're not seeing that much speedups because
it's annotation based, so if it's just a library call that's emitted, the annotation won't actually
do anything to paralyze it, and the CUDA back end works and can upload code that was not optimized
for any GPU or anything, and we are also pretty happy that we're beating in some cases the CUPA implementation
already. So, but as I said in the intro, mapping is the heart of the problem, so how do we map
something such as this to various accelerators? For example, here a really a CUDA architecture,
which is actually pretty similar to a normal CPU in terms of memory topology and address space,
it's one shared address space, any processor can access any data, it just determines the speed in
which it does it. But we can also go crazy, and for example, look at the tense-turned architecture,
which is just a mesh of many connected cores or need to sum up the neighbors, so here if we
wanted to get access to some data, we need to move that data through all those cores,
we need to consider all the bottlenecks we created on the way, on each core, and even if we get the
data to its target, it's still not a normal core, it's still five cores that need to move
on stream data around between the various accelerators. So, the mapping is at its height
about parallelism, on the different levels, CPU cores, threats, vectors, it's about the memory
topology and locality, and it's about data movement, and you can also further do data format
and position as we have before. And if you do all of this right, you get a good performing corner.
Right now, the standard way to do is to basically build a library of hard-coated or hand-optimized
kernel from experts for a given accelerator. Very popular math functions, you can just build a
library for a few of those and have good performance. The hard-coated, agnostic way, sorry,
I turned those around. Yeah, there are already solutions to do this hard-coated, agnostic way
that automatically map this, but they don't get close in performance to the manual performance
engineering. Okay, so to explain manual performance engineering, let's look at the short example
again, the simple mathematical application, and this would be the Nvidia example how to run
that fast in CUDA, and just to show you what goes into it. So, here's code to initialize the hardware
on allocate memory. We have data transfers to the accelerator because it's a separate memory.
We have to match the exact hardware size we have. Here, this hides the two outer loops because
the hardware managers are those. Then this is where the code that runs on the accelerator starts.
We allocate a new memory and pre-fetch data into that memory before we calculate it. We pre-fetch
in a different order to maximally utilize cache lines, so we fetch all the data linearly.
We still accumulate in the existing order. We need to manage multiple threats that we didn't
have before, and we're saving us of right and to memory. So a lot of stuff, but also we can break
this down as a collection of building blocks. Most of those optimizations are like textbook optimizations.
They have parameters and you to figure out on what to apply it and when, but it's basically building
blocks. So if you are enough of an expert in the hardware, you can pick the standard text book
transformations, pick their parameters, pick on which parts of the algorithm you apply to,
and then the order of it, and then instead of looking at the complex kernel I showed on the right
side, you can look at the basic pseudo code on the left side and a recipe that is specialized
for the algorithm and the specific hardware. We'll get you the same result, but this way you can look
at the base algorithm and the recipe is what determines what hardware you can run it on.
And as long as you get the recipe at a run anywhere as fast as it can.
An example for the recipe left the simple code and the transformation we're applying is loop
tiling. So we're splitting up loop into two. The outer loop will make bigger jumps, blocks of
32 in this case, and the inner loop will iterate over the elements sequentially as before.
And we can do this to the outer loop as well. Then for example, we could run loop interchange,
we just swap two loops. They don't have their own code, so we can do that. Without affecting
the inner most code, we get that and why do we do this? Okay, now here we can see the green loops
basically just operate on the rect angular on the right side. And the outer loops work over the
rest of the data. So what that gets us is very localized data, which for example could fit into
a cache of one processor core. So basically this is what we would run on the multiple cores,
and the inner code is what would run on one of them cores, to optimally paralyze something like this.
And we can repeat the same concept on the inner loop to exploit vectorization, for example.
But where do we get the fast recipe? We talked about how we can apply it, but not where we
do actually get that. We have the building blocks. Okay, there is a standard solution for this,
it's called outer tuning. But what outer tuning is is basically just trying all the possibilities,
spend a lot of time on it, and eventually you find the best performing solution for whatever
performance target you had. This takes a long time, it's infeasable to do this every time you
want to compile something or build an application. So what we really want is we want to cache on
this remote server that you don't have to see those results, the recipes. So once we have it,
everybody can use the recipe without having to know anything about the accelerator, and as long as
the recipe exists, you can switch the accelerator and gain the best performance for that hardware,
as well. Okay, so, but there are many accelerators, many different architectures, and even for
the same architecture, the optimizations can change, because there are microtactual implementations
like cache sizes, which change how to get the best performance out of it. Okay, so what we want
on top of the cache is similarity based to search. We don't just want to find the exact match,
we want to find any algorithm that was similar on similar hardware and get the recipe for this
because it will get you similar performance on code we haven't seen before. This would look
something like this, a database of similar algorithms, and we can port over the
optimization recipes from a slightly different code. Okay, so we are using a neural network to
classify the code, so to get the coordinates in that diagram, and we store for each point in
the diagram, the optimization recipe, the results we observed, and we can do that with multiple
target architectures, and then you basically need a server farm to do auto tuning nonstop
for every architecture that you want to support to find the recipes for every code you know of.
And this is self learning, if you feed it new code, it can figure out how to optimize that,
if you give it a new target architecture and you accelerate, it can learn how to do that as well.
To give you preliminary showcase of how well that is working out, we're using polybench,
this is running our internal C++ LLVM front end, so we're running on C code instead of the
python I showed before, it shows the speed up of already vectorized code by us if we apply our
transfer tuning to it, and the transfer tuning database we use is currently very limited because
it's very fresh. Okay, same type of graph, so you can see there are some benchmarks that
profit by as much as fact the two of the transfer tuning, and locally it was not much work for the
compiler to do this because somebody else already ran the auto tuning and figured out what to do,
but there are some outliers where we make it slower by a lot. We hope to fix that soon.
Okay, so in terms of building blocks, we're adding also the remote optimization and our public
repository includes a demo server that basically responds with three set responses of what to optimize for
because the server is not live yet in that way, but the client code is already there.
And internally, as I mentioned, we are already working on doing that from C++ code directly
and the transfer tuning database, but those are not that radiate. Yeah, and we hope to gain
contributions for different accelerators if we don't get to them ourselves and new front ends,
and if you're curious and want to give it a try, these are repositories, the above one is the two
and just yesterday we published the Python project that does the transformation that I showed in the
beginning. Thank you for your time.