Scale By The Bay 2020: Ryan Orendorff & Daniel Hensley, Programming machine learning algorithms...
Yeah, as Nicholas mentioned, we're going to be talking about how to program machine learning into hardware, specifically FPGAs, using Haskell and Rust. Um, so let me introduce the two of us. Um, my name is Ryan. I am currently a research scientist at Facebook Reality Labs Research, working on brain-computer interface research. Um, I've been working with actually Daniel for the last 10 years on different medical imaging technologies. Uh, and on the side, I really like to poke around with dependent types, uh, Agda, just programming language theory in general. I do have to mention that this talk was completely done on my own time. It's not endorsed or like sponsored by Facebook in any form
Um, and then my co-presenter, Daniel, works at a company called Magnetic Insight, which is commercializing a technology that is about, um, called magnetic particle imaging, that seeks to be something kind of new in the medical imaging field, not just in addition to, say, PET, SPECT, MRI, and those kinds of things. So, he has a background in, uh, signal processing, image reconstruction, applied math, and electromagnetic physics. Um, he's worked a lot on embedded hardware, uh, hence the Rust components. And, uh, he's worked a lot with FPGAs and Rust to great effect. And, yeah, he's just a fantastic engineer, so I'm excited to give the presentation with him. So, what we'll go through in this talk is we'll go through basically a simple neural network. We're only going to be doing forward propagation, um, in order to figure out, kind of, given some inputs, what would a label be. And we're going to talk about how we go from that neural network into something that can be programmed onto hardware
So, basically converted from code into, um, logic gates, flip-flops, all those fun jazz, automatically, in a way that's easy to handle. And then we'll talk about how to access the FPGA resource safely using Rust. So, the general overview, the big overview, is that we have what's called an Intel Cyclone 5 chip. And on this chip, there's two different pieces, an ARM processor and some FPGA components. So, you can program the FPGA components to kind of have any interconnect between the logic gates that are on there from the ARM system. And so, what we'll do is that we'll program, using a programming language called Clash, which is derived from Haskell, a neural network and then access it from Rust. Reasons why you might want to do this include, uh, you can get really low-power FPGAs. Um, you can have really high throughput on the order of what you would maybe do on a GPU, or even better if you have a very specific, specific neural network that you're implementing
Um, and you can also get deterministic timing. Since there is no operating system, everything kind of just runs against the hardware clock. So, let's go over the basics of, kind of, a neural network. So, everyone's probably familiar with this diagram at this point in time, but I do want to go through a little bit because oftentimes it's a little hard to see what's going on. You have a series of circles and lines connecting them. But what does that mean? How do we get data from the left side of this figure throughout to the right side of this figure? Well, what we do is that to transition from one of those layers to the next, we basically just solve this simple equation, or we just evaluate the simple equation. What it is is to say that from one layer X to the next layer Y, which are both vectors, each vector representing a set of nodes, we apply this equation Y equals some G activation function on top of MX + B. So, graphically, what does this look like? When we try to, um, when we look at this, we'll say that we have X1 and X2, say input is our two-element vector
The M is really the connections. It's the lines between the nodes. And what it does is that it's a weighting of how much you want to add these two nodes together. Now, there's this extra term in here, B, which basically allows you to also add a bias offset to this addition. Um, sometimes this is represented as one of the nodes in the input layer itself, but here we just have it kind of off to the side. And then finally, to get our output, our first output of Y, the first element of Y, we apply the activation function, which we've called G here, in order to get our output. Now, in order to evaluate the rest of the network, in order to get the other nodes in Y, we simply repeat this process uh, with different connections inside M. So, now if we see here, what we have is the same setup, except for with more nodes and more of the M green lines, which indicate the weightings of the different input values
So, in order to complete this, in order to evaluate this equation, we're going to need two components. We're going to need vector addition, and we're also going to need matrix-vector multiply. Um, so let's first talk about vector addition. Um, so if you've, say, added two lists before in Haskell, you know, a convenient way to do this would be with zipWith. And you say zipWith just the addition function. Um, but one thing to note here that's different between this zip and, say, a standard zip over a standard list is that there's a type parameter N, and specifically it has this constraint that it's a KnownNat. So, what this means is that we're actually able to encode natural numbers at the type level and use that as part of our type. This is part of the dependent typing of Clash, and actually this is available just in standard Haskell as well
So, what it means, what this function means, is that it takes in two vectors, an X vector and a Y vector, say, and they have to be of size N. They can't be mismatched in size in any form, and it returns to you also something that is the same size. So, that's kind of neat. Besides being from a programmatic correct standpoint, which is helpful to program against, the FPGA will also need to know exactly how much resources and what connections to make later on. So, this number is actually required in order to program down to the FPGA. Now, if we look at the dot product, which is the start of how we get to matrix-vector multiply, it's a similar operation. What we do is that we zip two, uh, vectors together of the same length with the multiply, but instead of just leaving it at that, what we do is we fold over that resulting vector and reduce it down to a single value. So, here you can see the output is actually now just a type A
And we still haven't specified what this type A is, just that it has to be some type of number. So, we can actually get Clash to tell us to actually print out what the circuit will look like when we actually try to put it on the FPGA. So, here's the diagram that it produces. I've added some colors. But basically, what you can see is that it takes in some X and Y vectors, and it first finds out where the elements are in those vectors. That's what everything in that second column is. Uh, cuz it's otherwise kind of smashed together in terms of bits. It then adds the correct components together in the zipWith layer, that's the purple parts, and then folds over the rest of it, adding each component until you get to the output stage
So, this is pretty neat. You can physically see what is actually going to be, what's going to be connected when actually gets put down to the FPGA. Now, in order to get to matrix-vector multiply, what we need to know is basically what a matrix is. So, a matrix is a vector of vectors. So, in this case, uh, it's row major form, so that means that we have M rows of N columns. And in order to make matrix-vector multiply, what we do is that we take each row in the matrix and we do the dot product of it with our incoming vector. So, in this case, we can use the same higher-order function that you're probably used to, map, fmap, um, in order to actually, uh, do this operation. Now, let's encode everything that we've just talked about in a, into a layer transition type
So, what does it mean to go from one layer to the next? Well, we've described this type called layer transition that takes, uh, goes from an input size I to an output size O of type A, and its inputs are the matrix of connections, the bias vector, and the activation function that we want to run. Note that the bias vector has to match the output size. And similarly, the matrix that we use, that takes that input and converts it into the like you have to have the correct number of rows and columns in order to construct this type. Now, how do we generate a network out of this? So, what we're going to do is we're going to use, uh, a GADT encoding of the network. What we're going to say is a network has a few components. Described in the type is an input size, a list where each size, each number inside the list, each natural number, is the size of the hidden layer inside that network, and the output number of nodes. So, in order to do this, we're going to basically make what's equivalently, equivalently to a non-empty list. So, the first case is we have just an output layer
So, for an output layer, we have the layer transition that goes from some size I to some size O, and we make a network that has no hidden layers in it whatsoever. That's what the tick, uh, brackets means. Now, in order to add a layer onto the front of this network, what we do is that we take in a new layer, layer transition, and then we take in an existing network, and we basically prepend it to the network that we have. And I want to call out here that, again, the shapes have to match. So, in order to add a layer to the front of this network, I have to, the input lay transition has an output of size H. That output size H has to match the input size H to the network when we construct this type. Now, note the final output also includes the new input to the network is the input to the layer that we're adding, and that size of the intermediate layer gets added to the front of our list of hidden uh, layers inside the network. So now we have everything we need
We've encoded layer transitions. We've encoded a network. Let's see how we can run this thing. So in order to run a specific layer, we have our layer transition and we're going to take some input number of nodes I and return an output number of nodes O. So all we're really going to do is deconstruct the type that we have that we've been passing around and apply our equation Y = G MX + B, right? And that's exactly what we have here. We see we map the activation function over each node and then we have the matrix vector multiply of M and X and then we just add our bias values. So it's as simple as that. To run a network, what we're essentially going to do is what you would might do to deconstruct a list is we're going to fold over it
So there's two cases here. In order to fold down a network, we're going to either have an output layer, in which case we know how to handle that, right? We're just going to run the layer itself. There's nothing else that we can really do with that object. Otherwise, we have a network with a layer on the front and the rest of the network. So the way we're going to proceed is we're going to pull off that first layer, run our input vector through it, get a new output vector and then run that through the nest of the network recursively. Note that in Clash the support for GADTs and recursions on GADTs is actually somewhat new. So this is a cool new feature that you can use. So how would you compose say a four-layer network? So say we have these weights
Well, you would just call the data constructors that we talked about previously. In order to generate a network, we would simply say layer one attached to layer two attached to layer three attached to layer four. Now note if I try to mix up any of the layers here, say I tried to replace layer two with layer three, I would end up getting a compile time error. And the reason why is because the number two and the number three, when it tries to unify the types, will not match. It will give me an output so that says basically type two does not equal type three, which is correct. Um so the type level natural numbers, besides being required in order to generate FPGA hardware, also make sure that we generate networks that are valid. You cannot generate an invalid network using this construction. So how do we finally get this down into hardware? We talked about we lived in a polymorphic land, which is great, but unfortunately on FPGA hardware, it's very hard to be polymorphic with very little space on that FPGA
So what we do is that we're going to now have to choose a size that we're going to use for our networks and we're also going to have to choose the type. So here we've chosen SFix. It's essentially a floating point number. FPGAs don't usually have access to floating point, sometimes, but rarely. Um and we'll basically take a size two input, an X and Y coordinate, and return a classification. So you can see all we need to do this is we just need to say we're going to run the network that we've generated so far in the code. So as I had mentioned, we need to specify a specific number and the reason why is because when it goes down to hardware, it needs to know what to do with those bits. We otherwise are just sending random bits, so it needs to know what to do with them
And we had to specify the size. So experienced FPGA developers might notice that there are some problems with the with this construction in some forms in the fact that you could potentially generate a very large network and you wouldn't be able to run it. The problem is is the network would have too many components, too many flip-flops and logic gates in order to run within one clock cycle. But Um so we can show like basically we have a our example with four we had before. Um so this is like a small example. This is the dot product. But imagine that if we scaled this up to a million elements, it wouldn't quite work. So in order to do this, you can break your computation into pieces and when you do so, you can use these things called registers
So registers basically store the intermediate values of your computation as you're running. Now in this case, the registers come down to this mux block and this addif block that occur here and here. And what they'll do is that they'll basically run what will happen in this network is you'll run the zip part, which is the part on the left, the zip with star or multiply, save that for a clock cycle or save that until the next clock cycle occurs and then run the fold part. So what's important about this example is in order to work with FPGAs, you'll still need to know quite a bit about FPGA programming and hardware and stuff like that, but Haskell actually makes the definition very easy to read and it looks very similar to what uh you would have before and remains fully compositional and polymorphic until you instantiate it onto the hardware. So um since Clash is mostly Haskell, you get a ton of benefits. You get a strong dependently typed type system. So um and it has some great features for specifically solving natural number constraints. Um complicated FPGA state can be modeled by something called a Mealy machine, which is something equivalent to state T
So you can do more complicated state machines inside Clash. Uh all your standard abstractions can be used in order to generate this code. The prior example used an applicative instance. Um but you can also have monad um type constraints or type classes on top of this in order to organize your code in a coherent manner. The base functions that you have are pure Haskell, which means that they can just be typed which means that they can just be tested like any other Haskell function. Say with QuickCheck. And additionally, Clash contains some components that allows you to simulate on a clock level exactly what happens to the hardware such that you can make sure that your timing constraints are met. So the last thing I want to say before handing off to Daniel is that on the FPGA, we have that Intel Cyclone V
We put a five-layer network and what it does is that it basically takes an input coordinate X and Y and returns positive one if you're in the first or third quadrant and negative one if you're in the second and fourth quadrant. So with that, I'd like Daniel to take it away. So we've just got this awesome network that we've built and are running on this FPGA, but how do we inter- interface with it? So a very common paradigm with FPGAs or hardware and you know, a mixture, you're going to have a host CPU that's kind of the brains. It's going to delegate tasks down to the FPGA. Um and so, you know, in this type of environment, how are we going to do that interaction, right? So what we're going to talk about now is building this session API in Rust. Um and this can live on top of a very low-level IO. Might be an FPGA driver, CDLL, could be more Rust. In the case of the application today, it's actually a memory-mapped file IO
All right. So the key concepts that we're going to come up next is we're going to shoot to use the Rust type system to do a few things. We really want to encode and enforce hardware invariants. We want to push as much as possible to compile time checks. We're going to use some, you know, cool type stuff, but we have to maintain ergonomics for the developers. So in the end, we want something that looks a bit like this. Very simple API where, you know, you get your session, then you just have these simple read and write methods to write down to the FPGA and read data back, right? The same time, this very simple code is going to have a lot of compile time guarantees behind it. So we will build this session API and here I've just got a figure that kind of has just some of the components of the Rust type system that I found really great and helpful in general, but especially for mapping onto hardware and controlling hardware
So the green ones are what we're going to talk about today pretty explicitly. So type states, affine types in Rust, traits and generics. Um the other ones here, error handling via types and the borrow checker with lifetime, but those are also great for hardware um abstraction, but we just won't have time to get into that today. All right. So the major design components of our session API is number one, we're going to have types and the ability to express our application specific resources. So these are registers, FIFOs, structures that are on the FPGA that we want to write data down into and we want to read data back from. And then we're going to have our session, right? So this is a type that's going to from the host application's point of view, it is the FPGA. It's going to wrap the FPGA and all of our interaction with it
Then we're going to link these two together with a type system and again in a way that's ergonomic. So first we'll talk about generically modeling FPGA resources and we're going to do that with traits and type states. So traits are one of the real anchors of the Rust type system. They allow you to define shared behavior and constraints for a set of types and they're similar to interfaces and type classes in your favorite language. And so this data trait is the trait we have to represent, you know, any FPGA data type in our application. And so we just have to implement four methods for this trait and basically there're two sets of methods. One that says, you know, if I have a sequence of bytes, how do I convert that into my type? And the opposite, how do I take my type and convert it into a vector of bytes? Um and if you look at the GitHub repository that's a companion with this um piece, then you'll see we've implemented it for all of the basic primitives of different bit width integers, floating point, etc. Um but if you're working on your own application, you have your own custom data type, all you have to do is implement data and you can plug into the session API
And you'll see that we actually did that for a tuple pair of fixed point numbers. All right. So another thing that's really important is to control read/write access to these FPGA resources. And so we're going to do that with a simple type state. And so I've got a slide here that just talk a little bit about type states, which I think are really great way to encode invariants in general, but especially working with hardware systems. Um and this is a, you know, a type compile time level concept, so there's no runtime cost, which is great. And again, if if, you know, whatever you've encoded with these type states, if it compiles, you know, they're guaranteed. Um there's a couple of patterns here that, you know, we've used a lot
Um on the left, this is kind of this idea where you have maybe a common runtime representation. So we're going to see this with our resource struct. Um and then we fan out into a number of different type states that are kind of orthogonal. And those separate different operations that are actually available. And then the other one is kind of a type-level finite state machine. Um this is helpful. Actually, the fair, you know, the common builder pattern in Rust is a simple example of this. But here you kind of map out in the type system transformations that are allowed through these different type states
And again, each type state has its own set of operations, and you control that. You also control, you know, how you enter, you know, the state machine, and then how you get to another node. Um interestingly with Rust, again, because of the affine type system, you have this nice thing where you have a variable, if you transform it into the next type state, it's consumed in the process, uh which helps. Okay, so this is uh the simple type state we'll use in in this case, but this also shows the mechanics of how you do this in Rust. Uh so we have an empty type, uh empty marker trait to start this type state pattern. And then for each specific type state, we just implement an zero-sized type, in this case, uh an empty enum. And so we're going to have read-only and read-write as our type states. So putting this together, this is the generic struct in our application that represents any, you know, FPGA resource
And you'll see the only runtime fields that we have are, you know, a useful name, um and then the byte offset. And then, you know, but we have the type state to DNI are actually what determine the available operations that are associated with the FPGA and that you do through the session. Um and then this phantom type, by the way, this is just how you tell the Rust compiler that these are type states. So what does this give us? So in application code, in a type sense, we cannot possibly send down the wrong bytes to the FPGA or interpret incoming bytes incorrectly. Uh compare that with maybe some void pointer C DLL type uh scenarios. And we can't mutate uh a read-only FPGA resource. So this is uh code that you'll see, you know, you can go see in the example files we have. We get our session, we define the input to the FPGA to the network as a tuple of these fixed-point values, that's what this IS I7F25 is
Our output classification is just a single um fixed-point value, and it's read-only. And then, if you see these codes, the code below, the three lines just will fail at compile time. So in the first case, we tried to write a floating-point pair when we said it's a fixed-point pair to the input. Uh then we try to read an integer from uh the classified the classifier register, that's not going to happen. And then finally, we try to write to what we've defined as a read-only resource, and that's not going to happen. Okay, so now we're going to move on to the session type. Um and so this is again the opaque type that represents the FPGA and our interaction with it. And one of the things we can do is encode the sort of singular nature of the hardware and the importance of maintaining appropriate state with the help of the type system
So we don't want to let ourselves or our devs accidentally spun up duplicate sessions and all this kind of stuff. There's only one FPGA, uh one piece of hardware, and racing on that is not something we want to do. Um so we can use Rust's version of the singleton pattern for this, which you can kind of see here. Um I guess the main point I'll just make is from the live, you know, using this library, the only public function is this take FPGA session function, um and then the application code, you can call it once you get the session, but if you call it again later, you're going to get a runtime panic. All right, and so another thing is again, FPGA and and other hardware uh are very stateful entities, and it it's very important to maintain proper state at different times in the real world, right? Um and so we have a lot of tools in Rust to help us do that. So first, you know, and this is common in other languages as well, we can have a constructor that's the only way to get a session. We just make sure we have our uh hardware initialization stuff in there. But what's maybe a little more unique to Rust is on the finalization side of things
Um and so what we'll show is that you have to implement what's called the drop trait in Rust. Um and this basically makes us figure out, okay, put here where how you need to finalize the hardware. And then once you do, it's impossible to forget to finalize that, no matter what happens in your program. So we're going to see the full session trait here in a second, but you'll notice here that um the trait, any type that you want to implement the session trait for, you have to also implement drop, as I mentioned. So we're going to have this memory map session. And so you can see here that we have to implement drop. And again, this is where we put this critical hardware invariants for the FPGA. Um and the what what happens here is this this method call is run right before the destructor um in Rust
And you know, Rust has this really nice aspect where it you don't have to ever call destructors manually. And so this basically, once you implement this trait, it'll just sort of all work out in terms of when this will be called. It's great. So basically, what this looks like is, let's say we have um some code here. We get the session, we do a bunch of cool stuff with our FPGA, somewhere, beginning, middle, end, it doesn't matter. We do something silly, we call a risky function, and we also assert that the happy path happened. So what's going to happen? You know, what's going to happen to our FPGA? Well, it doesn't matter in the sense of whether we panic or not, our session will be dropped appropriately and returned to uh the right state um because of the implementation of drop. Okay, so now we'll bring everything together on the session API side
Um and so we have a couple helper traits that we define, one for a readable FPGA resource and one for a writable FPGA resource. Both of these have an associated type value, which implements the data the data uh uh trait that we talked about earlier. And then you can see down here, we can just implement readable for, say, the resource read-only, resource read-write, but we only implement writable for the read-write resource. And then finally, bringing it all together to the actual session trait, you can see here, in the end, it's pretty simple. You as you saw, we depend on the drop trait, um but we just have two generic methods, read and write. And they take read, for example, takes in a readable resource, and in the happy path, it will return the associated value type of that resource. Um and then on the right side, it's similar, we take in a resource that's writable, and we take in a value that must be the associated value of that writable uh resource's type. Um and then it'll write it down to the FPGA
And so what we have here, again, is this very simple and ergonomic API with just read and write methods, but we've wired up a lot of traits, constraints, and type states so that there's a lot being enforced at compile time. And by the way, what Rust is going to do in this case is the compiler will monomorphize concrete read and write functions uh based on what resources are defined and used in a given application. And so here I actually just show, and again, you can check out our repo, um the implementation of the session trait for a specific memory map session. Um it's pretty simple. On the read side, we just basically get the byte offsets into the map, um and then we go grab that slice of bytes, and we had to implement this earlier, remember, on the data trait, so then we just call the appropriate, in this case, our board is a little Indian um system, so we just call the appropriate method, and good to go. On the right side, it's similar, we get the right offset into the map, we take the value that was provided in, um you know, uh convert it to the little Indian bytes, and then write it into the map. Okay, and so the final thing that we'll do here is show, you know, talk about this quadra classifier that we've got working on the on our dev board, and you can see in the repository. So on the right here, this is actual data points colored um based on the classification from our network running on the FPGA, so it works
Um and you again, you can check this out, and you've already seen most of this, um but how this all works, you get we get our session, we define explicitly the input point uh resource, the output classifier resource. We want to test, say, quadrant one and quadrant two, make sure those are correct. Just looks like this. Again, we write down a quadrant one point to the FPGA, and then read back what it classified that as. In this case, it should be 1.0. We can also test the quadrant two, which should be minus one. And again, as you've seen, it works. And here's just some terminal output uh from the program that we have in the repository
Um so in in summary, um FPGAs are really neat. They allow you to map your hardware on, you know, onto the algorithm. It's just it's a, you know, really nice thing. Really fun to do that. Um and then Clash provides a sane way to write FPGA programs, and we've talked about how Rust can provide a sane way to manage interactions with the FPGA. And then we demonstrated this in a simple point classifier application. Um I think a big takeaway for us is type systems and architectures that push guarantees to compile time are just really great for modeling and interacting with hardware. And we also want to say a big thanks and shout-out to our friend colleague and FPGA expert Wayne Ketterling for a lot of his help on this project
And so again, you can find um these slides and a lot of other stuff at this repository, and yeah, we'll conclude the talk and go to the questions.