Mykhailo Moroz

Writing an optimizing tensor compiler from scratch

2024-09-11T00:00:00+00:00

In this blog post I want to talk about the research and development results for a library that I started working on more than a year ago - TensorFrost. Under the hood it’s a static optimizing tensor compiler with a focus on being able to do more “shader-like” things while still keeping the ability to do high level linear algebra for ML in Numpy-like syntax with automatic differentiation support. (Click on the example GIF’s for more details!)

For documentation on basic functionality, read the README file in the repo.

So why make a new library?
Architecture
Python frontend
Backends
- Codegen
- Runtimes
Examples using TensorFrost
What’s the current performance compared to other tensor libraries?
What is left to do
Conclusion

I started working on this library around 14 months ago, initially I didn’t really plan to do much more than a few matrix operations for an optimization algorithm I wanted to implement in Unity, but there were quite a few things that I wanted to have on top of all of this and it sidetracked me into a writing an entire compiler (hello scope creep 👋).

The thing is, it’s not the first time I tried to make a tensor library, the first time was a whole 5 years ago and used OpenCL, as I didn’t have an Nvidia GPU at the time. To be honest I’ve been completely unprepared to the magnitude of what it required, and while I did get it to a “somewhat” working state like having basic kernels and somewhat working autodiff using a tape, the lack of good debug tools and actual problems that I wanted to solve pretty much killed it. And for the things that I did want to work on, it was completely unsuited for, as I usually write simulations or graphics, and the overhead of doing a kernel per operation (especially unoptimized kernels) is just too bad to be useful.

Since that time I’ve had a lot of ideas of what I would like a library like that to even look like, and wanted to try working on it again. However I did know, that for this project to survive I would need to make it suitable for projects I usually do, like the ones I usually do on Shadertoy. It might seem weird to you as to why I would make a specifically “tensor” library for something that is basically equivalent to writing shaders. But to be honest, shaders are not actually a perfect place for what I do, and a lot of simulation/rendering algorithms can map quite well to high level “tensor-like” operations. While the limitations might force you to come up with creative solutions, for really large or complicated projects it just becomes more of a problem, as it’s very hard to iterate on quickly. This was also one of the main reasons I didn’t really touch ML too much for most of my pet projects (except stuff like Neural Implicit Representations), ML algorithms are usually quite orthogonal to the way you write shaders, usually being split into hundreds of kernel dispatches, while those shader algorithms are effectively just one megakernel most of the time. The only reasonable way to integrate neural networks into those is unrolling the network into a single scalar function, which can be quite unoptimal and limits their size. Not even talking about the fact that training them is completely out of the question. This brings up another problem, shaders don’t have automatic differentiation, which is surprisingly much more useful here than you might think. While its usually used for optimization algorithms like SGD, having the analytic gradient can also be useful for computing normals/curvature of analytic shapes, or forces from potentials in simulations.

So in this library, I hoped to somehow extend the applicability range of a Tensor library to more “shader-like” use cases like rendering and simulations.

And right now I can actually say that at least to some partial degree it did work out. Currently the library is something of a mix of slightly more low-level Numpy-like operations with shader-like control-flow and operations (most of the built-in scalar shader functions are passed through to Python). In terms of where it stands right now, its more low-level than something like JAX or PyTorch, but still not as low-level as Taichi as you still technically operate on something similar to tensors.

So why make a new library?

It will indeed take a inordinate amount of work to make a library from scratch and get it to a useful state, as I’ve already experienced. Why would I not just use an existing library, as there are seemingly thousands of them? There are a few reasons mostly applicable to my use cases which make both using pure ML libraries or pure shaders annoying.

1. Performance scales poorly for non-ML specific operations

Of course, nothing stops me from using your usual ML libraries like PyTorch or JAX, but as I mentioned before, they weren’t really designed to be applied to problems that I have. These libraries effectively live in their own realm with their own rules and syntax fine tuned for ML, and hide some of the features the GPU has from the user, they pretty much have 0 crossover with how shaders operate. While technically you can write any algorithm you want in plain tensors, including graphics and simulations, depending on the number or complexity of operations - the performance could get terrible.

Most of the Tensor compiler research that I’ve seen focuses on ML bottlenecks, like efficiently utilizing cache, correctly aligning data for maximum performance of matrix multiplications, convolutions etc. Those usually aren’t a bottleneck when dealing with simulations or rendering, the bottleneck becomes the dynamic nature of the code and its complexity.

2. Dynamic control flow is very tedious (if possible) to implement in classic ML libraries

The performance is actually not the only issue when writing simulations or graphics, control flow can be very prevalent, but unfortunately it’s quite inconvenient to express in these libraries, if even possible, as native loops in JAX for example, require doing stuff like:

def factorial(n):
    def cond_fun(state):
        i, fact = state
        return i < n

    def body_fun(state):
        i, fact = state
        return i + 1, fact * (i + 1)

    initial_state = (0, 1)
    final_state = jax.lax.while_loop(cond_fun, body_fun, initial_state)
    return final_state[1]

Which looks completely cursed. Now imagine multiple stacked loops - yikes

Such native loops are required if you want a varying iteration count for some operation, in simulations this could be, for example, summing a force from a varying number of particle neighbors. Even something like Gaussian Splatting requires a variable loop per tile. From a purely classic ML standpoint, this isn’t really a problem, since such cases practically never happen and the computational graph is absolutely static, and whats worse, autodiff gradients of such loops might potentially have atrocious performance (or might just be uncomputable if the loop can be theoretically infinite, as the compiler might not know the specific context where its being used).

You could alternatively write a dynamic mask that will depend on the iteration, and unroll the loop, but this would simply be slower and less readable. Like here, I once tried to make a vectorized Numpy function to computes the mandelbulb SDF:

def mandelbulb_sdf(pos, iter_num=mandelbulb_iter_num, power=mandelbulb_power):
    z = pos
    dr = 1.0
    r = 0.0
    for i in range(iter_num):
        r = np.linalg.norm(z, axis=-1)
        not_mask = ~(r > 1.5)
        
        theta = np.arccos(z[..., 2] / r)
        phi = np.arctan2(z[..., 1], z[..., 0])
        dr = np.where(not_mask, r**(power - 1.0) * power * dr + 1.0, dr)
        
        zr = r**power
        theta *= power
        phi *= power
        
        z = np.where(not_mask[:, :, :, np.newaxis], pos + zr[:, :, :, np.newaxis] * np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]).transpose(1, 2, 3, 0), z)
    return 0.5 * np.log(r) * r / dr

3. Extending ML libraries with custom high-performance kernels requires using external tools/APIs like CUDA or Triton

Usually, when hitting a performance bottleneck, you have to write custom CUDA kernels, which is doable, but quite inconvenient, as it forces you to use separate environments, CUDA and Python.

There are actually domain specific languages (DSLs) that allow you to write kernels in relatively high-level Python, like Taichi or Nvidia’s Warp, they are really nice for simulations or graphics, and Taichi specifically is fine tuned for high performance physics simulations and even has built in automatically optimized sparse grids. And while they are also differentiable, they still are a bit off from what I would consider “perfect”, as you don’t have access to the giant library of ML operations that PyTorch has for example. Of course you can still write ML using them if you really want, there are some Nerf implementations for Taichi, but they require much more code to represent. As a compromise you could also interoperate them with PyTorch for example, but once again, it makes it less convenient to work with. I can also mention Triton here, but it’s usually used more like a backend for other libraries (like PyTorch) rather than a standalone DSL, at least as far as I have seen.

There is also Slang, which is quite different from all these from above, as its an improved shader language with added differentiability. It would be nice if it was widely supported. But its even more low level than something like Taichi.

4. ML libraries don’t have an easy built-in way to make real-time visualizations with optional interactivity

When I want to do some advanced visualizations in Python, the options that are available in ML libraries are often hilariously bad. Usually you just end up making a bunch of matplotlib plots, which are not only slow, if you want to render like hundreds of millions of points or an animation, but also not interactive. (They are fine for papers tho)

In the world of real-time graphics, you can render those points at 60fps, and you could even interact with them in real time. Seeing the things you are working on in real time is in some cases quite useful for quick iteration, and I feel like this is something you miss in your classic ML libs, where the usual interaction you have with your model - is staring at the loss graph for hours.

Though, while I am stating these things, most large ML models are simply not visualizable in real time, and the ones that are, are usually not easy to usefully interpret. Visualizations are usually most applicable to the intersection of ML/Physics/Graphics, like NERFs, diffusion, neural implicit representations, etc. But I still think that even changing hyperparameters in real time and seeing its result on the training loss can also be somewhat interesting, though you do need the model to be rather performant for that.

The lack of a native way to output graphical data from these libraries is even more annoying when you remember that GPU’s are called Graphics Processing Units, not Tensor Processing Units. And they have all the required hardware to work with and output graphics.

PS. Taichi actually does have a way to this! It has integration with GLFW and ImGUI.

5. Writing simulations or graphics in a high-level language is could be much easier to iterate on rather than in pure shaders

On the other side, in the world of real-time simulations and graphics, I’ve written custom kernels for every specific algorithm something needed: radix sorts, linear solvers, numerical integrators, etc. So when I’m prototyping or having a new idea how to optimize the algorithm globally, it can get annoying to make global changes in the code structure, since they usually require a partial rewrite, creating new kernels and so on, and I don’t really see why this couldn’t be automated from higher-level operations.

In the end I was wondering: can I somehow combine the best of both worlds? Being able to do both Numpy-like operations while also doing more shader-like things in one place sounds somewhat impossible on paper, but I thought that maybe if you tuned the kernel generation algorithm to specifically be optimal for shader-like operation it might at least work for my use cases? Afterwards I could still support ML use cases well enough even if I just shoehorned the matmul/reduction/convolution kernels separately.

And even if it is impossible, combining everything into a single language would be nice, because the GPU development infrastructure is scattered all over the place, some features are available only in some places, while others are not, they can sometimes not even be interoperable. The development environments are completely separate, there are ML-specific debug tools on Linux, and graphics API specific debug tools on Windows only - and all of these things are using the same hardware - the GPU!

Architecture

Kernel fusion

When thinking about the architecture of what a library like this could be, I’ve first thought from the point of view of simple per-element operations and indexing. Nothing really stops you from writing Numpy code as if it was a shader, for example. However that would be extremely slow, why? Numpy usually executes a single operation over the entire tensor at a time, i.e. loads the data, does the single operation, then puts it back into memory. For a modern computer this is extremely inefficient, since the bandwidth and latency of system memory are nowhere near sufficient to keep up with the computational power of the processor. If you want to utilize all available resources - you need to utilize the cache hierarchy to its full extent. If you could cache intermediate results close to the ALU’s of the processor and keep the intermediate results as long as possible there, you could get a massive speedup by reducing latencies by orders of magnitude. This is why modern processors have multiple levels of cache - they try to solve the issue of memory improvements not keeping up with increasing performance. As a matter of fact since 2017 top tier consumer GPU memory bandwidth has only increased from 500Gb/s to 1Tb/s, while performance has spiked from 10 TFlops to 80 TFlops.


Tiers of memory of a typical GPU

But optimizing this specific aspect when dealing with tensors is actually nothing new, the so called kernel fusion has been around for a while, and is used in tensor compilers like XLA or PyTorch’s TorchInductor. And the degree to which they can do it is perfectly fine for most ML applications. However, to my knowledge, they do not fuse operations with already existing optimized kernels, even though in some cases this might be beneficial (Well, in TorchInductor there are only 2 operations that don’t - conv and matmul).

If you tried to write algorithms, like the ones I write in Shadertoy, you will eventually start to hit the limits these compilers have. The number of operation nodes you could fuse now rises to the order of thousands, not even mentioning the complex control flow, and ideally they should fit into a single kernel, but its highly likely you will end up with a lot of smaller kernels if you apply fusion naively.

First prototype

When I initally started prototyping the operation graph compiler in C# in Unity (not exactly your typical place for a compiler prototype, I know), I kept the graph just as a simple Directed Acyclic Graph (DAG), where each node was a single operation with some tensor shape. When I began testing the kernel clustering even on simple physics or rendering, the clustering algorithm quickly started to get out of hand.

Here is an example of a compiled graph with the kernels clusterized.

This was the operation/kernel cluster graph for this fluid simulation:

It did work, unfortunately there was no good way to procedurally generate shader code in Unity, so I did something rather stupid and implemented a Virtual Machine right in the shader, which made it particularly slow. However, of the good things, there was practically 0 compile time, but unforuntately that was not the goal I had in mind. The way I did reductions to get the average energy here was also rather questionable - I did it purely with atomics. For some cases they are good, but reduction is not one of them. Especially given that there are no natively supported float atomics in HLSL, so you need to emulate them through InterlockedExchangeCompare, which makes it even slower for this particular use case. One time it even crashed my computer when trying to add a few million floats in parallel to a single element.

The VM basically only knows of 1 dimensional loads and stores, so I also needed to add a compilation stage that converts multidimensional indices into operations that compute the “flattened” index.

At this point the graph had only a few operation types, and I implemented backward-mode autodiff here, which was surprisingly easy. The only not super obvious thing initially was load/store gradients. But those are effectively just atomic_add/load respectively. So in theory, you could implement any ML algorithm even here, but that would be comically slow. Matrix multiplication gradient would be 2 atomic adds per thread in a 3D kernel.

Another thing that turned out to be a huge problem was that the operation graph did not have a “stable” order, it was only topologically sorted (I actually resorted it every time I did something with the graph!), which is good for normal operations, but for inplace operations like stores and atomics this leads to randomly varying results, which is a no-go.

Adding the problem of exponentially growing number of possible kernel fusions and the graphs quicky became an undebuggable node soup, which made this approach not very appealing:

I also wanted to have at least some sort of control flow, which wasn’t obvious how to add into this specific graph for me, at the time.

Second prototype

Any sufficiently complicated C or Fortran program contains an ad hoc, informally-specified, bug-ridden, slow implementation of half of Common Lisp.

This time I wrote it in C++ and decided to instead enforce the order of operations to make in-place operations have properly specified order. I constructed the IR like a linked list, order of which is taken from the way the code is parsed. This simplifies the kernel fusion problem to one of finding ranges of fusable operations, and in my case I fuse everything except any node pair that violated rules based on order of memory access, shape of the operations, etc. In some cases this effectively can convert the entire graph into a single kernel, which is exactly what I’m looking for.

Ordering the operations is not the only thing. To implement control flow I needed to have child and parent nodes while still keeping a uniquely specified ordering, I implemented this by using a multi-level linked list. This way operaitons controlled by a conditional statement or by loops become its children. In multi-level linked lists kernel fusion becomes slightly more tricky, but effectively its just a recursive process of finding the ranges of fusable operations starting from the lowest level of the list, and fusing them level by level.

Having children or parents for a tensor is a rather unusual notion, but it is totally fine as long as you make sure that the shape of this node is broadcastable to all of its parents. So you can, for example, have a scalar loop with children of any shape. Or, lets look at a more complex example, a loop of shape [N, 1], with a child loop of shape [] with a child load operation of shape [M]. Totally different shapes, but they are all broadcastable between each other so the shape a generated kernel here will be [N, M] for all these nodes. There are also cases when you have broadcastable shape between parents and children and incompatible shape between neighbors - those are simply incorrect and should lead to a compilation error, as there is no way to generate a kernel here without complex masking, which I don’t plan to support now. Example would be a loop of shape [N] with 2 children of shape [5, N] and [2, N].

The kernels are also represented in the same IR, as children of a “kernel” node. At the code generation stage, everything outside of the kernel nodes is converted into host code (right now C++), and the kernels are converted into device code (OpenMP C++ or GLSL/HLSL). This way the entire program is represented in a single IR, and the compiler can optimize it globally both for CPU and GPU parts.

Since the IR is the same for all compilation stages, you could for example input both high level tensor operations together with explicitly specified kernels, and it can already be done like this:

A = tf.input([-1, -1], tf.float32)
B = tf.input(A.shape, tf.float32)

C = tf.buffer(A.shape, tf.float32)

with tf.kernel(A.shape) as (i, j):
    C[i, j] = A[i, j] + B[i, j]

C = (C @ A.T) + tf.unsqueeze(tf.sum(tf.sin(B @ A),axis=-1),axis=-1)

This is already quite nice. (I’ll explain the Python syntax a bit later)

(Note: of course, if you tried to compute the gradient here, the compiler would fail, at least at this point in time, as general gradients over control flow are not trivial)

Another interesting aspect of this representation is that kernels can be created as children of control flow nodes, meaning you can create a loop of kernels for an iterative algorithm, like for example a bitonic sort!

log2N = tf.ceil(tf.log2(tf.float(element_count)))
Nround = tf.int(tf.exp2(log2N))
sort_id = tf.indices([Nround/2])[0]
steps = tf.int(log2N*(log2N + 1.0)/2.0)

with tf.loop(steps) as step:
    j = tf.floor(tf.sqrt(tf.float(2*step) + 1.0) - 0.5)
    n = tf.round(tf.float(step) - 0.5*j*(j+1.0))
    B = tf.int(tf.round(tf.exp2(j-n)))
    mask = tf.select(n < 0.5, 2*B - 1, B)
    e1 = sort_id%B + 2*B*(sort_id/B)
    e2 = e1 ^ mask

    with tf.if_cond((e1 < element_count) & (e2 < element_count)):
        key1, key2 = keys[e1], keys[e2]

        with tf.if_cond(key1 < key2):
            val1, val2 = values[e1], values[e2]
            keys[e1] = key2
            keys[e2] = key1
            values[e1] = val2
            values[e2] = val1

In this case the compiler can spot that you can’t fuse the insides of the loop since they are reading and writing from the same memory, thus creating a kernel region under the loop.

However, while you can do that, in this particular case, I woudn’t recommend relying on the compiler too much, and would put an explicit kernel under the loop, since even changing the loading order from before the stores, to the middle, will split the kernel in 2 and potentially break it right now.

I had one specific case when it was an issue when I tried to optimize a software sphere rasterizer, by adding an additional read to check if the atomic min can be skipped I effectively made the compiler think that this part of the code needs to be split into parts and simply just broke the rendering.

This shows that, while powerful, such inference of how the program is structured does not always work, or requires a much more advanced compiler than I have here. On the other hand, in the case of the 2D fluid simulation example, kernel generation worked quite well, since there is no control flow to confuse the compiler.

Optimization and generation of the kernels

Simply splitting the IR into kernel regions is actually not enough to make sure you dont have a million unneeded loads or stores. One very basic way to optimize the kernels is effectively just copying computations that are cheaper to do than to load from memory. Such things like constants, or very simple arithmetic are examples of this. Doing this optimization not only reduces the number of global memory access a lot usually, but also removes a lot of unneeded kernels that would have just stored a constant or something similar into memory.

I also have the standard “removing of unused computation” here. Since we have the entire IR graph from input to ouput given, this additionally allows to figure out which parts of the computation are influencing the outputs. So I can effectively assume that everything else is unused and can be simply removed.

When generating compute kernels out of such an IR, you can not simply use the N dimensional shape of the kernel, and you need to map the tensor indices to the specific layout of the GPU. In this case, that’s the group indices, and the workgroup thread indices (in more advanced cases, in DX12/Vulkan/CUDA/etc there is also the warp sub-group, but I’ll ignore it for now). To do this mapping we ideally would need to figure out the shape of the workgroup, which must be a compile time constant, from the computations we do. But at the moment I simply estimate the group shape from the last 3 dimensions of the kernel, and clamp them to some predefined constants depending on dimensionality. This is suboptimal, but doing it better would either require having a VM that estimates the range of indices of memory accesses, or an autotuner. The first will take some time to implement, and is in my TODO list, as its also useful for other things, and the second, while easier I’m not currenty considering, as it could increase compile times quite significantly. (And they already reach 30 seconds on Windows for my Variational Monte Carlo solver!)

Algorithmic operations

The bare IR does not know about complex operations like matrix multiplication, reductions, etc. These are implemented as a special compiler pass, that converts the high level operation nodes into a series of simpler operations. For example reduction sums tf.sum are converted into a loop of adds, matrix multiplications into a loop of muls and adds, etc.

These simpler operations aren’t yet put into kernels, so the compiler can do additional kernel fusion optimizations on them. As they are written in the same IR and not as separate outside kernels also means that all possible optimizations that that the compiler has - can be applied to them, like inserting simple arithmetic instead of memory loads (useful for operations over procedural data, like random), or adding the activation at the end of the matrix multiplication loop, etc.

(Though to be fair, right now, the compiler optimizes them a bit too aggressively, and can put a lot of needless computation inside a matmul loop for example, this needs to be fixed in the future with better heuristics. Detailed example in IR under the hood section)

Advanced kernel fusion

Kernel fusion by splitting into ranges of the multi-level linked list works fine until you get to more complex situations, that do actually happen quite often in ML, for example, reductions of some expressions.

To optimize even these cases you can do something I call tensor load fusion - you replace a loading operation with the recomputed result of the load target with the given load indices replacing the target kernel indices.

This was tricky to implement, but it allows to fuse operations like tf.sum(A[i,k]*B[k,j], axis=2) into a single 2D kernel, instead of 3D + 2D kernels. This gives a massive speedup in some cases when writing operations in such a naive way. For example this allows you to write out a conv2d operation in a similar form as a sum over a 6D tensor, and the compiler could fuse it into a single 5D reduction + 4D reduction, while giving comparable performance to native pytorch (for small kernels). Here is an implementation in TensorFrost which is effectively equivalent to PyTorch’es conv2d without padding:

def conv2d(X, W):
    N, CIN, HI, WI = X.shape
    COUT, CIN, h, w = W.shape
    bi, cout, wi, hi, cin, it = tf.indices([N, COUT, HI - h + 1, WI - w + 1, CIN, h * w])
    i, j = it%w, it/w
    return  tf.sum(tf.sum(X[bi, cin, wi + i, hi + j] * W[cout, cin, i, j]))

At the moment only the first tf.sum(some operations) gets fused into a single kernel. Which is actually completely fine, as we dont want to merge multiple sum together and we want them to be staged. This is actually why I manually fused the kernel dimensions w and h together, as you need to have a balance between the size of the loop and number of reduction stages. This could theoretically be done automatically, but it would require autotune or more advanced heuristics.

Automatic differentiation

Looking at the example in the section above you might think that autograd would completely fail when differentiating such an expression. Which would be the case if done naively. Whats worse, in the programs I write, loads at addresses are extremely prevalent. If you apply autodiff dirrectly on multidimensional loads you get multidimensional atomic adds. In general, you can’t really do much with them, but in most cases, the indices of the atomic adds are simple, or are constants. In those cases you could check what dimension indices of the operation are not used in the atomic add address computation, and conclude that all threads of this atomic add for this dimension add to the same element. This would mean that we can optimize this by transforming this dimension into a sum over it, with a following atomic add to this element in the end. Even more, you could also check if the indices map 1 to 1, meaning that they do not make conflicting writes, and can just be replaced with a simple load, add and store operation. (However, I don’t actually do this at the moment as nonconflicting atomic adds are cheap enough to be ignored for now)

These optimizations improve performance of automatically differentiated operations of this kind by a lot, and makes prototyping differentiable convolution-like algorithms quite easy. (I’m interested in trying to fine tune a differentiable multi-grid Poisson equation solver, with boundary conditions, by making lots of small convolutions like these)

The rest of the automatic gradient algorithm is your run-of-the-mill backwards mode autodiff. The autodiff pass is before the algorithm insertion pass, so the gradients are computed for high-level operations if possible, as they are usually cheaper and more numerically stable. (And also I don’t have a way to compute gradients of control-flow at the moment, so it works as a substitute for that for now)

Right now gradiets are specified like tf.grad(a,b), and actually, a doesn’t need to be a scalar, the default vector jacobian product (VJP) input is always a 1 no matter the dimensionality of a. This is useful, for instance, when computing gradients of a potential, for instance:

dx = x1 - x2
dist = tf.sqrt(tf.sum(dx**2))
pot = 1.0 / dist
force = - tf.grad(pot, dx)

Quite nice when doing a particle simulation. In fact I also use this when computing the normals for the SDF in my path tracing example. I should note that this is only valid behaviour because these computations are independent, for pixels, or for particles, if they were depending on each other, the gradient would be invalid, and in that case you should use a scalar a.

The compiler when doing the autodiff searches for all unique a’s and does a full backprop over their dependencies, then all unused gradients are simply removed after this compilation pass.

I still plan to implement forward mode automatic differentiation, in its case its somewhat easier to implement. Since for example I can just do it after all the algorithmic passes, as the form of the computation will be exactly the same as the original, just with slightly different operations.

It does pose the question of what to do when doing a hybrid autodiff, like backward grad of forward grad. In that case the gradients need to be sorted by their order, and done one by one. Unfortunately in this case I would need to implement full jacobian vector product (JVP) (on top of the VJP’s) for all algorithmic operations, not just the base simple non-algorithmic ones, so I’ll probably leave that for the far future.

IR under the hood

Let’s look at how the IR looks in the compiler right now. We will look at the bitonic sort example from above to see how control flow is represented in the IR.

Parsed/traced input

Visualizing General Relativity

2022-08-21T00:00:00+00:00

When thinking about renders of things like warp drives and black holes we usually just expect to see a simple approximation or an artist rendition, assuming the math required to pull off something accurate would require someone with at least a PhD in Mathematical Physics. Which I won’t tell that its completely untrue, but in this blog post I’ll try to explain a way to do quite accurate visualizations within a 100 or so lines of code, for basically any kind of space-time for which you can write its metric as code. Although, the detailed mathematical derivation of this approach might be somewhat math heavy.

The main ingredient of any GR render is figuring out how the rays of light move around. Knowing how light moves we can trace rays from the camera into the scene and see where the light came from. So, to render the simplest scene without objects we simply trace a ray for each pixel and assign the color of the pixel to the color of the skybox in the direction in which the ray ends up pointing to.

What are geodesics?
Mathematical description of shortest path
Lagrangian description of a geodesic
Hamiltonian description of a geodesic
Writing a Hamiltonian geodesic tracer in GLSL
Conclusions
References

What are geodesics?

How exactly do we trace rays in curves space? Any object inside a curved space follows something called a geodesic.

A geodesic is just a fancy word for, in some sense, a path of shortest length between 2 points inside a space.

I should note that there could be multiple of such paths, which are locally minimal, in the sense that you can’t nudge the path to make it shorter, while globally there might be a shorter path. Also, in Minkowski space-time the definition is a bit more complicated, because of the imaginary distances (when \( ds^2 < 0 \) ).

In our case, instead of paths between 2 points, we are only interested in finding how a ray moves given an initial point and direction, but the definition above will still prove useful when deriving the equations describing a geodesic, which we will use here.

Mathematical description of shortest path

I’ll try to quickly go through the derivation. If you wish to skip the math part, jump to Writing a Hamiltonian geodesic tracer in GLSL.

Mathematically speaking we have some coordinate system, a path, and a way to compute distances between 2 points.

A coordinate system being a set of several numbers labeling each point in the space. A path is a function that takes in the path parameter and outputs a coordinate, in General Relativity the path parameter is usually proper time (like a clock moving with the object, labeling each point), but it can be anything really. The way to compute distances is called a metric (and it’s the main source of scary math here).

In physics, or more generally differential geometry, a metric is defined as an integral(“sum”) of something called the metric tensor. A metric tensor is a bilinear form \( g(a, b) \), it essentially maps pairs of vectors to real numbers, and is a generalization of dot product for curved spaces. So using a metric tensor we can find the length of a vector in space, and also distances \( ds \) between infinitely close points in space.

\begin{equation} ds^2 = g(dx, dx) \end{equation}

In our case, where we describe vectors as a set of numbers, a metric is simply a matrix product of some matrix \( g_{\mu \nu} \) times the vectors. For our infinitesimal distance \( ds \) we get this expression:

\begin{equation} ds^2 = \sum_{\mu \nu}^N g_{\mu \nu} dx^\mu dx^\nu \end{equation}

Usually the sum is just implicitly assumed by Einstein notation [1].

\begin{equation} ds^2 = g_{\mu \nu} dx^\mu dx^\nu \end{equation}

Here we can actually see that for some simple choice of \( g_{\mu \nu} \) we can get the distances by Pythagoras’ theorem. Specifically for the case when the metric tensor matrix is a unit matrix.

\begin{equation} ds^2 = dx_1^2 + dx_2^2 + dx_3^2 \end{equation}

For a flat space-time like space, we get something similar but with the exception that the time coordinate component is with a negative sign.

\begin{equation} ds^2 = - dx_0^2 + dx_1^2 + dx_2^2 + dx_3^2 \label{flat} \end{equation}

Here I used the \( (- + + +) \) signature, but signs can actually be flipped without changing the geodesics, and in some cases, like for particle physics, it makes more sense to use the opposite \( (+ - - -) \) signature.

Going back to the main question of computing distances, to compute the length between 2 points, A and B, along some path \( x^i(t) \) we simply need to sum the infinitesimal distances together using an integral:

\begin{equation} l = \int_A^B \sqrt{g_{\mu \nu} dx^\mu dx^\nu} = \int_A^B \sqrt{g_{\mu \nu} dx^\mu dx^\nu} \frac{dt}{dt} = \int_A^B \sqrt{g_{\mu \nu} \frac{dx^\mu}{dt} \frac{dx^\nu}{dt}} dt \end{equation}

Where \( \frac{dx^i}{dt} \) is simply how fast the coordinate x changes with respect to the path parameter (“clock”), in some sense can be interpreted as the velocity.

Now our main question is how do we minimize the path length? Here is where we introduce a thing called calculus of variations, which is roughly speaking a way to find how a functional(distance) changes by infinitesimally small variations of its input function(path). Such derivatives have similar properties to normal function derivatives. And in fact, similarly to calculus, to find the extremum of a function (min, max or stationary point), we simply need to equate the variation to 0.

Lagrangian description of a geodesic

There is an entire branch of physics related to variational principles, which states that any kind of physical system has a value it likes to minimize (or more generally make unchanging under small variations of path, i.e. “stationary”). That value is called action, and the function under the integral is called the Lagrangian function of the system. The branch of physics studying Lagrangian functions of systems is called Lagrangian mechanics.

In our case the Lagrangian can be written like this:

\begin{equation} L = \sqrt{g_{\mu \nu} \frac{dx^\mu}{dt} \frac{dx^\nu}{dt}} \end{equation}

Turns out we don’t need the square root for the minimum of the functional to be a geodesic, and we can simply remove it from our geodesic Lagrangian:

\begin{equation} L = g_{\mu \nu} \frac{dx^\mu}{dt} \frac{dx^\nu}{dt} \end{equation}

The proof of this you can find here [2]. The only difference such simplification makes is that the parametrization of the path might be different, but the path itself will be the same. Also notably, the equations for this specific case turn out to be the same, meaning the parametrization is also the same.

Additionally, we want this with a 1/2 factor, to simplify the equations down the line.

\begin{equation} L = \frac{1}{2} g_{\mu \nu} \frac{dx^\mu}{dt} \frac{dx^\nu}{dt} \end{equation}

So, our goal right now is to minimize this functional:

\begin{equation} S = \int_A^B \frac{1}{2} g_{\mu \nu} \frac{dx^\mu}{dt} \frac{dx^\nu}{dt} dt \end{equation}

In general the minimum of a functional like this can be found by applying the Euler-Lagrange equations [3]:

\begin{equation} \frac{\partial L}{\partial x^i} - \frac{d}{dt} \frac{\partial L}{\partial \frac{dx^i}{dt}} = 0 \label{el} \end{equation}

Euler-Lagrange equation derivation

Reintegration tracking

2020-08-31T00:00:00+00:00

In this blog post I’ll explain this advection algorithm and how to use it to make advanced fluid simulations like the ones I made, including Paint streams and Everflow. Before starting, I should give a big thanks to my friend Wyatt for giving useful suggestions on building this algorithm.

Looks really smooth, doesn’t it? It even captures droplets with almost close to pixel level precision. And while it does model a fluid with a boundary it does not use particles directly like in SPH, but is actually a semi-Lagrangian grid based algorithm. Oh, and I forgot to say - it’s also super fast.

Let’s start with a brief history, the initial idea for this algorithm came from trying to extend screen space voronoi particle tracking to avoid particle loss, indeed when trying to store the particle state in screen space as close to the particle location as we can, in the case of overlap one of the particles will be lost, and there is no way to avoid it. But actually we can reformulate the problem, what if we don’t try to avoid particle overlap but tried to conserve the total mass? Like by adding the overlapping particle masses together and weighting their velocities and positions by mass. Actually that was already tried by stb, but it immediately becomes apparent that the total number of particles drops proportionally to the particle/pixel density because of them combining, you might ask if there is a way to separate them back, or if there is a way to achieve approximate particle number conservation? In fact, it is possible, but let’s overview how he does the particle tracking first.

Cellular automaton particle tracking

This algorithm is similar to a lattice gas automaton where the state of each cell is defined by a few discrete levels, is the particle in the cell and in which discrete direction it is moving. In our case instead of using a few discrete states we store the position, velocity and mass of the particle within the cell as floats (or any other number type, depends on the precision you want to achieve, I actually used 2 int16’s per channel since Shadertoy only has a 4 channel output per pixel and I needed to store at least 5 numbers)


Initial Frame	Next frame

The main idea of the algorithm is to loop over all neighbors of the current cell(including itself), integrate the position of each neighbor particle and add it if it ends up in this cell. Above is a visualization of how it looks like. Each cell has a particle with a mass, the mass is shown by opacity, so 0 mass is invisible and 1 is completely dark. The red cell is our current cell for which we want to find its future state, the future position of the particles is shown by the arrow. We can see that there is more than one particle moving into the red cell. So in the end all the particles that end up in the same cell are summed and their positions and velocities are averaged. In mathematical form it can be written as:

\[ M_{i}^{t+1} = \sum_{j}^\textrm{neighbors} K_{i}(\vec{X}_ {j}^{t} + \Delta t \vec{V}_ {j}^{t}) m_{j}^{t} — \textrm{updated mass} \] \[ \vec{X}_ {i}^{t+1} = \frac{1}{M_ {i}^{t+1}} \sum_{j}^\textrm{neighbors} K_{i}(\vec{X}_ {j}^{t} + \Delta t \vec{V}_ {j}^{t}) (\vec{X}_ {j}^{t} + \Delta t \vec{V}_ {j}^{t}) m_{j}^{t} — \textrm{updated center of mass} \] \[ \vec{V}_ {i}^{t+1} = \frac{1}{M_ {i}^{t+1}} \sum_{j}^\textrm{neighbors} K_{i}(\vec{X}_ {j}^{t} + \Delta t \vec{V}_ {j}^{t}) \vec{V}_ {j}^{t} m_{j}^{t} — \textrm{updated velocity} \] Where \(M_{i}^{t}\) is the mass of the particle in cell i on time step t, \(\vec{X}_ {i}^{t}\) is the position of the particle and \(\vec{V}_ {i}^{t}\) is the velocity. The function \(K_{i}(\vec{X})\) is equal to 1 if the point \(\vec{X}\) is inside the cell i, and zero otherwise. \( \Delta t \) is the timestep.

For a square cell the K function is simply \[ K_{i}(\vec{X}) = H(\vec{X} - \vec{C}_ {i} + 0.5) H(\vec{C}_ {i} + 0.5 - \vec{X}) \] Where \(H\) is the multivariate Heaviside step function. \(\vec{C}_{i}\) is the center of cell i.

It’s quite a simple algorithm, but it has one limitation, to ensure that we counted every possible particle that might end up in this cell we would either need to loop over the entire grid, which is highly expensive, or limit the maximum velocity of the particles to make the search radius finite, and hopefully only 1 pixel wide.

Obviously the second option is much better in our context, and in fact if we want to track particles with velocities so fast that they traverse the grid in 1 frame it’s actually cheaper to do lots of smaller steps with a 1 pixel neighborhood instead of counting all cells in one step, since \( 9 \sqrt{N} < N \) , where N is the number of cells, and the square root is because we need to apply the operation only a linear amount of times, instead of applying it for every cell. And the 9 is the number of neighbors. In 3D it would equivalently be \( 27 N^{1/3} < N \) which is even more efficient. I should note that more steps does not mean better quality, since the particles may combine together on the way. Also we can just do the tracking on CPU, and do it in a forward way, just loop over all particles and add them into the right cells, then we would not care about those things, but we obviously want to use that GPU power to our advantage.

Here is a pseudo-code glsl implementation of this algorithm:

//this cell position
ivec2 pos;
//values stored in the cell
vec2 velocity = vec2(0.), position = vec2(0.);
float mass = 0.;
//find and average the particles 
//that land in this cell after a time step dt
for(int x = -R; x <= R; x++) //only check the neighbors at radius R
  for(int y = -R; y <= R; y++)
  {
      //get the particle in this neighbor cell from the previous frame
      particle P = getParticle(pos + ivec2(x,y));
      //integrate the particle position
      P.X += P.V*dt;
      //check if the particle is inside of this cell
      if(inCell(P, pos))
      {
        mass += P.M; //add the particle mass to this cell
        position += P.X*P.M; //add the particle position weighted by mass
        velocity += P.V*P.M; //add the particle velocity weighted by mass(momentum)
      }
  } 

//normalize
if(mass > 0.0) //if not vacuum
{
  position /= mass; //center of mass
  velocity /= mass; //average velocity
}

Dividing particles

To combat the problem of a decreasing particle number we can try and divide each particle into M virtual particles with distributed positions that might end up in different cells and thus increase the total number of particles.

The radius of the distribution defines how likely is the particle is to multiply. If the all the virtual particles end up inside of a single cell the particle would not divide and stay essentially the same(if the average particle of the virtual particles is equal to the original particle). To make it properly conservative we just need to divide the mass of the particle into a number of equal chunks, and make sure the distribution average is zero.

On the picture above and in the code below we see an example for a 5 virtual particle distribution, the average of the distribution directions is zero and its radius is 0.1 pixels.

Here is the code, the only addition are the diffusion directions and a loop for each of the virtual particles.

//this cell position
ivec2 pos;
//values stored in the cell
vec2 velocity = vec2(0.), position = vec2(0.);
float mass = 0.;

//diffusion radius
float difR = 0.1;
//diffusion directions
vec2 difDir[5] = {vec2(0,0),vec2(1,0),vec2(-1,0),vec2(0,1),vec2(0,-1)};

//find and average the particles 
//that land in this cell after a time step dt
for(int x = -R; x <= R; x++) //only check the neighbors at radius R
  for(int y = -R; y <= R; y++)
  {
      //get the particle in this neighbor cell from the previous frame
      particle P = getParticle(pos + ivec2(x,y));
      //integrate the particle position
      P.X += P.V*dt;
      for(int i = 0; i < 5; i++) //divide particle into 5
      {
        particle difP = P;
        difP.X += difR*difDir[i]; //move particle in one of the diffusion directions
        difP.M /= 5.0; //divide mass into 5 particles

        //check if the divided particle is inside of this cell
        if(inCell(difP, pos))
        {
          mass += difP.M; //add the particle mass to this cell
          position += difP.X*difP.M; //add the particle position weighted by mass
          velocity += difP.V*difP.M; //add the particle velocity weighted by mass(momentum)
        }
      }
  } 

//normalize
if(mass > 0.0) //if not vacuum
{
  position /= mass; //center of mass
  velocity /= mass; //average velocity
}

You might ask if its possible to achieve a perfect particle number conservation, and not just total mass conservation. Actually to do that we just need to divide the mass into integer chunks that sum into the original mass, the virtual particles need not be the same, so for example one particle with mass 3 can divide into 2 particles with mass 2 and 1. Dividing into more than 2 virtual particles is pretty much the same, even though a bit more complicated.

Using particle distributions

The previous algorithm did solve the problem sufficiently well, but did have some shortcomings, like the additional loop which makes that algorithm far slower to execute and produces discontinuities in density that made it harder to implement fluid simulations. We can actually do better. Lets look at the limiting case of an infinite number of virtual particles, in such a case we end up with a continuous distribution of mass that we can call \( \rho(\vec{X}) \) (the distribution should also be centered on the particle position).

So to find the amount of mass (and its center) each particle deposits into this cell we need to integrate the distribution within the current cell bounds. (And that is where I got the name - reintegrating tracked particle distributions) Our equations are: \[ M = \int \int_{\Omega} \rho(\vec{X})d\vec{X} — \textrm{deposited mass} \] \[ \vec{C} = \frac{1}{M} \int \int_{\Omega} \vec{X}\rho(\vec{X})d\vec{X} — \textrm{deposited center of mass} \] Where \( \Omega \) is the cell region. But which distribution should we use? If we try to use a normal distribution we will end up with the problem of how to compute it, since there is no analytical solution for such an integral, and we would need to numerically integrate the distribution, which is not much better performance-wise than the previous algorithm. The simplest distribution that gives an analytical solution we can use is actually a uniform axis aligned box, for which the mass and the center of mass are trivial to compute analytically. We can also try a uniform circular distribution or a nonuniform circular distribution equal to \( 1 - |\vec{X_0} - \vec{X}|^2 \) for \( |\vec{X_0} - \vec{X}| \leqslant 1\) and equal to \( 0 \) for \( |\vec{X_0} - \vec{X}| > 1 \) where \( \vec{X_0} \) is the particle position, but let’s try out the simplest approach.


Diffusion radius 0.2	Diffusion radius 0.55

To find the mass and center of mass of the distribution within the bounds of the cell we only need to figure out the overlap box of the cell and the particle distribution. Its relative area will be the relative mass and its center is just the center of mass. It can be implemented like this:

//this cell position
ivec2 pos;
//values stored in the cell
vec2 velocity = vec2(0.), position = vec2(0.);
float mass = 0.;
//find and average the particles 
//that land in this cell after a time step dt
for(int x = -R; x <= R; x++) //only check the neighbors at radius R
  for(int y = -R; y <= R; y++)
  {
      //get the particle in this neighbor cell from the previous frame
      particle P = getParticle(pos + ivec2(x,y));
      //integrate the particle position
      P.X += P.V*dt;
      //find the overlap of the diffused particle distribution with this cell
      
      vec3 ovrlp = overlap(P.X, pos, diffusion_radius);
      float overlapRelativeArea = ovrlp.z;
      vec2 overlapCenterOfMass = ovrlp.xy;
      float overlapMass = overlapRelativeArea*P.M;

      mass += overlapMass; //add the overlap mass to this cell
      position += overlapCenterOfMass*overlapMass; //add the overlap center weighted by mass
      velocity += P.V*overlapMass; //add the particle velocity weighted by overlap mass(momentum)
  } 

//normalize
if(mass > 0.0) //if not vacuum
{
  position /= mass; //center of mass
  velocity /= mass; //average velocity
}

The axis alligned box overlap calculation is rather straightforward

vec3 overlap(vec2 x, vec2 p, float diffusion_radius)
{
    vec4 aabb0 = vec4(p - 0.5, p + 0.5); //cell box
    vec4 aabb1 = vec4(x - diffusion_radius, x + diffusion_radius); //particle box
    vec4 aabbX = vec4(max(aabb0.xy, aabb1.xy), min(aabb0.zw, aabb1.zw)); //overlap box
    vec2 center = 0.5*(aabbX.xy + aabbX.zw); //center of mass 
    vec2 size = max(aabbX.zw - aabbX.xy, 0.); //only positive
    float m = size.x*size.y/(4.0*diffusion_radius*diffusion_radius); //relative area
    //if any of the dimensions are 0 then the mass ratio is 0
    return vec3(center, m);
}

As you can see we don’t need to loop over virtual particles anymore since the solution is analytical, so the performance of this approach is the same as in the original cellular automaton particle tracker with the added benefit of giving much smoother results.

And here is a real time visualization of the reintegration process with diffusion radius 0.35

In fact this algorithm has some quite interesting properties, depending on the radius of the distribution the behaviour can change from particle-like to field-like as shown in the simulation below(you need to unpause it). The distribution diameter oscillates between 0.75 and 1.25:

We can see that for a diameter less than 1 pixel the behaviour tends to be particle-like and for a larger one it behaves more like usual advection with numerical diffusion.

Another interesting fact is that if we fix the particle positions to the cell centers and set the distribution radius to 0.5 so that the distribution is exactly as big as the cell we’ll get exactly forward Euler advection! Since we are technically just integrating the cells forward and depositing their contents.

Using the SPH formulation instead of finite differences to compute forces

Now what are we going to do with this algorithm? We can use the grid and compute finite difference gradients to compute forces. But we are actually losing the sub-cell distribution information - the cell centers of mass.

Since it can model particle systems we can try to adapt particle algorithms here, for example molecular dynamics(it would need exact particle count conservation), or maybe how about using smoothed particle hydrodynamics(SPH)? Actually this algorithm gives pretty much the perfect conditions for SPH, the particles are already uniformly distributed, around 1 particle per cell, and we can quite easily find the particle neighbors, since the grid itself is an acceleration structure! And that is pretty much exactly what was done in Paint streams or Everflow. With a large enough distribution radius (0.55-0.6) the natural diffusion is smoothing the particles so that we get away with a relatively small smoothing kernel, about 1.5 pixels wide, and we only need to compute the forces from the closest neighbors which makes it even more efficient.

To implement SPH we just need to integrate the new cell velocity using the reintegrated particle distributions.

\[ \vec{V}_ {i}^{t+1} = \vec{V}_ {i}^{t} +\Delta t \frac{\vec{F}_ {i}}{M_ {i}^{t}} — \textrm{updated velocity} \] Where \(\vec{F}_ {j}\) is the computed SPH force [1] computed for the particle in cell i.

\[ \vec{F}_ {i} =M_ {i}^{t} \sum_{j}^\textrm{neighbors} M_ {j}^{t} \left( \frac{P_{i}}{\rho_{i}^2} + \frac{P_{j}}{\rho_{j}^2} \right) \nabla_{i} W(\vec{X}_ {j}^{t} - \vec{X}_ {i}^{t}) \]

Where \(P_{i}\) is the pressure in cell i, and \(W(\vec{X})\) is the smoothing kernel. For the density \( \rho_{i} \) we can just use the mass of the cell divided by its volume, assuming the volume is 1 we can just place the mass. It’s ok to do so if the the distribution radius is big enough to smooth out the mass.

Pressure for each cell can just be computed using an equation of state. For a gas its simply just proportional to the cell density times the temperature, but let’s consider only constant temperatures:

\[ P_{i} = k \rho_{i} \] Where k is some proportionality constant.

For a fluid we can use the Cole equation of state \[ P_{i} = k \left( \left(\frac{\rho_{i}}{\rho_{0}} \right)^ \gamma - 1 \right) \]

Where \(\gamma\) is the adiabatic index (\(\gamma = 7.0\) for water), \(\rho_{0}\) is the reference fluid density.

In most of my simulations I just used the following pressure, which worked quite well in this setting. \[ P_{i} = k \rho_{i} (\rho_{i} - \rho_{0}) \]

Storing more properties inside a cell

There is nothing holding us from storing a more advanced description of the insides of the cell, we can store not just a particle - but an entire distribution, and we can also update its size depending on the variations of the centers of mass of other distributions that fell into this cell. We can also store the angular momentum of such a distribution to preserve vorticity in fluids for example, sadly averaging angular momentums of distribution parts is a bit complicated, and my experiments are not entirely stable and require angular momentum clamping suggesting that the way I added them was not exact (the colors show the curl value in the fluid). If I happen to successfully implement such a summation I will write a follow-up article, since perfect angular momentum tracking is really important for nice vortices.

Conclusions

This is a really cool algorithm, and I wanted to share it with other people the moment I had the first successful results, Wyatt has already used it for some really cool simulations, including multi-substance interactions. Thanks to the mass conserving quality of the advection it can be used to model self-gravitating gas too (the angular momentum there is total whack tho, but looks cool)

Other interesting use cases are:

Modelling boilling - very approximately, the equation of state is not that good there.
Fluid advection with a pressure computed using the Poisson equation.
Rocket Mach diamonds - modelling supersonic gas using a gas equation of state.
Slime molds - I should probably make a blog post on this too.
Life-like cells - ????. Another post needed, yeah.

Also here is a processing implementation of reintegration tracking (no SPH forces)

References

[1] Smoothed Particle Hydrodynamics