Aaron Watters - Using WebGL2 transform/feedback in Jupyter widgets | JupyterCon 2020

Key Takeaways

hello my name is aaron waters and i'm a senior software engineer with the flatiron institute in new york city with the simons foundation this talk describes how webgl 2 makes it possible to use the gpu of a standard computer to do highly parallel computations by making use of a feature called transform feedback a programmer can capture the uh outputs from a gpu computation and either send them into a traditional gpu processing stage another one or extract them extract the data for other uses this talk describes support libraries i built which simplify the process of building transform feedback processing stages and how transform feedback stages can be created and controlled from within a jupiter notebook the talk describes an implementation of the marching tetrahedra method for generating three-dimensional iso surfaces from dense arrays as a concrete example of the marching tetrahedra of transform feedback the marching tetrahedra method is implemented in two feedback stages that i will be showing you so let me share my screen now [Music] and i will start my talk all right first off i'd like to thank some people um i always run out of time and forget to thank people so i want to thank people up front i have a a lot of collaborators that i've worked with at the simons foundation who helped me who inspired me to build these things and also uh helped me debug them and offered suggestions and also i have an intern uh uh ms lu luo who helped me with this presentation and i've had a lot of help from the jupiter community particularly the jupiter widgets community and notably jason grout and martin brettels and also people like brian granger and uh and many others and also i've i've used i'm i've been inspired by three.js which is a three webgl library that's available in javascript it's an amazing piece of software which i will be talking about here all right so let's talk about uh use cases for transform feedback this is the biggest use case which is a rep a uh viewer for three-dimensional matrices so you have a three-dimensional matrices of values which represent something in this case they represent the density of a inside a human's human body as determined by mri processing so first we get the matrix then we load the matrix into the widget and that's all you have to do and then you have a view of the widget and you can rotate the view [Music] and you can change the intensity level so here we have the skin if we go back a little bit we get back to the bones and then in addition to that we can focus in to some part that we're interested in like the spine and we can rotate that around or zoom in a little closer rotate it in and look and see what the spine looks like this data was derived from mri data from the national institute of health and i sub-sampled it down so that it was a little bit smaller you can do larger uh arrays but the performance begins to get jumpier and uh so let's so that's the uh the most the the basic use case uh how is this implemented so the um let's talk a little bit more about so just to review to use the the widget you need a matrix in this case i'm i'm creating a matrix which is just a very simple matrix where every entry has the the distance to the center of the matrix so what does that give you when you view it in the viewer it gives you a bunch of spheres so here if we go down here we have a sphere we rotate it around and if we go up the sphere begins to insect intersect the edges of the matrix so you have a truncated sphere so the various components here this is a 3.js view this is a 3.js view this is a these three slices in the x y and z dimensions are canvases html5 canvases and they're built using something called jp doodle which is a library i built so let's talk about what the components are with the uh the top level component we were looking at is in a module called volume volume uses a module of the surface rendering module and the surface rendering module uses webgl3 for presentation and a module called feedback which implements the html file or the webgl 2 feedback mechanism i'm talking about today and then of course those are both based in webgl 2 which is an underlying which is the underlying functionality also it's where the object itself is a widget which is built using something called jp proxy widget which is a helper func helper module i built and it's which is built on top of i pi widgets which again is a part of part of jupiter which of course uses other features of html5 and i'm also drawing the you know the x y and z slices are drawn using something called jp doodle which uses canvas 2d so that's uh sort of a 10 000 foot view of the where things fit together now let's talk about um [Music] what uh the uh the basic webgl so what the what i'm talking about is um today is um transform feedback in webgl 2 which is built on top of basic webgl so what is basic webgl it's an amazing library which allows a programmer to um to do things like 3d modeling in a web browser and there are two uh it has a standard pipeline and the standard pipeline the input to the standard pipeline is usually an array is usually usually a collection of triangles these go into something called a vertex shader which which does things like move the the vertices of the triangles to different positions or uh sometimes besides colors and normals and things like that can do a lot of different computations that goes into some components that you're not allowed to change and ultimately enter something called a fragment shader which decides on the color for each individual um pixel with of of each triangle and then that gets output to the screen now we're mainly interested here in the vertex shader because that's what that's where the uh that's that's the part that we can access using the feedback mechanism a very simple vertex shader is shown here all it does is a couple matrix it takes an input vector and uh two input vectors actually and multiplies them by two matrices basically that's not very interesting from a mathematical point of view the interesting thing from a computer science point of view is that you can send in literally tens of millions of vertices into into the into the processing pipeline and this vertex shader will process them all in parallel roughly speaking on thousands of processors instantaneously uh so you can do things like uh take a complex geometric figure and rotate it in real time as i've already demonstrated with the uh the widget i showed you so for example if you want if you wanted to represent a chimpanzee the first thing you need to do is find triangles which represent the surface of the chimpanzee that's called a mesh and then you take the take the mesh so you take the the the chimpanzee divide it into triangles and then you load all those triangles into the gpu and then you run the rendering pipeline with an appropriate with appropriate shader programs and the vertex processing stage will broadcast your rotation matrix to all the vertices of the chimpanzee in parallel and it looks like this so you have a a single matrix which is shared by all processes and then millions of vertices down here millions of processing um millions or thousands of processors handle each each of the vertices in sequence and outcome millions of rotated vertices and then if you want to change the rotation what you do is you change the matrix and leave everything else alone tell the processing pipeline to run again and you've got a new rotated uh chimp now uh if you want to there's also a feature of webgl that we're going to be that i use which is called instancing and for instancing you don't want one chimp you want thousands or millions of chimps that are all um described by the same mesh by the same vertices and in this case you you can describe the mesh once and so the mesh is just the vertices of the mesh are on along the top here and then you can change the position of each so you can specify the position for each chimp uh separately along a separate dimension and then have the common uh any common attributes like the matrix broadcast so here we have each processing element gets one vertex from the mesh one position from the position ray and uh the common matrix combines them into a rotated translated uh vertex and produces an output and and this is all done very quickly in my case i don't want to generate chimps but what i want to generate is isosurfaces and the way that works is you take the voxels of the matrix each voxel is broken into six tetrahedra with two triangles each with three vertices in each triangle and so um so here we have and so here we have the the corners of each of each voxel along the vertical dimension and then along the horizontal dimension is is all the triangles it gets divided into so that's the outer look that's how it looks in a uh in a so that's how the um the iso surface maps into the um uh the the the instancing model now let's look a little bit about talk a little bit about how feedback works okay transform feedback is the idea that so for example i have a method which produces iso surfaces i don't want to have to i write a program which which did which uh you know there are there are libraries out there like 3.js which are very good at presenting triangles once you have them so i know how to get the triangles but once i have the triangles i don't want to have to figure out how to color each individual each individual pixel i want to use 3.js which is a library to do that and the way you can do that is you put something in the front of the processing pipeline so here i'm going to stick something in the front call which is called which is the iso surface triangle generation and then the triangles and that runs on the gpu and when it's done it feeds all the triangles into the standard pipeline that's created by 3.js in order to produce the result in the final result that's displayed so um so the way the way it looks is it looks a lot like a standard vertex shader you have uniforms vertex attributes instance attributes and a shader that produces some output but we can capture that output call it feedback output one and then we can send that output into the next processing stage so here we have the output from the previous processing stage being fed in where the as instance attributes to the next processing stage and that will produce feedback output too which can go into the standard process processing pipeline or it could go into a third stage potentially it can go into many different stages and the the name feedback comes back comes from the idea that you can actually take the output of the feedback and feed it back into itself and that will allow you to implement things like particle systems now feedback to program this it's actually quite difficult if you use the direct the direct libraries provided by opengl so what i've done is i have created something called the feed webgl wrapper library which does all the standard things in a standard order for example here is the run function which i'm not going to talk about in detail but basically it does a bunch of operations that you always want to do in the same order for every run you know basically it's it captures the the you know the basic underlying operations the the the underlying operations are very much like uh programming a state machine or an assembly language uh program uh this is a little bit more higher level it groups things into objects and objects which can be linked together and um but it doesn't try to come up with a a new paradigm it lives within the same paradigm but it's just a slightly higher level uh gets rid of the boilerplate code that you tend to need to use when you write feedback programs so i'm going to look at a trivial example now so the trivial example looks like this we have a single vector a three three dimensional vector called shift and then a whole bunch of input vectors and we want to send it into a shader program which produces the input vector plus the shift as an output and uh the so we i'm going to use a standard a very simple shader program to implement it and then the in order to get the the wrapper library to run this program what you need to do is you need to specify what the variables are and how they are linked to data inputs and outputs so that's here's the program the program says okay there's an input vertex and there is a shift vertex and there's an output vertex and what i want you to do is add the input vertex to the shift vertex and produce an output vertex so that's the program very stupid program as a you know standard hello world example now we need to link this to values so first the shift factor is a shared value so in this case i'm going to ship uh here the the value is going to be -1 0 and 1. and this chunk of code here says oh shift vector is a uniform and here is its value when you link the program together now the input the input vertex is not a single vertex it's actually an array of vertices uh and the array of vertices comes from a buffer so we declare the buffer down here and then up here we say oh the input vertex variable is an input variable and it comes from it is it should be bound to the buffer that's called the input vertices buffer so that's how you link things together and finally the output var vertex is specified in this chunk of code the output vertex is a feedback output and it has three components each so basically you need to tell the the the system uh what the variables are and how they are linked to data sources in a declarative manner and it runs like this we import the module we say okay basically we're we build and compile the program and then we run the program and then we get the feedback outputs back now what is 0 2 4 well 0 2 4 is so 0 to 4 is 1 2 3 plus minus 1 0 1. now once we've got the program on the gpu and it's all compiled we can run it again so let's run it again we're going to change the value of the shift vector to 300 200 100 and then we will uh run it again and then uh so we get 301 201 as as the next level so of course this is a totally artificial example uh if you uh it's much slower than using numpy to do the same thing of course because there's so much overhead but if we had 30 million vectors it might be faster and in particular it might be faster if the operation we were implementing was much much much more complicated which it generally is so so just very quickly i'm going to talk about how to do use this system for a more complex problem which is the isosurface generation so again just to review to generate an iso surface you start with an array and what you want to do is convert it to a bunch of triangles that look like this here are the triangles and actually the you want to produce triangles and normals which are so here's a diagram which shows the the that each each triangle has a normal vector sticking out of it those are the little arrows coming out of the triangles uh those are used for coloration all right uh so uh in order to do this the way it's done with marching tetrahedra which is very simple method conceptually is you take each voxel here we have a voxel with two negative values and six or three four five yeah six positive values and what we want to do is find triangles which cut cut between them at the zero threshold uh interpolation um what you do is for the voxel you divide the voxel into into six tetrahedra and then for each tetrahedron you decide oh does it have two above the threshold and two below the threshold in which case you use two triangles or does it have one above and three below or three below or three above and one of one below in which case you use one triangle uh and and that's how you determine the triangles except that there's another case which is all of them are either above or below the threshold in which case there's zero triangles um and uh you go through all the all the boxes and you produce those triangles and that produces your isosurface the so so that's the basic idea how is that implemented as a transform feedback pipeline all right so once again we have an array we want to produce a bunch of triangles like that so we need to break down the steps for producing the normal vectors in the triangles okay first of all on the gpu you identify the voxels the corners of in the array which possib which contains some piece of the isosurface and you do that by looking at the corners and seeing if if it has one corner above and one corner below then it contains some piece of the iso surface so that looks like this you have you send in a shape uh you send in the threshold value you send in the corners of each voxel you test to see if there is a crossing and you return the corners and an indicator saying one or zero is it crossing or not now the next stage at the next stage um oh so so the output of the of the of the test stage is the filled the filled voxels are have some piece of the of the isosurface inside them so the next in the next stage which is actually implemented in javascript you throw away everything except the interesting voxels and that usually throws away almost all the voxels and that's a very important optimization step right there and then and then finally you apply the method i described marching tetrahedra to each of the remaining uh to each of the remaining boxes and that produces the triangles and that's so that's the basic step so there are there are basically two steps you identify the voxels you generate the tetrahedra once you generate the tetrahedra you get a bunch of triangles and then you send it into the standard pipeline so that's the um that's the basic sequence of steps all right now summary and conclusions so i've described something called webgl to feedback which implements highly parallel computations uh it doesn't implement an arbitrary computation because you have to have something which has very few inputs or few inputs and very few outputs and that there's no commun no communications needed between processing elements uh at least in a given stage you can communicate at the next you can you can feed back into the next stage and the next stage can see what the previous stage did but you can't um you can't communicate within the same stage uh the whole so it looks like each processor sees something that looks like this there are uniforms and textures are interesting there are small number of arrays which can be large so you can send in a large amount of data but it has to be only in the form of arrays and so the each processor sees a bunch of uniforms and arrays and possibly instance attributes definitely vertex attributes and it produces some sort of feedback which gets captured by some of some downstream processing the overall architecture at the gpu level looks like this you have uh the uh the vertex you know you have many vertex attributes and many instance attributes that get where they get sent to individual processing elements and then the shared parameters are broadcast to the processing elements and then the the it produces a matrix array of outputs a lot of operations look like this and uh the obvious one would be something like a particle simulation uh or um isosurface generation like i've demonstrated but uh i am now at 36 minutes so i think i will stop and thank you very much and thanks i'm delighted to have given this presentation i'd like to thank the organizers and all the people who helped me out with this work thank you so much