A pure go, high performance, linear algebra library
Machine Learning in Go
I really like the Go language. The tooling is great, the syntax is simple, and the features powerful. Recently, I wanted to try a new machine learning algorithm, and I checked how much machine learning was possible using pure Go. It turns out, not that much!
Gonum
The reference package in the language is called gonum. Broadly, this package has two layers. One is called "BLAS", for Basic Linear Algebra Subprogram. Simply speaking, this layer contains all the primitives required to perform linear algebra calculations. The second is called "LAPACK", for Linear Algebra Package. This one contains higher level functions that users will call directly. It relies on the blas layer to provide all the tools it needs.
Example of primitives are:
daxpy, for A.X + Y. This routine will scale a vector and add another.ddot, for dot product. That one will take care of doign two vectors dot products. It is used everywhere.
Gonum pure go implementation is optimized, but remains constrained by the language's limitations (garbage collection, etc.). If one wants to get better perofmances, Gonum makes it easy to rely on openblas, a C implementation of the blas layer. By doing so, however, one looses a lot of the nicest aspect of Go's tooling. I did not want to do that, and was looking for an alternative.
Serendipity
Well, it turns out that Go is able to execute assembly code natively! No need for any bindings of any kind, we stay in pure go territory, but we get access to the lowest form of interaction with the machine.
When Anthropic released their Fable model, I wondered if such a powerful model would be able to generate a new blas layer in assembly code that we then could call from a Go program. I was expecting very mixed results and a very different blog post to reflect on my findings. Turns out I was wrong.
In a couple days (fortunately, all before Anthropic shut down the model), the LLM has been able to plan the required architecture and code it.
Now. Anybody in this situation would say: "Well...I am sure the LLM generated something. But how good is it?"
Correctness, Performance and Interoperability
As the code plan was executed, I was following along and trying to understand the overall logic. However, I do not know assmebly at all, and the technique used in the code were very advanced. I was unable to assess if the resulting code was any good. There were two components to it: the performace aspect (thw whole point of the library) and the correctness aspect (who cares if it is fast if the result is wrong!).
Performance
Regarding the performance aspect, the LLM did a good job at running the same calculations with gonum(with and without the C blas layer) and with goblas. Given those benchmark it seems like goblas is significantly faster than gonum's pure go version, and almost as fast as the C implementation. I was sure to make the benchmarks as easily reproducible as possible.
A blas implementation has three layers. Conceptually, they build on each other (L1 is vector on vector, L2 is matrix on vector, L3 is matrix on matrix). However, L1 and L2 are memory bound and no amount of assembly would be able to speed it up. L3 is different, and this is where all the optimization will take place. It means that the L3 routines are not reusing the ones from L1 and L2, because they need a different logic for optimization.
Correctness
Regarding the correctness, the idea is to execute the same calculations with a library you trust (here, gonum) and compare the results with the ones computed by goblas. Again, those tests were made during the implementation to make sure we were on the right track. However, afterwards, I wanted an additional verification. I was looking for a tool that would give me some inputs and the expected outputs, so that I could compare.
Such a tool did not exist. So...I made one! It's so easy nowadays! It is called goblast, and is available for free on Github! I think this is quite cool. This tool take the same ideas but allow anyone who would want to test their blas layer correctness to do so. The tool will generate a list of input files representing different gotchas and edge cases that could trump any of those subroutines. The expected output is generated by gonum. Then it's up to the blas layer's maintainer to run their routines on the input files and export the results in text files as well. Once done, the tool is able to compare the results and create a report to let you know which part of the tool can be improved.
Interoperability
So far we have talked about the blas layer. What about the lapack one? Gonum has a vey nice API that allows you to switch one blas layer for another. Since the goal is not to reinvent the wheel, goblas is fully cmopatible! It takes literally one line to switch to goblas, and every other functions will work out of the box.
Float64 vs Float32
This project really reminded me that computers are bad at computing float numbers. No operation will ever be perfect. Small errors compound, and any correctness verification need to take this into account.
From what I understood, any scientific calculation will be done using float64 to minimize those errors. Neural network training apparently handles those nicely, and since float32 numbers on a 64 bit machine allows for even more optimizations, neural netwrok libraries will use float32. goblas has both available.
Conclusion
This project is strange. LLMs hallucinate. It is a fact. Their output needs to be verified. Which usually means that someone who knows how to do the task should supervise the result. Here, I was more like an observer. I have asked all the questions and clarifications I could, and an LLM can be a fantastic teacher. However, I am not able to tell you if any of the assembly code is correct. All I was able to do was to consider the library as a black blox and ensure that the right outputs are provided given known inputs.
On the other hand, I think this project was particularly well suited for an LLM. The routines are well scoped, small, self contained, and easily testable. The LLM had a very clear way to ensure they were on the right track. It does not prevent mistakes, but helps.
I think there are very good applications to this new library. Having the ability to train and run machine learning models in pure go can be very useful. Overall, I hope that enough people will find it interesting to attract the attention of someone with enough skills to assess the quality of what was done.
Finally, it is understood that this library is meant to be used on a CPU. The goal is not too train your own language model, but there are tons of other algorithms to try out there!