Software Engineering and Computer Science with the Wolfram Compiler

Tom Wickham-Jones walks through how the Wolfram compiler works: its 2,000 source files, SSA intermediate representation, type inference pipeline, and nightly runs of 220k-250k verification tests. The compiler optimizes immutable data to avoid copies and compiles Wolfram Language functions into CUDA kernels.
Chapters
What I'm going to talk about now is some ideas of software engineering and computer science based around the Wolfram compiler. So let's just dive in. So showing a little bit more. There's a Wolfram compiler. It's built from the Wolfram language. It's got about 2,000 source files. It uses data structure, not the things in the data structure libraries, but things like data structures to represent some of the internal structures of the compiler. And I will be looking at some of those if you don't mind, because I think they're interesting. And another nice feature of the compiler is it uses the compiler to compile the compiler. That's called bootstrapping the compiler. And it's a nice powerful way of innovate, making sure we can compile Wolfram language and making the compiler faster. And we actually do that. It's quite interesting how we do that because we compile it with an interpreted version of the compiler and then we reap to make something that is fully optimized, but it runs slowly but generates optimized code and then we recompile it a second time to make a fully optimized type of thing. So that's kind of interesting to me at least. So let's look at testing system testing of the. So the compiler, it uses the Wolfram verification test mechanism. I would say it's one of the best tested parts of the Wolfram system. It lends itself nicely. So our test suite is currently around just under a quarter of a million tests that run nightly on a very wide variety of platforms. And it runs incredibly cleanly. And so here's like a little test output. Don't know if you can read all this, but out of 220,000 tests there were 32 failures. And we're looking to track those down. Very often those are kind of like things that, you know, like some, sometimes, occasionally there are glitches. We try to write robust tests and things. So it's nice, very well tested. We also run a benchmarking system, so we have a set of standard operations and we run twice daily timing tests on a machine that nobody else is allowed to touch because this machine must run deterministically and it can never break down or never be upgraded or anything because we collect this data forever. And it uses, it's a nice system actually use all nice cloud use cloud notebooks and data drop and things like that. And so I'll just show you, this is. So this is an output. So it makes this cloud notebook and these are the current. Some of the tests, these are tests of the execution of code and you get a little summary there. And so these are the functions, you know, you can sort of drill down and see, see things like this. Occasionally you get these glitches. There are glitches sometime. This is timing. This is not a real time operating system. So of course there's other things happening on the computer at the same time and we want to look at this, you know, what happened there. We see these glitches and then if they start persisting then we have to go in and figure out why. And then this is the thing that's really interesting. This is showing the timings going back to, you know, sort of for, you know, just the end of. So we've got like almost two years of these timing tests. So this is quite a useful resource. Data is, data is a powerful thing that people have collected since William the Conqueror wrote his Doomsday book. He was collecting data on his new kingdom. And this is our timing test. So, so this, this is quite useful. That's the type. Runtime timings this. So we occasionally we accept, we improve some of the operations and we decide to accept the. A timing glitch and then we have compile time gain showing speed ups and. Or the operation of. So there is some potential slowdown here that we're currently investigating. And but this is the thing that's nice. This is proving that over time the compiler has got faster. So that's, that's, that's all you know, since tradition. It's, it's, it's got faster. So. Yes, yes, yeah, this is time. So smaller is good. And. Then they get into terrible trouble. So I really don't want that to happen. If we would. No, it's a Linux machine, it's going to live forever. No, no, yes, that's a real. That is an issue. We might have a new machine and run them concurrently and then have some thing. There are some issues, future proofing, questions like that. Still this is quite a good resource and we can look back and see how the system was running two years ago is actually quite useful. So. Right, let's keep going because time is precious. So that's some of the benchmarking system we want. Now I want to talk a little bit about how the compiler works. So we have this pipeline of things. I'm probably not going to get through everything. So we generate this intermediate representation it generates and then we operate on this intermediate representation and we use LLVM to generate executables. So let's have a Look at this intermediate form. So it generates these things. So it's compiled a function, but it puts it in this thing called the program module. And we can see here, this is the code that it's produced. And this is the lowest, the original form of the intermediate code. So it knows what the type of the input is, but it doesn't know anything else. Now, this has an interesting form. It's called single static assignment. And so it's got these variables and the code is broken up into these blocks of code that are called basic blocks. And this has some interesting properties. So each variable is only declared once. The code is arranged in blocks and branches or loops are implemented as jumps between blocks. And each variable can only be used if every path to that use of that variable has defined it. So this thing here, each variable only declared once. If there's any lambda calculus people, this obviously max is redolent of lambda calculus. And in fact this SSA form is kind of isomorphic. It's a strange way of writing lambda functions, but you can write lambda functions like that. And the other thing about these jumps and blocks, and this is again, if there's any graph theory experts, this is very redolent about dominators and control flow graphs and things like that. And the whole internals and things is very much thought in terms of graph theory. So, and this is again a nice. This shows I've got a program here. This is like some, you know, not. It's got like loops and branches. And so this is quite hard to see. And sometimes you want to see some different version of this. And so we have like a nice way of getting a control flow graph. And sometimes this is what you want to see. And this is showing that the Wolfram, all these other tools in the Wolfram language are super useful for this. So how do you work on these, you know, to do type inference things? We have these things called compiler passes. And the passes are designed to be very simple and mono, and you rerun them many times. And this is a philosophy that LLVM sort of prototyped. So let's have a little look at that. So we've generated this thing and now I'm going to run a type inference pass. And this is added type annotations. So we know what this, you know, we still don't know what version of not it is, but we know that it's going to return a boolean. And then we do this again. We run another pass that resolves function calls. And so this is inserted like this is expanded all the functions down to the very basic level, maybe intrinsics from LLVM or the processor, but there's still lots of sort of strange jumps and things. So it's added the actual implementation and then when we run the whole thing, it's now cleaned up the code and made various jumps and things. So that's all pretty nice. Let's keep going nice and fast. So in top level we can script this in Wolfram language and we can get. So here I'm pulling out an individual instruction. So this is the call instruction, this is an argument, and that's where the result or the target of that instruction is. And we can see these things. These are the sort of variables and this is the. Oops, my mouse is. This is the arguments. So I can work with this in interpreted Wolfram code. And it's quite useful to do that. Now a key property of these variables is looking at, say the uses. So this is every instruction that uses a variable and this one here is the instruction that defined a variable. And these simple properties are very, very useful in transforming code. And they get used to a lot and there's passes that update and regenerate the use def values. And since they get used a lot, they're compiled. They're one of the first things that we compiled. And that was a nice use of compiling the compiler. So let's just have a quick look at that. Probably don't have lots of time. I'm going to generate. This is now I'm adding a compiled component. That's the compiled, the compilation of the compiler. And once I've done that, I can start to compile functions. So here's something that I've compiled. It's still running. Hope this won't take too long. And this is taking a call instruction as an argument. And then I'm getting the target and the uses of that target. So I can execute that compiled function and it returns a list of the return instructions. But in terms of types, it's called a thing. Like it's called the list expression, which is. It's really implemented as an expression, but it's given a type and a contents. And so that means it fits nicely with compiled code, but it's also very interoperable with the evaluated code. So that's kind of a nice way of doing this. Let's have a look at a few. Probably running short on time, but I'm going to run down to the wire. Yeah, so I'm looking now at some more computer sciencey things. Some of the optimizations that we do. So I'm going to talk now about. This is like nested function functions, so lambda functions. And this is the way that you pass function pointers around in your code. Now, if you actually pass them as function pointers, that's actually slow. I don't like that because it's slow, because you just get an address and it says, at this address there's a function, call it. But I want to know what function it is because then I can inline it. And these lambda functions are astonishingly useful. So they. Sometimes people think, oh, you know, I don't really use them very much, but they are useful. They're such powerful ways of encapsulating and abstracting. Now, so here I've got something, here I've got a function and then I'm calling it. And so I'm going to run this at a sort of low level of optimization. And you can see here it's calling the function through a lot like a variable. So it's not really knowing much about this. Now if I increase the optimization, then now what it's doing, it's actually figured out I'm calling a particular function. Once it does that, it can inline it. And then, in fact, when you pass this to llvm, it's like thrown everything away. It's like, oh, we'll evaporate this, we'll inline everything. There's just one thing left. And it's astonishing, when you look at this underlying LLVM code, sometimes what you see is massive and sometimes what you see is tiny. It's. It's super interesting. And this is when you want to know a little bit about how much of the internals it's using. This is the way that. One way that we would use to look at this. So now I'm going to show very short. Oh, no, good for time. So, right, so I'm not going to sort of evaluate this, but I'm talking. This is another sort of researchy topic that. A nice Internet. Alfie Richards was working with me over the summer and this is one of the things that he was looking at, was trying to convert immutable operations to mutable to avoid data copying. So here I've got a function passing it in a packed array. I'm putting that packed array into computing the sine of the packed array and then the cosine of that. So when we look at this code that we see here, we see that the output of the sine is being generated new, but it's only used in the cosine and we can prove that from the uses use def chain. So that means we could rewrite that we could have when we wrote this immutable operations pass. And then when we run that, it's taken the output of the sine. But now it's calling a different cosine of a packed array, it's calling one where it's given the output and the output that it's given is. So it's is this argument here. So it's giving an output array and an input array. So it's basically overwriting itself and that's saving data, copying. So that was quite interesting. Project is not finished yet. Another interesting project was integrating with the CUDA. And the CUDA toolkit. This is part of a rebranding of CUDA Link that we're going along. And here's the basic workflow of GPUs. So you start load data on the GPU, operate on the data, retrieve the result and then free the data on the gpu. Copying the data to the GPU is very slow, but working there is fast. And so these are things we added. We added these CUDA vector different sorts of things that fit in nicely into memory management. Now what we want to do is to compile Wolfram language into cuda. So here we've got a function and it's got some rather strange CUDA intrinsics that are part of it. And so we have this thing compiled CUDA kernel, and then we see here the CUDA code that's produced here. And then we can execute this and write some other compiled code, and then we can execute this function, downloads the data to the gpu, execute it there, and such like. But even better is to integrate all this, where we embed the CUDA function inside the other function. And so we can do that here. So this is a bit sort of embedded here. So in here we have a function that's been set to be a CUDA kernel. And then we're going to run that with this run CUDA kernel object. And then we achieve this by having a pass that looks for these CUDA kernels. And this is doing a 2 system. So part of the program, it's compiling for a different architecture and then it's embedding that back into the original program. And we only just got this working, but that's pretty nice. This is what the CUDA compiler itself does, things like that. So when we do that, we had some nice things where at least for vector vector multiplication, we were able to achieve similar performance to Nvidia's blast library matrix multiplication is much more complicated because it has various blocking and things, but still this was a useful sort of proof of the whole thing. So it's implemented with a pass and the pass runs just before type inferencing and it compiles CUDA kernel functions to PTX and inserts this. So the open architecture of the Wolfram compiler makes it easy to write these passes to modified code and learn about new techniques. So I will just summarize. Probably don't have any time for questions, but I looked at construction of the Wolfram compiler and testing and benchmarking tools. We looked at some of the pipeline structure of intermediate code and how we can compile that and then looked at some applications of this such as CUDA integration and immutable operations. So thank you very much. I think we probably do not have time for questions because there's another talk about to start, but I'll be around to answer questions later. So thank you very much.
More from Guy

Cheryl Strauss Einhorn | The Human Edge: Smarter Decisions in the Age of AI | Talks at Google

Rachel Glennerster | Innovation and the Center for Global Development | Talks at Google

Rebecca Winthrop | Rethinking the Purpose of Education in the Age of AI | Talks at Google

Dr. Robert Wachter | A Giant Leap: How AI Is Transforming Healthcare... | Talks at Google






