SSA Graph representation


#1

So I asked this on the irc channel, but maybe here more people will see it and provide more ideas, one suggestion is to use the Sea of Nodes representation, but actually I’m not sure how that will get implemented in Rust. I will explain the problem in some what elaborate detail, with the hope that someone can suggest good options for implementation in Rust.

The main goal is to build an SSA compiler for tensor operations. The main goal is to make something close to what LLVM is for compilers, but rather for Machine Learning and HPC for staticly defined programs. Note that the syntax is much more simplified than standard languages - each variable is a Tensor (large by assumption, but can be scalar). All operation are a list of predefined mathematical functions - e.g. Add, Div, Gemm, Trace, Tanh, Log etc… There are no control flows, like if/else, but there could be loops (but without control flows these are easy). Additionally, the compiler will have a built in autodiff capability, which is something very important and we care about in ML. I already have a working prototype in C++, but was wondering why not move this and work on it in Rust, as it is more portable, have already crates which make python integration a few macros away, and well I like Rust as well.

So getting back to the problem. In C++, my SSA representation is as a graph, where each node is owned by the graph, while it also has weak pointers to parents/children. This is of course doable in Rust as well. However, in the future, when we get to the optimization bit, what is needed is a backtracking operator. What this is since you defined many different optimizations independently, you Queue all of the nodes in some topological ordering, and you apply them in order. When an optimization changes the graph, you queue on the top everything that is affected and is not in the Queue (e.g. you have already passed it). The benefit of this is that you do not need to sweep trough the graph many times, and this can give significantly faster compile times. A possibly better explanation you can find here, although there the compiler in hand has control flows, thus is a more complicated problem.

To some extend I do not know if it is possible to achieve this in Rust, and please bare in mind I’m using Rust for a just a week and a half. The thing I lack to see how to implement this Queue which will have some reference to the nodes of the graph, yet you will be allow to mutate the graph at each step. So I was hoping if someone can give me some good suggestion on this or in general any advices on the topic.