Iterative Mojo Programming with %%moso

Lately, I've been interested in learning GPU programming using Mojo. I've followed Modular for a while and have always wanted to get some hands-on experience with GPU programming. It's just really cool to be able to write some custom kernels to speed up deep learning models.

I initially intended to follow the modular/mojo-gpu-puzzles, but the workflow didn't really work for me. I almost always use Jupyter Notebooks when I start learning something new. The interactive environment allows me to explore and understand concepts more deeply than any other method I've tried. Because Mojo is a compiled language, we can't natively execute some code, pause, and then run more code like we do with Python. We need a Mojo kernel to bridge this gap in Jupyter, but currently Mojo doesn't officially support one yet.

Answer.AI created mojokernel based on the Mojo REPL; it works fine on the CPU, but since the official Mojo REPL doesn't support the GPU yet, mojokernel unfortunately doesn't work on the GPU either.

Mojo does have its own `%%mojo` magic for Jupyter Notebooks, but it only compiles and runs a single cell. This means you have to constantly copy and paste all the previous cells you've developed into your new cell, which isn't convenient for fast experimenting.

So, I decided to write `%%moso` magic. It's primarily designed for SolveIt, a Jupyter-like environment where humans and AI can collaborate together. In SolveIt, %%moso collects only the cells from the beginning of the notebook up to the currently executing cell. In a standard Jupyter Notebook, because there is no native way to know the index of the currently executing cell, this magic has to compile and run all %%moso cells in the notebook.

This %%moso magic can also execute the code on a remote machine via SSH, which is very convenient to use on SolveIt where the instance currently doesn't have a GPU. We just need to specify the SSH host and the command to run the Mojo code. For example, we can configure it like this:

setup_moso(
    host="user@host", 
    cmd="cd ~/mojo-gpu-puzzles && /root/.pixi/bin/pixi run mojo run {path}"
)

I have already finished the first GPU puzzle and am quite satisfied with the results so far. I will include a detailed example of how it works in my next post!