Smuggling NumPy Into Go

A couple of days ago, writing software for Tillered, I ran into a situation where I needed to use a C++ library in my Go codebase. I’ll spare you the why. Go can’t really call C++, but it can call C, which turns out to be a fun door to walk through. The trick is writing a thin shim in C to act as a C++ translation layer that includes the library. This shim hides its classes behind opaque pointers, and exposes flat C functions wrapped in extern "C" so Go has something C-shaped to call. It works, though the build path is its own kind of cursed C++ lib -> C shim -> Go codebase, three languages deep before anything runs. But hey, anything for a PoC.

Around the same time, in my spare time (which I’m trying to force myself to have more of), I’d been contributing to the UQCS Discord bot, specifically hardening its Python Docker image, stripping it down as far as it’ll go without falling over. Which meant poking at CPython and its extension libraries to see what’s actually needed.

Then it hit me. CPython is also just a C library.

And that’s the thought that sent me down a rabbit hole of mild insanity.

Go Speaks C (The Easy Bit)

Running C from Go is quite easy, which is the first lie the language tells you. But the syntax is cursed and my LSP lights up like my holter monitor does when I go for a walk, but hey it works.

package main

/*
#include <stdio.h>

void helloworld() {
    printf("Hello, World!\n");
}
*/
import "C"

func main() {
    C.helloworld()
}

That comment isn’t a comment. cgo reads the C sitting directly above import "C" as a preamble, compiles it, and hands it to you under the C. namespace, which is why it’s C.helloworld() and not helloworld(). Leave a blank line between the comment and the import and the whole thing throws a fit.

The catch is that this needs CGO_ENABLED=1, so no static binary and a libc dependency is required on the host. Usually a fine trade as it means you have a smaller Go binary but you lose the portability to run it in scratch containers and/or embeded systems.

Smuggling In the Interpreter

If Go will link C, and CPython is just a C library, then Go will link CPython (my math skills at work). This just feels like an entirely cursed thing to be allowed to do. With the Python C API Reference Manual open and a Docker container to wrangle the dependencies, it turned out to be less code than it has any right to be:

package main

/*
#cgo pkg-config: python3-embed
#include <Python.h>
*/
import "C"

import (
	"runtime"
	"unsafe"
)

const python_source = "print('Hello, World!')"

func main() {
	// 1. Keep interpreter calls on one OS thread.
	runtime.LockOSThread()

	// 2. Initialise the interpreter.
	C.Py_Initialize()

	// 3. Allocate the source string in C memory.
	code := C.CString(python_source)
	defer C.free(unsafe.Pointer(code))

	// 4. Run the Python source.
	C.PyRun_SimpleString(code)

	// 5. Shut the interpreter down and release memory.
	C.Py_FinalizeEx()
}

Run it in the container and it prints Hello, World!. It’s pretty crazy that I can just spin up a python interpreter directly in a go codebase in a handful lines of code. I’ve stripped every safety check and error path to get it to the happy path, more sane version has a lot more guarding so it doesn’t wander into a bad state and take the process down with it.

Looking back at it, nothing interesting is actually happening here. Go handed Python a string and Python went off and talked to itself. No Go data went in, nothing came back, for no better words, it’s a party trick. To make it worth the jank, I want Go to call a real python library function and get a real result back. Which is where the difficulty stops being polite.

Thirty Lines to Say ones(5)

One thing first. Everything you just saw earlier; the thread lock, the init, and the finalise are identical every time for every use of the cpython in go. Because of this I’ve folded it into a wrapper called run(f func()) so the NumPy code isn’t mostly boilerplate. Nothing new in there, just the lock-init-finalise from the last section in a box. Read it as “do Python in here.”

So, goal 1. Can I call numpy.ones(5) and get [1, 1, 1, 1, 1] back? One function, one argument?

package main

/*
#cgo pkg-config: python3-embed
#include <Python.h>
*/
import "C"

import (
	"fmt"
	"unsafe"
)

func main() {
	run(func() {
		// 1. Import numpy.
		cname := C.CString("numpy")
		np := C.PyImport_ImportModule(cname)
		C.free(unsafe.Pointer(cname))
		defer C.Py_DecRef(np)

		// 2. Grab numpy.ones as a reference.
		fname := C.CString("ones")
		onesFn := C.PyObject_GetAttrString(np, fname)
		C.free(unsafe.Pointer(fname))
		defer C.Py_DecRef(onesFn)

		// 3. Build the argument tuple for ones(5).
		t := C.PyTuple_New(1)
		C.PyTuple_SetItem(t, 0, AsPyObj(5))

		// 4. Call numpy.ones(5).
		res := C.PyObject_CallObject(onesFn, t)
		defer C.Py_DecRef(res)
		C.Py_DecRef(t) // free the args tuple

		// 5. Convert the result to a string.
		s := C.PyObject_Str(res)
		val := C.GoString(C.PyUnicode_AsUTF8(s))
		C.Py_DecRef(s)

		// 6. Print it.
		fmt.Printf("np.ones(5) = %s\n", val)
	})
}

That’s roughly thirty lines, six steps, three manual Py_DecRefs and a free, to express the thing Python writes as np.ones(5). Import the module by C string. Reach in for the ones attribute by another C string. Build a one-element tuple to carry the argument, because PyObject_CallObject wants its args as a tuple. Hand it over, get a PyObject back, then convert “that” to a string by hand, because Go has no concept of a NumPy array. And every object along the way is a reference you’re now personally responsible for releasing. As you can see, I’m going to have to create more abstractions and boiler plate soon.

One real trap worth calling out: PyTuple_SetItem “steals” the reference to whatever you hand it, so you must not release that value yourself. That’s why there’s no decref for the 5. I would love to tell you I didn’t panic the progam numerous times when trying to get this to work.

And it works. Output is np.ones(5) = [1. 1. 1. 1. 1.], which is not really what I thought it would output. I’d forgotten that ones defaults to float64 due to being a math library, so everything’s a float, and the bracketed with no commas thing is just how NumPy renders an array. I believe this is the repr output of the object instead of an actual Python list. I’m sure thats fine.

Giving It an Actual Job

ones(5) was easy in comparison. A simple copy everything op to test if I can even do it. So the real test needed two things:

• Enough work that the boundary cost disappears into it.
• Data big enough that copying it would actually hurt.

Matrix multiplication (matmul) is the obvious pick as it is something “easy and small”, and I didnt want to bother doing something large like a Mandelbrot or Neural Network calc. For those compsci and stat people out there it’s O(n^3) compute over O(n^2) of data. It’s also one of the main usecases for NumPy. Matmul is also something I can just write a simple Go implementation for to test a comparison with as well.

The move that I can’t be bothered doing, but would be the “correct” thing to do is to setup a zero-copy. This is handing NumPy a pointer to memory that Go has already allocated and owns, instead of rebuilding the arrays and matrix values. According to the docs we can use PyArray_SimpleNewFromData to do this but it also means we need to do some messing around with the Golang runtime garbage collector to keep the memory alive and not freeing it when we are using it.

But I didn’t do this as my usecase is just small enough to not really have any benefit. I’m not going to go through everything I did, there is just alot of abstractions to make the system useable and you can check it out in the repo. I will however show you the simple matmul in main.

package main

import (
	"fmt"
	"strings"

	"github.com/lcox74/numbat/python"
)

// oneLine collapses numpy's multi-line array str onto a single line.
func oneLine(s string) string {
	rows := strings.Split(s, "\n")
	for i := range rows {
		rows[i] = strings.TrimSpace(rows[i])
	}

	return strings.Join(rows, " ")
}

func main() {
	python.Run(func() {
		np := python.Import("numpy")
		array := np.Attr("array")
		f64 := np.Attr("float64")

		// a = np.array([[1, 2], [3, 4]], dtype=np.float64)
		a := array.CallKw(
			[]any{[]any{[]any{1, 2}, []any{3, 4}}},
			map[string]any{"dtype": f64},
		)
		fmt.Printf("a = %s\n", oneLine(a.String()))

		// b = np.array([[5, 6], [7, 8]], dtype=np.float64)
		b := array.CallKw(
			[]any{[]any{[]any{5, 6}, []any{7, 8}}},
			map[string]any{"dtype": f64},
		)
		fmt.Printf("b = %s\n", oneLine(b.String()))

		// a @ b
		res := a.MatMul(b)
		fmt.Printf("a @ b = %s\n", oneLine(res.String()))
	})
}

Having this work was a good starting point as I was able to use the inbuilt Golang Benchmark tests to properly test this against a naive attempt of matmul in Go as well as the gonum which is a numpy but for Go.

The Race

Three contestants.

• Naive Go: A hand-written triple loop, single-threaded, cache-sane but no SIMD.
• NumPy: The whole smuggled stack, Go -> cgo -> CPython -> NumPy -> OpenBLAS.
• gonum Wired to that same OpenBLAS, so Go calls roughly the exact kernel NumPy does, just without an interpreter sitting in the middle.

I did add a throughput metric to the Go Benchmark of throughput, calculated in FLOP/s. Which for those who don’t know, its “Floating Point Operations Per Second”. So we arent looking at massive numbers, I’m just going to convert to GFLOP/s, bigger is better:

Naive Go is a flat line. 4.4 GFLOP/s from 64 all the way to 2048, completely indifferent to size. That’s one core running sequentially with no SIMD, so it never gets any faster no matter how much data you throw at it. The other two climb past 100 GFLOP/s, because OpenBLAS spreads the work across every core and vectorises it.

Here’s the thing that matters, and it’s the whole post. NumPy and gonum land pretty much in the same place. They both use openBLAS, which is an optmised linear algebra subsystem to do the actual calculations, and it has to be linked in at compile time, so the Go binary can’t be statically compiled for either option.

Which is the actual point. When the cursed CPython route and plain old Go both land on the same number, the language was never the thing doing the work. It was openBLAS the entire time, and NumPy and gonum were just two different ways of typing at it. NumPy isn’t fast because of Python, and gonum isn’t fast because of Go.

You could also tell the moment NumPy or gonum kicked in, because Docker Desktop would have absolte coniptions. At one point it reported that the container was using 852.36% of 800% CPU, on an M2 MacBook Air, melting it’s lifespan.

Going back to my initial question and thought of “what does the cgo boundary actually cost?” Well, the worse performance on the smaller matrices is the tell. Every call pays a fixed cost to cross from the Go code into the C library, and the smaller matrix benchmarks do far more iterations in the same benchmark time as the larger ones, so there are just more calls and more memory crossing over.

In the large matricies there are far fewer crossings, as each call stays in the C lib longer and gets to utilise its full performance. So the overhead is real, you can see it suppressing the small-N numbers, and it amortises away the moment each call does enough work to bury it.

Was It Worth It?

Eh, probably not. I can see a possible usecase in using an existing Python library if there is no Go version available. An example I guess would be tensorflow, but aparently that now has an official Go library which is just a C binding abstraction which is probably done alot better than what you could do with a cpython abstraction. I think it is probably safe to say that it isn’t really worth it, the level of abstractions that are required to make it easy to use in your native go codebase is pretty intense. That and trying to remember to manage memory from data allocations to strong references to be able to use the functions and arguments.

All in all, this was a fun experiment and a good test to see what some limitations of calling/managing C libraries are in Golang. I might play with it a bit more and do some more complicated examples, though I wouldn’t hold your breath. Check out the project at Numbat.