Scale By The Bay 2019: Shoumik Palkar, Weld: An Optimizing Runtime for High Performance Data
[Music] yeah thanks for the intro no I absolutely what he said I my name is Cho Makai a fifth-year PhD certain now at Stanford and I'll talk about well there's an optimizing runtime for high performance data analytics and this is work that I've done with my advisor Mateus Acharya and many collaborators so the motivation for the Weld project is that modern data analytics applications combine many disjoint processing libraries and functions so as an example in the Python ecosystem a typical data science pipeline may load some data using pandas normalize it in some way using numpy and then and then maybe train a model or something using scikit-learn this is nice because on one hand you get great results leveraging the work of thousands of authors who are implementing the state-of-the-art algorithms in each domain but on the other hand you don't get any optimization across these individually written and optimized functions so how bad is it that we don't get end-to-end optimization across these independently optimized functions well it turns out that the growing gap between memory access speeds and processing speeds on modern hardware actually makes this rigid function call interface to be as libraries worse and worse over time so as an example of this here on the slide is a typical pandas program that's performed some really simple data cleaning so this program is parsing in CSV then it's dropping some missing values and then it's computing a mean so each operation is actually highly optimized so under the hood it's implemented in C and it's using things like Tim D vectorization sometimes but because each function executes independently over in memory values the overall execution is repeatedly scanning these values from main memory and rather than you know collectively operating on values that reside inside the CPU caches and because of these repeated stands over the data even when each library is already optimized much of the execution time can actually just go into moving data between main memory and the CPU so together with other missed opportunities like eliminating common sub-expressions for example we've observed up to 30x slowdowns when combining popular libraries compared to an optimized c implementation of the same application so there's a lot of performance left on the table with what you get using these libraries versus kind of just rolling your own solution I'll bake with a lot more engineering work this is kind of the state of the art of data science today a data scientist will pip install or acquire libraries that they need to perform some analysis then they'll observe some performance issues in pipelines composed of these fast data science tools just because they're composing them and then they'll hire engineers or file support tickets to optimize the pipeline or to do things like leverage new hardware such as GPUs so the vision for this project weld is to basically enable the spare metal performance out of the box even when you're kind of composing these arbitrarily already optimized libraries so weld approach for achieving this is to provide a common runtime for data analytics libraries so with weld libraries from different domains such as sequel and machine learning and graph processing can each submit computations to this common parallel runtime the runtime can then analyze entire applications end to end and then JIT optimized code for each application to target different kinds of parallel Hardware so that can be you know multi-core CPUs GPUs or even future emerging accelerators so under the hood weld has a number of kind of key components the first component is this runtime api that developers of these libraries use to actually submit computations into the runtime the second main component is this ir intermediate representation this is like a small kind of parallel programming language that Weld exposes that libraries are expected to kind of express their core algorithms and then the third component is an optimizer that basically takes this ir and then it takes the kind of fragments of code that these libraries them in and makes sure that they run efficiently together so it's an optimizer that takes the IR and transforms it into kind of a more efficient version of itself and then there's also these backends that will actually do code generation by mapping the wild ir code into machine code for these various hardware packets and taking this kind of end-to-end approach can actually lead to some pretty significant speed ups so this is a data set that we took from Kaggle it's pretty simple it's just filtering some data computing a linear regression model and then it's aggregating some indices at the end and it uses numpy and pandas so both of these are already pretty fast they're both backed by C so each function is already you know fair we optimized but we can kind of compare their performance when just using these libraries versus using weld so if we just use numpy on pandas on a single thread you know it finishes in eighty seconds a pretty small workload so it's not too bad but allowing well to do its optimizations just within each library so just optimizing the numpy functions together and optimizing the pandas functions together already leads to an ABEC speed up because you eliminate things like allocating unnecessary intermediate values and also do things like loop fusion and things like vectorization of cross function calls if you allow well to actually optimize across libraries as well you get a further three X beta so now we're almost thirty times faster than this implemented implementation that we started with which is already leveraging these already optimized operators and because of the way weld is designed anything you submit to it cannot kind of trivially be parallelized so anything that you expressed in the IR well doe kind of know how to parallelize because of constraints that exist within the IR so if we run this on twelve thread CPU we get a further speed-up so it happens to be 180 X faster than the original execution in this case so that's a huge jump in productivity when you have to wait you know half a second versus eighty seconds to wait for these kinds of tasks to finish so when the rest of this talk I'll first discuss some details about the runtime API then I'll dive into the world I are a little bit more and then finally I'll talk about world's optimizer and then I'll talk about some performance results how should we skip the demo because I can't see the slides up there so a little like to do okay so the runtime API uses lazy evaluation to collect work from across libraries um so this is pretty similar to API some other popular system such as spark or desk data frames which also use lazy API X so as an example let's say we have a user application with three well enabled functions right so wait wait born f1 Lib to that map and lead 3 F 3 when the application calls these functions they don't actually compute anything over their inputs instead each function will just use its inputs to generate a little snippet of welds IR to represent a computation and then it's amidst these snippets and any data dependency is to weld using this runtime API and then weld will kind of internally track dependencies among individual snippets between that are submitted by these different functions when you actually want to evaluate something this might be when you want to print a result or write something to disk or maybe send it over the network well it will actually use this dependency graph to stitch these IR fragments into a single program and it'll lift those data dependencies up and then it'll basically run its optimizer over the ir program and then compile it to machine code and that's what the data in your application will actually interact with so just to illustrate this a bit further here's a super simple program it's basically computing like a euclidean distance so it computes some squares and then a sum but with kind of the everyday pattern that we do this you have some data sitting around and each function will call will just kind of eagerly produce a new value so each call is reading and writing from memory with well the runtime API basically has this concept of a weld object which is basically this thing that encapsulates this dag of dependencies so every time they'll call a function you're just registering it with this object and that object is basically passed around to other weld enabled functions so every time you call something we register a computation and at the end when we actually want to see a result we compile that into an optimized program and we evaluate the optimized program just once so the next kind of natural question is how do we actually express these computations so we can actually do some optimizations over them and this is where the wild IR comes in so welds IR is designed to meet kind of three high-level goals the first goal is just generality so we want to be able to support a diverse set of workloads because usually these data science type pipelines are composing functions from different domains so that might be sequel it might be some machine learning it might be some graph processing for example the second goal is the ability to basically Express powerful optimizations especially when you're kind of stitching these redundant fragments for IR together so we want to be able to do things like loop fusion over the same data easily but also support more hardware specific optimizations just vectorization or loop tiling as well I mean a lot of design of vir kind of kind of goes into making these kinds of optimizations easier and finally we want parallelism to be an explicit concept within the IR I mean this is just because modern hardware is just inherently parallel whether it's a multi-threaded CPU or a GPU which has you know thousands of Hardware threads so and I are designed for performance should kind of keep that in mind and design around it so internally welds IO is pretty small it's it's kind of like a functional programming language and it has like the usual constructs which is branching and arithmetic but it's kind of built around two-million constructs the first construct is pretty easy it's a parallel loop so this is just a looping operator that will iterate over your data set in parallel but the second construct is a little bit more interesting and it's this thing called a builder which is a declarative object used to produce results and when I say declarative object that just means that a builder will say something like I want to append items to a list or I want to compute a sum but it won't actually provide a implementation for how to do that within dir so whereas in C you would have to say you know I want to spawn these many threads and manage state in this particular way the Builder will just say this is the result that I want to produce but in the IR it won't actually say how to do that and the nice thing about this is that it allows for different implementations tailored to different kinds of hardware so builders also have this additional constraint that makes parallelism easier which says that you're only allowed to read values out of them after all rights from them are complete and because we're kind of working in this functional world this basically allows us to enable mutable state when implementing these builders and also simplify some kinds of analyses in the compiler as well so if formally these builders are known as linear types if you know about those but this this constraint basically exists to make a compiler analysis easier and just using these two constructs we can capture a pretty wide variety of workloads so we can do things like relational algebra we can capture basically all the functional API subsystems like spark we can also do things like linear and tensor algebra and most importantly we can also perform optimizations when you have a composition of all of these different things so just to show an example of what the IR kind of looks like we can implement some simple functional operators using it so here's a map function so this just takes data set and a function f and it applies F to each element in data so the way this works in Weld is that you create this new builder so there's a builder type in Weld called an appender so when you do writes to it it just depends items to a list and when you read back from it it gives you back that final list and then we use this parallel for loop to iterate over data in parallel and in the body of the for loop we use this merge operator which is like doing a write into the builder and we write the value f of X into into the Builder and then when we're done with that we just call a result which will give us back the final value so one example of how this might be implemented under the hood is that merge will kind of spawn per thread vectors and kind of tack on values onto those vectors and when you call result and you know result won't be called while you're kind of doing writes you can just take those vectors and concatenate them or sort them in some weight before concatenate them so it's one reason why this rights before reads constraint makes makes things a little more efficient if you want to implement reduce it's pretty similar we just replaced a pender builder with this other one called a murder which basically applies an associative function to anything that you merge into it so other than that this this function looks pretty similar to the one above okay so one reason that all these kind of design decisions remained was because we actually want to be able to do meaningful optimizations over the IR and that's where this automatic optimizing comes in so the optimizer is kind of built with two goals in mind I mean these goals kind of come from world setting of trying to optimize these ad hoc data science libraries so the first goal is just to remove redundancy caused by composing independently written libraries and functions so well there's in a language that you'd expect like a programmer to write it's basically snippets of code that are kind of glued together at the end so things like data movement or redundant traverses over the data we're going to be Supercop in and we want the optimizer to kind of eliminate these things automatically and the second kind of big goal is being able to generate fast code but without having to rely on external statistics so a lot of the optimizations I'll talk about are actually similar to optimizations that exist and say a sequel database but sequel databases have these nice catalogs and histograms that you can rely on to basically guide what kinds of optimizations you want to apply based on what your data looks like but we're trying to optimize things like pandas and numpy and kind of these ad hoc libraries that people use so we don't have access to these kinds of statistics so we actually have an adaptive way of kind of collecting statistics at runtime and then kind of deciding which which version of the code to run so over all the optimizer design looks something like this we have these runtime ir fragments that are stitched together into a single combined ir program we then apply a standard rule-based optimizer to this ir so these rules are kind of tailored toward removing redundancy and kind of applying things that are you know kind of always good to applies if you have like expressions that appear more than once it'll and eliminate those and then it'll also do things that standard compilers to use just holding constants for example then we have this adaptive optimizer and it's adaptive because it'll actually insert code into your ir program that will kind of choose whether you want to run one version of the program or the other depending on what your data at runtime looks like so it's adaptive because it adapts to the kind of data that you're streaming through at runtime then you can kind of go back and run your rule-based optimizer again to fix point or you can generate code so we use LLVM to do CPU code generation to translate ir code into into machine code okay so the first category of optimizations are these redundancy removing ones so these are rule-based optimizations that basically just pattern match on patterns in the ir and replace them with kind of more efficient sub programs so here's a really simple kind of program so you can think of each line in this program is being generated by a separate function so the map call can be like an x times X and numpy the reduce can be like a sum and numpy and then again there's like another map there and right off the bat we can see a number of problems with this so we're unnecessarily kind of materializing temp assuming that our only outputs that are of interest are some and map there's two traversals of data here which isn't great because we're not exploiting things like cash locality you know and there's other things that just might be missing in the underlying implementations so there might not be vectorization so we might not be using cindy vectors for example in the cpu and there's also like things like you know some other things like size inference so we don't want to keep calling malloc for example to grow the vector we can kind of infer the size of that but upfront so after weld applies its optimizations this is what the program will look like and it looks pretty different and kind of the most obvious changes that we've replaced these map and reduce macros with their kind of loop and build a representation but there's some other nice things that adopt misers done as well first you can see here that we've used loop fusion to actually pipeline to loops so this temp vector that's in the left hand side program actually isn't materialized anywhere we just make this merger which remember creates aggregations and then we merge in x times X directly instead of kind of storing it in this intermediate thing and then kind of merging it into into a sum we've also fused these two loops over data to make a single pass over over over the input so in a single pass over the data where we're constructing both map and sum so this is actually an optimization that is a little bit data dependent or rather program dependent so you can actually there's effects like trashing the instruction cache and stuff so it's not always worth doing this but for simple programs Weld will basically run like a internal cost model and decide whether it wants to do it or not there's also some smaller things like just inferring the size of outputs to pre-allocate data and then it also basically signaled to the backends that some of these things can be vectorized using Sindhi so there's a Sindhi type in Weld and basically adding that to the loop will say hey this thing is safe to kind of execute in parallel using data parallel instructions in the hardware so there's a number of rule-based optimizations I won't go through all of them but many of these as I said kind of come from databases and compilers but they're just nicer to apply in our IR because you know you don't have to worry about things like alias analysis for example so the second class of optimizations is these adaptive ones and these ones come in when optimizations require data dependent decisions and the most classic example this optimization is something called predication where you might have a branch in your code and you're trying to decide whether you want to vectorize that branch at the cost of having to evaluate both branch targets and so to kind of dive into what that means here's a simple loop it has a branch so if X is greater than zero you call foo of X and otherwise you call bar of X so one way to possibly speed this up is to use Cindi vectors so the hardware can actually evaluate on multiple values at once but in order to use these same vectors you need to unconditionally evaluate foo and bar and then basically use this masking instruction in the CPU to select which ones you want to merge into your result but if you look closely you're actually doing more work on the left on the right side here because they're always evaluating foo in bar so if the branch is really predictable or it always goes one way or the other it might not actually be worth doing this if foo and bar are really expensive for example so it is dependent on the data it depends on how often X is greater than zero so the thing that well does to kind of handle these cases is to basically generate code to sample the predicate and determine at runtime which iteration it should run so it'll basically generate code like this where there's a sampling function on the input and it'll evaluate that branch condition basically to figure out how often that branch is expected to be taken so it'll get a probability and then it'll use that probability to determine using a cost model whether it should run the predicated and vectorized code or the branched code and then a lot of this stuff again is safe to do because we have this kind of nice functional ir so you can you know run this branch condition without having to worry about side effects for example okay so in terms of implementation well it has a pis and C Python Java and rust so these are the languages that you can use to kind of build your own world enabled libraries and they're also the things that you can just call weld from we have a full LLVM based cpu back-end that automatically generates parallel code we also have a partial envy PTX based GPU back-end so it'll actually generate CUDA kernels and manage things like streams to offload things under the GPU but it only supports a subset of the language right now in terms of implementation it's written in around 40,000 lines of rust with the backend and LLVM and we have partial prototypes with pandas and numpy which are kind of the most fleshed out but we also have kind of more packy prototypes with tensorflow and apache spark as well in terms of integration effort so how hard is this actually to use if you want to enable it in your own libraries there's a small upfront cost to enabling Weld integration so in our integrations it took roughly 500 lines of code for each library prototype so this is things like adding raisi evaluation into pandas and numpy which don't have lazy api's but also doing things like marshalling data between the kind of numpy array format and the well dry format and so these are usually 0 copies it's just like you have to figure out where to put like the pointer that world can access but it is some kind of glue code that you need to add after that each operator is kind of easy to port over so I think our biggest operator was around 30 lines of code we're like there's this pandas function that performs like a weird hash join or something so that one was around 30 lines of code the simpler ones are like a line or two of code each using the runtime API to actually submit IR fragments you know into the runtime and it's also incremental e deployable so the runtime API will kind of automatically figure out when it needs to convert between a well data type and a native library data type so you can actually interrupt with native native operators you don't have to go through and port your entire library you can kind of focus on the most important things and kind of incrementally add more things as you need to and if you're just an application developer then ideally you don't really have to change anything at all so you just import this new library and you run your code as always and it should just work if you want to try it out in Python you can pip install PI wells we have a subset of pandas in this package called grizzly which is a faster bear apparently and then we also have a well-defined numpy called Walden on pine so those are all available on pipeline and we also have a new release coming soon with a bunch of usability improvements as well so look out for that for sure so in terms of performance so we've shown that weld can actually accelerate several existing libraries so in this integration we did with spark sequel we were able to accelerate these T PCH queries by up to 3.5 X now so most of this speed-up actually came from just doing native code generation and introducing things like 70 vectorization so sparkle already do things like pipelining and loop fusion in its catalyst optimizer and it also actually does code generation I'll but it generates Java bytecode and not native code so by just generating native code that can leverage you know special CPU instructions we were able to get this meet up with numpy we were able to get a 4.5 x speed-up on one thread and a further speed-up on multiple threads so this mostly comes from eliminating things like like redundant allocations of data and also fusing loops so you don't have to unnecessarily kind of go out to memory and fetch values repeatedly and we were actually competitive with tensor flows excel a compiler as well so excel a is a compiler within tensor flow that basically optimizes these tensor computations so it's a similar idea to well but it's kind of more tailored towards just linear algebra so the fact that we were competitive with it but also able to support things like you know joins and sequel and stuff is pretty promising it shows that the compiler is kind of good enough to match these specialized systems as well but while really shines when accelerating multi library workflows so this is a result I showed already but in this data cleaning and linear algebra workload with pandas and numpy we were able to get a hundred and eighty X beat up on 12 threads partially from parallelization but also partially from doing things like compilation and eliminating unnecessary data movement for example in this workload we're doing image whitening as a pre-processing step in numpy and then running a linear regression classifier in tensorflow this is over images and here even though tensorflow is still using this excel a compiler we were able to get a 9x beat up because we can actually look a car across those library boundaries and do things like pipeline the image whitening step with the linear regression step so instead of kind of running these as like two different phases we looked at the whole program and compile it down into this one really efficient kernel and so we were able to outperform tensorflow pretty significantly here just by taking a more global view and then again in spark sequel we were able to speed up this linear model evaluation by around 6x by basically compiling a Scala UDF to native code and kind of co-opted amaizing it with Sparks mapping operator finally we also showed that incremental integration actually you know gives you incremental benefits this is on that black-scholes workload and it's meant to show that implementing more operators will actually kind of give you incremental benefits so it's not like implementing one Operator or implementing all the operators is required to get 90% of the benefit as you add more operators as long as you pick the operators that are most impactful you'll kind of get benefits slowly see this kind of points to the fact that you can take your most important operators and put them over first and as people need them you can you can add more operators as needed you got incremental benefits with incremental innovation sorry you mean why does it go up at around six yeah that's a good question so it goes up around six because basically the way that the operators are ported there's some operators that are repeated so after a point when you port some number of operators there's an extra step of going from well to numpy and back that happens more times at that specific point yeah so yeah if you're interested we also love contributors today we have around 30 or so contributors we have over 2,000 github stars and there's many many things to do ranging from compiler optimizations to better kind of Python integration to you know better the GPU support so no matter what your interests are I think there's something to do in this project we also have contributions by others in academia industry so CWI's University and Amsterdam that's kind of used it to build a new code generating database kkh is a university in sweden that uses it to build a new streaming system NEC actually has this custom vector accelerator and they used well to kind of target Python workloads for it and then data breaks exploded using it for SPARC as well and I also want to thank the Stanford team this project obviously wouldn't have been possible without work from a ton of other people so deepak james met a party prediction and rahul were also instrumental to building this out yeah so to conclude a common runtime can speed up analytics applications by 10 to 100 X on modern hardware and this is just by kind of taking a global view and analyzing the bottlenecks on modern hardware and kind of optimizing around them by you know not kind of dividing your code up artificially into these function calls and some the best performance you're going to have to kind of take an end-to-end view and so well this is kind of one way to do that my contact and the code is up there and with that I'm happy to take questions Thanks [Applause] all right thank you shall make are there any questions yep awesome talk you compare and contrast weld with its nearest competitors with it's nearest competitors if there are any yeah so I think the one that comes to mind for most people is number the biggest difference with number is that weld is kind of trying to look at functions that are called dynamically and then optimize across from them whereas in number the model is that you add this decorator over a Python function and that'll compile it for you so it's it's a similar model in some sense in that we're both doing compilation and kind of looking at a lot of code at once to make it more efficient well it is better for these kind of dynamic settings and numbers better for these settings where you have no what you're going to run then you can kind of decorate it and just like try to guess we don't that way all right do you have any comments on what support there might be in Weld for enabling of streaming and also stateful applications yeah that's a good question so in terms of stateful applications we've actually done some work to support that already so the kind of execution model is that you have some code that you compile and then you also have this other object called the context that lets you store state after a particular well run so actually the way that most of our integrations work is that you feed batches of data into weld and keep updating state in that context and then at the end you basically run a small program that materializes the final results and pumps them out in terms of streaming that group at kth I mentioned is actually working on streaming we're not doing any active work on it but it's something that people ask about a lot so we're planning to kind of reach out to them and collaborate with them or yeah have you looked into like torch integration PI torch integration for speeding up tensor computation we have yeah so I think one kind of way to make this easier to use is having like a nice front-end for it so people can kind of write like UDF's for example using weld so we've actually considered using torch script to do that some-something we're still looking at but definitely you we've looked at it yeah when you were comparing your product to spark SQL were you incorporating Scala native it all with Scala native was that the question were you really compiling with Scala native yeah on the spark said no we were not using Scala native we were just comparing against sparks kind of tungsten code generator by that that just generates Java bytecode so it's not it's not really native execution do you support any other GPU back and other than Nvidia not currently now all right thank you very much we're out of time can we think show me again [Applause] [Music]