2026-08-04

Exploring the Mojo Programming Language's RISC-V Assembly

Exploring the Mojo Programming Language's RISC-V Assembly thumbnail

A nerdy, exploratory dive into the assembly output from the Mojo compiler when targeting RISC-V CPUs, comparing scalar and vector codegen, and tinkering with Mojo's SIMD metaprogramming features.

In this article, we'll dive into the assembly output from the Mojo compiler, when targeting RISC-V CPUs, and tinker a bit with Mojo's metaprogramming features.

This post is a spiritual spin-off of sorts, to my blog post series on writing RISC-V-specific Linux kernel modules in Rust. For this spin-off, we'll be moving out of both the Rust and Linux context, and setting our sights on the Mojo programming language, and how it handles compiling assembly code for RISC-V architectures.

For reasons explained as we go along, this will be purely focussed on reading and analyzing the generated code, rather than actually executing or benchmarking Mojo binaries - so it'll get really nerdy, without a ton of bells and whistles, and none of those nice dopamine rushes when things compile. Luckily, the fun discoveries we'll make along the way make up for that, so if reading some assembly code sounds like a good time to you, let's get right into it.

As always, you can find the source code, and some very, very rough notes for the post on my GitHub repo for the project: https://github.com/grammeaway/mojo-risc-v-assembly-codegen-project

Table of Contents

Technologies and Terminology

Just to catch everyone up to speed, we'll quickly go over some of the main terminology and tech at play in this exercise:

  • Mojo: A new programming language, developed by the team at MODULAR, the latest entrepreneurship venture from one of the compiler GOATs; Mr. Chris Lattner. Developed as a superset of Python, Mojo is a compiled language, offering the sort of performance one would expect from a systems language like Rust or C++, but leveraging a, to many developers, more friendly and familiar Python-like syntax. The language has rich support for metaprogramming, to make it well-suited for GPU programming, and it binds easily to Python programs. I highly recommend Mr. Lattner's guest episode on the Pragmatic Engineer podcast (From Swift to Mojo and high-performance AI Engineering with Chris Lattner), to learn more about why Mojo came to be.
  • LLVM: A compiler backend used by a wide array of programming languages, with some notable names being Rust and, of course, Mojo. LLVM's development was, just like Mojo's, led by Chris Lattner.
  • Metaprogramming: A feature of some programming languages, which allows the developer to programmatically tell the compiler how to handle specific pieces of code. Allows for very fine-grained, case-by-case performance optimizations.
  • RISC-V: A CPU instruction set architecture. Sets itself apart from the more commonly known x86 and ARM architectures, by being a fully open standard - no need to pay any licenses to produce RISC-V chips. It has a lot of market share in IoT and academia, with constantly growing efforts towards e.g. data-center and personal computer use cases. You can read my full reasoning for wanting to dive into RISC-V in my original blog post on the matter.
  • Assembly: Code in more or less the closest-to-machine-code format we can produce. Assembly code contains raw CPU instructions, and very few of the luxuries that we're used to from even more low-level programming languages like C. For the same reason, it allows you the highest level of control when doing performance optimizations, since it allows you to manipulate memory and CPU registers exactly how you want to - the tough part then becomes to do this better than a compiler.
  • Kernel: In my past blog posts, we've been dealing with the term kernel in the context of Linux, where the kernel is the core of the operating system. In this case, we're using the term in the same way as it's typically used in data science, where it refers to a small, focused computational routine that does one specific math operation on an array of numbers.
  • SIMD: "Single instruction, multiple data", a CPU design feature allowing a chip to perform a math task on multiple pieces of data, at the same time. So multiple inputs are processed, but in one single CPU instruction, making it a powerful optimization technique. Mitchell Hashimoto (co-founder of HashiCorp, creator of ghostty) recently released a fantastic blog post on the subject, which I highly recommend giving a read: https://mitchellh.com/writing/everyone-should-know-simd

A Disclaimer of Sorts

Since we'll be analyzing, and in some cases attempting to "evaluate" the output assembly code from Mojo, I feel the need to quickly mention that I am by no means an expert in compilers, nor in RISC-V, nor in assembly code. These posts are almost always documenting some learning journey of mine, and any observations are nothing more than that - observations. So take it with a grain of salt, if I start pointing at things that seem suboptimal in the assembly output.

Mojo is under very active development, by a team of people who have forgotten more about compilers than I'll most likely ever learn. Moreover, the language is still on version 0.x.y (I'll be working with 0.26.2.0 in this post), so treating it as a fully completed language would be unfair, because it isn't.

So this is very much not an evaluation of the quality of Mojo, and more of a convenient excuse to get to play around with Mojo (which I'm very excited about, due to our backend at Tryp.com being primarily Python), while simultaneously learning about RISC-V assembly code.

Step by Step

Step 1: Setting up your environment

Since Mojo is a Python superset, you'll want to have a functioning Python environment set up. This guide will presume that you're managing dependencies through uv, which I can only recommend that you take for a spin if you're not familiar. If you prefer to avoid uv, a setup using just pip and any sort of virtual environment you prefer, should also do the trick.

It's worth noting that I went through the steps on an x86 Arch Linux machine, so your mileage may vary in terms of some of the commands.

If you're going with uv, adding Mojo to a project is as simple as:

uv add mojo

To verify a successful install, simply run:

uv run mojo --version

When I went through these steps, it was with version 0.26.2.0 of Mojo.

For your project dir, go ahead and make a helloworld directory, an asm directory for our assembly output, and a kernels directory for our kernel implementations:

mkdir helloworld kernels asm

And that's all we'll need in terms of our environment setup. We'll add a quick Makefile later, to speed up iteration, but first we need to figure out what to even put into that Makefile.

Step 2: A Mojo "Hello World!"

In your helloworld directory, go ahead and create a file named hello.mojo, and add the following code:

fn main():
    print("Hello, Mojo!")

And that's actually all a Mojo "Hello World" application needs! Without getting too deep into Mojo language syntax itself, note that Mojo supports defining functions using both the Mojo-native fn keyword, and the more Pythonic def keyword. When defining functions with fn, the Mojo compiler expects you to write the code in a statically typed manner, where each variable and constant gets assigned a type. When defined with def, the Mojo compiler lets you write code in a dynamically typed manner, like you'd be able to do in Python. We'll be opting for fn throughout this entire exercise.

With that written, we're ready to run our little program, using:

~ # uv run mojo helloworld/hello.mojo
Hello, Mojo!

And that's our "Hello World" done! Now on to where things get really interesting, and where some of the really cool features of Mojo's build system start coming through.

Step 3: Outputting RISC-V Assembly

Mojo has an incredibly richly customizable build system, allowing you to specify not just CPU architectures, but specific CPUs as your build target. First things first, we need to decide on our target triple (LLVM target triple format), and our target CPU.

To see the list of supported targets, run:

~ # uv run mojo build --print-supported-targets
Registered Targets:
  xtensa - Xtensa 32
  x86-64 - 64-bit X86: EM64T and AMD64
  x86 - 32-bit X86: Pentium-Pro and above
  riscv64be - 64-bit big endian RISC-V
  riscv32be - 32-bit big endian RISC-V
  riscv64 - 64-bit RISC-V
  riscv32 - 32-bit RISC-V
  nvptx64 - NVIDIA PTX 64-bit
  nvptx - NVIDIA PTX 32-bit
  hexagon - Hexagon
  amdgcn - AMD GCN GPUs
  r600 - AMD GPUs HD2XXX-HD6XXX
  aarch64_32 - AArch64 (little endian ILP32)
  aarch64_be - AArch64 (big endian)
  aarch64 - AArch64 (little endian)
  arm64_32 - ARM64 (little endian ILP32)
  arm64 - ARM64 (little endian)

For our purpose, we'll be going with 64-bit RISC-V. Since I'm on a Linux machine, my target triple ends up being riscv64-unknown-linux-gnu.

Next up is finding our target CPU, which Mojo also has a convenient list command for. The list command requires supplying the target triple as an argument, like so:

~ # uv run mojo build --print-supported-cpus --target-triple=riscv64-unknown-linux-gnu
Available CPUs for target riscv64-unknown-linux-gnu:
  an-erbium
  andes-ax25
  andes-ax45
  andes-ax45mpv
  andes-nx45
  et-soc1
  generic-rv64
  mips-p8700
  rocket-rv64
  sifive-p450
  sifive-p470
  sifive-p550
  sifive-p670
  sifive-p870
  sifive-s21
  sifive-s51
  sifive-s54
  sifive-s76
  sifive-u54
  sifive-u74
  sifive-x280
  sifive-x390
  spacemit-a100
  spacemit-x60
  spacemit-x100
  syntacore-scr3-rv64
  syntacore-scr4-rv64
  syntacore-scr5-rv64
  syntacore-scr7
  tt-ascalon-x
  veyron-v1
  xiangshan-kunminghu
  xiangshan-nanhu
  xt-c910v2
  xt-c920v2

Quite a selection, and for this little project, we'll be building for two of them: generic-rv64 and sifive-x280.

The reasoning behind working with two different target CPUs will become much clearer once we start reading the assembly code, but the primary reasoning pertains to the vector support of the CPU targets. The generic-rv64 has none (it actually doesn't have much of anything - more on that later), whereas sifive-x280 has native support for vector operations. This of course makes them really interesting CPU targets to compare, when talking about Assembly codegen.

With our build target configurations sorted out, the only remaining configuration to keep in mind is the --emit flag, which lets us request specific output types from the Mojo compiler. In this case, we'll be telling it to emit assembly code, by passing asm as the --emit argument.

So our final build command for our "Hello World", with generic-rv64 as the CPU target for this example, ends up being:

uv run mojo build helloworld/hello.mojo --target-triple=riscv64-unknown-linux-gnu --target-cpu=generic-rv64 --emit=asm -o asm/hello.s

This should produce a file named hello.s in the asm directory, full of RISC-V assembly code. We won't be analysing the "Hello World" assembly, this was mainly to end-to-end validate our Mojo build pipeline.

Let's get a proper kernel written, to get some more interesting assembly to analyse.

A Sidebar on Why We Can't Run the Program

In my original plans for this exercise, we were going to run a RISC-V binary output through QEMU, similar to how we were doing things in my other RISC-V posts. However, this got axed due to a pretty significant absence in the Mojo ecosystem: There is no RISC-V runtime for Mojo (or at the very least not one I was able to get my hands on). So while Mojo can output binaries made for RISC-V CPUs, there's no Mojo runtime library available for those binaries to interact with, making any operations that rely on the Mojo runtime unable to run.

This is luckily not a huge blocker for what we're trying to do in this exercise, and it does make a lot of sense for Mojo to currently be prioritizing support for more widely used instruction set architectures, especially due to Modular's heavy focus on AI workload hosting (which would primarily be happening in data centers, where x86 and ARM still reign supreme).

So as mentioned in the introduction, we won't be executing or benchmarking any of our kernels this time around. That's all good though - let's get to reading some assembly instead.

Step 4: A quantized int8 dot product kernel

For this post, we'll be building a small, but classic kernel: A quantized int8 dot product kernel, i.e., a kernel which produces the sum of element-wise products of two 8-bit integer arrays, accumulated into a 32-bit integer.

It scratches most of our itches, by being small enough that we won't be completely drowned in assembly output, while still allowing us to play around with optimizations through metaprogramming. It's also highly vectorizable, which is one of the main properties of Mojo codegen that we're trying to exercise.

In your kernels directory, create a file named dot_int8.mojo, and add the following code:

fn dot_int8(
   a: UnsafePointer[Int8, origin=...],
   b: UnsafePointer[Int8, origin=...],
   n: Int,
) -> Int32:
   var acc: Int32 = 0
   for i in range(n):
      acc += Int32(a[i]) * Int32(b[i])
   return acc


fn main():
   var a = List[Int8](length=8, fill=Int8(1))
   var b = List[Int8](length=8, fill=Int8(2))
   var result = dot_int8(a.unsafe_ptr(), b.unsafe_ptr(), 8)
   print(result)

Note: The inclusion of the main function is necessary for Mojo's compiler to not just ignore our kernel function - if it isn't invoked in the main execution loop, it won't be compiled. There might be ways around this, I just didn't happen to find one. It'll lead to a fair bit more output assembly code than we'd ideally want, but we'll manage.

Since we have the main function declared, you can run the kernel to validate the correctness of the code:

~ # uv run mojo kernels/dot_int8.mojo
16

This initial implementation of ours is quite naïve by design: It simply handles producing the dot product of the arrays in a loop, meaning that in a real ML context, with huge arrays of data, it would scale awfully performance-wise. But that's exactly what makes it an interesting baseline to start from - let's see what kind of assembly code we get from this kernel.

Run the builds towards the two CPU targets:

~ # uv run mojo build kernels/dot_int8.mojo --target-triple=riscv64-unknown-linux-gnu --target-cpu=generic-rv64 --emit=asm -o asm/dot_int8.generic.s

~ # uv run mojo build kernels/dot_int8.mojo --target-triple=riscv64-unknown-linux-gnu --target-cpu=sifive-x280 --emit=asm -o asm/dot_int8.x280.s

At this point, it makes sense to set up a Makefile for making building the contents of the kernels dir towards the two target CPUs a bit faster - you can whip one up yourself, or borrow mine from the project GitHub repository.

And this should produce two new assembly files in your asm directory. Let's dive into them, and see what we learn.

Step 5: Analyzing assembly from our naïve kernel

We'll be having a look at a few points of interest in both the two assembly outputs, starting with the generic-rv64 output, and then moving on to the sifive-x280 output. I want to once again stress that this is all very new territory for me, so do take some time to do your own deep-dives into the assembly code - you're pretty likely to find some interesting stuff that I missed!

Step 5.1: Analysing the generic-rv64 assembly

If you open up asm/dot_int8.generic.s in your preferred text editor, you'll be greeted by just shy of 1000 lines of assembly (912 in my case). This is pretty overwhelming, but note that most of it actually comes from our harness (i.e. the main function implementation), rather than our kernel code itself.

The first small point of interest, is already on line 2:

        .attribute     5, "rv64i2p1"

.attribute 5 declares the target architecture, and the enabled ISA (instruction set architecture) extensions. For our generic-rv64 CPU target, we have zero enabled ISA extensions - living up to the "generic" moniker.

One small annoyance from an analysis perspective, is that our function does not have a dedicated function block in the assembly code - the compiler made the choice of inlining the function into the main function, since our kernel function doesn't get called anywhere else, and is quite small. Perfectly reasonable optimization for the compiler to make, but it does mean that we have to dig a bit to find our kernel code.

If your output matches mine, our kernel code starts showing up at line 103, in two main blocks labeled as .LBB2_6: and .LBB2_7:. I've given them a few explanatory comments, to help understand what's going on, line for line:

.LBB2_6:
      li      s2, 0            # acc = 0
      addi    s5, s1, 8        # end pointer = b + 8
      mv      s6, s0           # cursor_a = a
      mv      s7, s1           # cursor_b = b
.LBB2_7:
      lb      a1, 0(s6)        # load int8 from *cursor_a (sign-extended to 64)
      lb      a0, 0(s7)        # load int8 from *cursor_b (sign-extended to 64)
      call    __muldi3         # 64-bit multiply via library call
      addw    s2, a0, s2       # acc += result (32-bit add with sign extension)
      addi    s7, s7, 1        # cursor_b++
      addi    s6, s6, 1        # cursor_a++
      bne     s7, s5, .LBB2_7  # if cursor_b != end, loop

Note: .LBB2_7 has been shortened slightly, you should have 3 additional lines in there that we'll ignore for this exercise.

If you, just like me, aren't exactly used to reading assembly, this'll be a lot to take in already. I found it to be a pretty fun exercise to read through how looping works in this assembly code - not relevant to what we're analyzing, but there's a greater appreciation for the humble for loop to be found by seeing just how many operations it abstracts away.

In this little bit of generic-rv64 assembly, there are a couple of interesting facts emerging already:

  • call __muldi3: The assembly code executes a library call to do multiplication. Remember the previously highlighted .attribute 5 line? On the long list of ISA extensions not active for this CPU, you'll find m: multiplication. So this CPU has no native support for running multiplication, meaning that it has to rely on a software library to do so. Not a knock on any parts of the stack here - the compiler is working with what it has available, which in this case is very little. Doing a full library call is the only viable option, but also incredibly inefficient for multiplication purposes.
  • The lack of autovectorization: Not shocking, as this particular CPU target also doesn't have the v extension enabled, meaning that it doesn't support vector operations in the first place. But it's something worth noting for our next bit of assembly, since that will be built for a Vector-enabled CPU target. What we'll be hoping to see there, is that this perfectly scalar loop will be vectorized by the compiler by default, to optimize performance.

So this little bit of assembly analysis has already shown us examples of the compiler working around the limitations of the target CPU, and given us a chance to look at some somewhat non-threatening assembly code. Let's take a look at the exact same kernel, when compiled for a much more advanced CPU target.

Step 5.2: Analysing the sifive-x280 assembly

Let's have a look at asm/dot_int8.x280.s instead. You'll notice that this assembly code is marginally shorter, clocking in at 892 lines of code.

Starting out with the exact same point of interest on line 2, let's have a look at the ISA extensions available to us on this CPU:

        .attribute     5, "rv64i2p1_m2p0_a2p1_f2p2_d2p2_c2p0_v1p0_zicsr2p0_zifencei2p0_zmmul1p0_zaamo1p0_zalrsc1p0_zfh1p0_zfhmin1p0_zca1p0_zcd1p0_zba1p0_zbb1p0_zve32f1p0_zve32x1p0_zve64d1p0_zve64f1p0_zve64x1p0_zvfh1p0_zvfhmin1p0_zvl128b1p0_zvl256b1p0_zvl32b1p0_zvl512b1p0_zvl64b1p0"

Now that's a CPU we can actually have the compiler optimize around. Once again, this is barely legible to me, but just like before, we're targeting a 64-bit RISC-V machine. In the chain of letters that follows, separated by _ characters, you'll notice that we now have luxuries like M (multiply), A (atomics), F+D (float), C compressed, and V (vector) enabled. So much more potential for optimizations in our codegen out of the box, and a CPU much more representative of what actual RISC-V silicon looks like.

So how will this change our kernel output?

You'll find the start to our kernel around line 99, in code blocks named the same as on the generic CPU target:

.LBB2_6:
     li    s0, 0         # acc = 0
     addi  a0, s1, 8     # end_pointer = b + 8
     mv    a1, s2        # cursor_a = a
     mv    a2, s1        # cursor_b = b
.LBB2_7:
     lb    a3, 0(a1)     # load byte from *cursor_a → a3 (sign-ext to 64)
     addi  a1, a1, 1     # cursor_a++
     lb    a4, 0(a2)     # load byte from *cursor_b → a4
     addi  a2, a2, 1     # cursor_b++
     mul   a3, a4, a3    # a3 = a4 * a3 - real hardware multiply!
     addw  s0, s0, a3    # acc = acc + a3 (32-bit add, sign-extended)
     bne   a2, a0, .LBB2_7  # if cursor_b != end, loop back

We once again have a couple of interesting findings, both relating back to our analysis of the generic CPU output:

  • mul a3, a4, a3: We got a proper hardware-level multiply this time! Since this CPU supports it, the compiler opted for the much faster option of having the CPU perform the multiply for us, rather than an external software library.
  • Still no autovectorization: This one is very interesting, because we know that the CPU supports vector operations, and that the code is a prime candidate for optimization through vectorization.

Especially the second point, is worth noting when talking about the RISC-V assembly codegen from Mojo, and something we can note down as a finding of sorts:

Scalar Mojo loops don't auto-vectorize on RISC-V (sifive-x280, Mojo 0.26.2.0). We built a dot_int8 kernel, targeting both generic-rv64 and sifive-x280.

generic-rv64 emits a 7-instruction scalar loop with software multiply (__muldi3). sifive-x280, despite its attribute string including V (Risc-V Vector extension 1.0) plus Zvl512b (512-bit vector length) and the full Zve* family of subset extensions, emits a scalar loop using hardware mul. No v* instructions appear anywhere in main. The autovectorizer chose not to engage.

Note: This isn't really any sort of flack against the LLVM compiler. Auto-vectorization is a somewhat famously difficult domain within compiler design, and there aren't, as far as I'm aware, any compiler backends or compilers solving it perfectly. All the more reason to explicitly declare SIMD usage, which sets us up really nicely for the next steps, where we'll use Mojo's metaprogramming features to force the compiler into leveraging vectorization.

Step 6: A SIMD-optimized kernel

In the kernels directory, make a new file named dot_int8_simd.mojo. To exercise the Mojo compiler a bit further, we'll be writing two different implementations in this kernel: One where we explicitly run a loop, and one where we leverage the strengths of the Mojo language's SIMD primitives, i.e., one of the strong meta programming features of the language.

Write the following kernel implementation to the file:

# Pure SIMD primitive: no loops, no pointers, just vector arithmetic.
fn dot_int8_simd_primitive(
   a: SIMD[DType.int8, 16],
   b: SIMD[DType.int8, 16],
) -> Int32:
   # Widen to int32 element-wise (sixteen int8 -> sixteen int32)
   var a32 = a.cast[DType.int32]()
   var b32 = b.cast[DType.int32]()
   # Element-wise multiply, then horizontal sum across all 16 lanes
   var products = a32 * b32
   return products.reduce_add()


# Looped version: iterates over a buffer of int8s in chunks of 16.
# Assumes n is a multiple of 16 to simplify the loop body
fn dot_int8_simd_loop(
   a: UnsafePointer[Int8, origin=...],
   b: UnsafePointer[Int8, origin=...],
   n: Int,
) -> Int32:
   comptime WIDTH = 16
   var acc = SIMD[DType.int32, WIDTH](0)
   var i = 0
   while i < n:
      # Load 16 contiguous bytes from each pointer
      var va = (a + i).load[width=WIDTH]()
      var vb = (b + i).load[width=WIDTH]()
      # Widen + multiply + accumulate into the vector accumulator
      acc += va.cast[DType.int32]() * vb.cast[DType.int32]()
      i += WIDTH
   return acc.reduce_add()


fn main():
   # Anchor for codegen — same harness pattern as the scalar kernel.
   var a = List[Int8](length=32, fill=Int8(1))
   var b = List[Int8](length=32, fill=Int8(2))

   # Call the primitive — needs us to construct SIMD values from somewhere
   var va = (a.unsafe_ptr()).load[width=16]()
   var vb = (b.unsafe_ptr()).load[width=16]()
   var result_primitive = dot_int8_simd_primitive(va, vb)

   # Call the looped version
   var result_loop = dot_int8_simd_loop(a.unsafe_ptr(), b.unsafe_ptr(), 32)

   # Force observability so the compiler can't elide the calls
   print(result_primitive)
   print(result_loop)

In the primitive implementation, we leverage the built-in SIMD structure, a part of Mojo's metaprogramming features, to perform the dot product math in very few lines of code. In the loop version, we go for a bit more elaborate route, and loop over the input arrays, accumulating them into the accumulator vector. Just like before, we'll be including a main function, to ensure that our kernel implementations end up getting compiled.

Run your build Makefile flow again, and verify that you end up with new files in your asm directory - if you're using my Makefile, you should end up with a file named dot_int8_simd.generic.s, and one named dot_int8_simd.x280.s.

For this SIMD-based implementation, we'll start off analyzing the output from the build targeting the sifive-x280 CPU.

Step 6.1: Analysing the sifive-x280 loop assembly

At this point, we're dealing with quite a lot of Assembly output, especially since we have two kernels in one but let's break it down into smaller points of interest, and try to draw some learnings about both RISC-V assembly itself, as well as how the Mojo compiler handles specific code snippets.

At around line 106, in the blocks LBB2_6 through block LBB2_7, you'll find the main assembly code for the dot_int8_simd_loop() function. The LBB2_6 block is our setup block, initializing the things our loop implementation will need:

.LBB2_6:
    csrr     a0, vlenb                       # Read the hardware's vector register length
    vsetivli zero, 16, e32, m1, ta, ma       # Configure the vector unit: 16 lanes of 32-bit integers, using one full vector register, in tail-agnostic and mask-agnostic mode
    vle8.v v8, (s4)
    vmv.v.i v11, 0                           # initialize accumulator to zero
    ...
    vs1r.v v8, (a0)                          # spill to stack
    addi     a0, sp, 64
    vle8.v v8, (s1)
    vs1r.v v8, (a0)
    li    a0, 0                              # loop counter = 0

Unlike our previous sifive-x280 output, this is abundant with vector operations already, all denoted by the instruction starting with a v. As expected, using the SIMD metaprogramming features has led to Mojo's compiler playing to the strengths of the hardware, leveraging the V extension on the CPU instruction set.

A fun thing that anyone who followed my post series about writing RISC-V-specific kernel modules will notice, is the usage of the csrr instruction, which we back then leveraged to read out hardware data like the CPU's cycle count. This time, it's being used to fetch out the vector register length of this particular RISC-V CPU, and load that into the register a0.

Another point of interest is the vsetivli instruction, where we configure our vector unit. In this case, the compiler passes the inputs zero, 16, e32, m1, ta, ma, which tells the CPU that for the next batch of vector instructions, we'll be operating on 16 lanes of 32-bit integers, using one full vector register, in tail-agnostic and mask-agnostic mode.

An interesting thing worth noting, is that this block also prepares for our SIMD-primitive kernel: The vle8.v + vs1r.v pairs are for the primitive kernel's inputs, and get spilled to the stack to survive the end of the loop's execution.

Moving on to the slightly more action-heavy part, in the form of the loop itself.

.LBB2_7:
    add       a1, s4, a0             # cursor_a = base_a + offset
    add       a0, a0, s1             # cursor_b = base_b + offset
    vsetvli zero, zero, e16, mf2, ta, ma # reconfigure: 16 lanes of int16, half-register
    vle8.v v8, (a1)                  # load 16 int8s from a
    andi      a1, s0, 1              # a1 = s0 & 1 (loop control)
    vle8.v v9, (a0)                  # load 16 int8s from b
    li      a0, 16                   # next offset = 16
    li      s0, 0                    # zero out s0 for next iteration's control
    vsext.vf2 v10, v8                # widen int8 → int16 in v10
    vsext.vf2 v8, v9                 # widen int8 → int16 in v8
    vwmacc.vv v11, v8, v10           # v11 += v8 * v10 (widened to int32)
    bnez      a1, .LBB2_7            # loop if a1 != 0

Once again we've got a lot to unpack, but skimming through it does show a flow matching what we wrote in our Mojo kernel - step by step, we:

  1. Compute the offset into a and b (cursor_a, cursor_b).
  2. Load 16 int8s from each pointer.
  3. Widen them to int32 and multiply.
  4. Accumulate the products into the running sum.
  5. Advance to the next iteration.

But naturally, our point of interest in this post is more geared towards how we do it, and the choices made by the Mojo compiler. The main points of interest here are two important steps: A reconfiguration of our vector unit, and the powerful vwmacc.vv operation.

The reconfiguration once again calls the vsetvli operation, essentially re-doing the original configuration, by instead requesting 16 lanes of 16-bit integers (as opposed to the previous 32-bit integer lanes), using half a vector register (mf2 tells it to allocate half a register, since 16 lanes of 16-bit integers is half a register's worth of data on this CPU).

Why this re-configuration? Because of the inputs expected by the ABI (application binary interface) of the vwmacc.vv (widening multiply-accumulate) operation. Its inputs must be at the smaller element width, and it widens to write to a destination at 2x the width. To use vwmacc.vv to produce int32 outputs (as our accumulator v11 requires), we need to feed it int16 inputs.

This is also why we see the two vsext.vf2 calls: these are "sign-extend by a factor of 2" instructions, converting the freshly-loaded 8-bit integer values into 16-bit integer values, which can then subsequently be widened into our target 32-bit integer results by the vwmacc.vv call.

With these pre-emptive steps done, our data is ready for the main action of the kernel: The vwmacc.vv call.

vwmacc.vv v11, v8, v10

This single instruction performs the following actions:

  1. Multiplies each lane of 16-bit integers stored in v8, with the corresponding lane in v10.
  2. Widens each product to being a 32-bit integer.
  3. Adds each widened product to the corresponding lane in our 32-bit integer accumulator (v11).
  4. And finally, it writes the result back to the v11 register.

So after the prepwork, this absolutely powerhouse of an instruction goes ahead and gives us 16 multiplications, 16 widenings, and 16 additions in one single instruction call - pretty much as much of a dictionary-definition-worthy example of SIMD as we could have asked for. We'll be contrasting this against the generic RISC-V CPU codegen in a few, to show just how incredibly powerful of a feature proper vectorization is.

What we've gotten to see here, is that one of Mojo's SIMD abstractions, when going through LLVM's RISC-V backend, produces powerful, well-optimized vectorization Assembly code, highlighting the strength of the entire end-to-end codegen pipeline.

Let's try having a look at our primitive-based kernel.

Step 6.2: Analysing the sifive-x280 primitive assembly

Recall that our primitive-based kernel is defined in the function dot_int8_simd_primitive(). I recommend taking a second to re-read the Mojo source code, just to refresh how it differs implementation-wise from our loop-based kernel - note that they both leverage Mojo's metaprogramming features, just in different manners!

In the same codegen output as in 6.1, we'll be looking for the block named .Lpcrel_hi3, which in my case starts at around line 155. However, we need to go just a bit into this block before our primitive-based kernel begins.

The main action happens in just 14 lines of Assembly, running from setivli zero, 16, e16 to the final vmv.x.s a0, v9:

vsetivli  zero, 16, e16, mf2, ta, ma  # configure: 16 lanes int16, half register
vsext.vf2 v10, v11                    # widen input a: int8 → int16
vsext.vf2 v11, v12                    # widen input b: int8 → int16
vwmul.vv    v12, v11, v10             # widening multiply: v12 = v11 * v10 (int32 result)
vsetivli  zero, 8, e32, m1, ta, ma    # reduction begins: 8 lanes int32
vslidedown.vi v10, v12, 8             # shift upper 8 lanes down
vsetivli     zero, 8, e32, mf2, ta, ma
vadd.vv       v10, v12, v10           # add halves: 16 → 8 partial sums
vsetivli     zero, 4, e32, mf2, ta, ma
vslidedown.vi v11, v10, 4
vadd.vv       v10, v10, v11           # 8 → 4 partial sums
vsetivli     zero, 2, e32, mf2, ta, ma
vslidedown.vi v11, v10, 2
vadd.vv       v10, v10, v11           # 4 → 2 partial sums
vredsum.vs v9, v10, v9                # final: 2 lanes → scalar
vmv.x.s       a0, v9                  # extract scalar to a0

One obvious point standing out already, is of course our absence of any loop control mechanisms. Leveraging the SIMD primitive from Mojo, has let us reduce the workload down to a single string of instructions, where the flow of operations until and including the widening multiply can be roughly boiled down to:

  1. vsetivli setting up 16 lanes of int16 in half a register - similar to what we saw in the loop-based kernel, however not executed as a re-configuration this time around.
  2. vsext.vf2 widens our int8 inputs to int16, to make them match our int16 lanes, and prepare them for the upcoming widening to the int32 target output.
  3. vwmul.vv executes a widening multiply, executing all multiplies and a widening to int32 in one single instruction.

There's more to analyse in the codegen, but we'll quickly dive into one point of interest before proceeding: In our loop-based kernel, the widening multiply was executed with vwmacc.vv, due to it having to fold the result into a vector accumulator. With us not having a loop here, our compiler can leverage a slightly simpler widening multiply-only instruction.

Moving on, let's look at the reduction portion of our primitive kernel, i.e. the execution of our final products.reduce_add() call. Continuing from the vwmul.vv instruction, we now have all our products from the multiplied lanes, and we need those reduced down to a single scalar value to return from our kernel: 4. Rather than using vredsum.vs across all 16 lanes (i.e., a fully horizontal reduction), the LLVM backend opts for a pairwise tree reduction. This is achieved with the vslidedown.vi instruction, which slides the upper 8 lanes down onto the lower 8. 5. Following the slide down, the vadd.vv instruction sums them. This initially brings us from 16 to 8 partial sums. 6. In-between these instructions, vsetivli is used to continuously reconfigure the vector unit size, to match our new outputs. 7. This flow is executed two more times, bringing us from 8 partial sums to 4, and then from 4 partial sums to 2. 8. With just two to go, the compiler finally opts to reach for vredsum.vs, for the final reduction, with vmv.x.s popping that final scalar value out of the vector register, and into a0.

And that wraps up the execution flow!

One interesting question to ponder, is why the compiler didn't just use vredsum.vs to reduce the vector. The instruction is fully capable of achieving what we were after, and it could have done so in one instruction. So why the vslidedown.vi calls instead?

Up until the reduction point in our implementation, we were beautifully leveraging the strengths of the hardware, by performing vertical (i.e., parallelisable) operations. The reduction into one scalar value however, inherently doesn't fit into this mode of operating in a neat manner. This makes it a slightly awkward problem for our hardware to solve.

Calling vredsum.vs would sort it out, but it would do so in a potentially not super performance-scalable manner, going through the result lanes one by one - a horizontal solution. By reaching for a pairwise tree-reduction, the LLVM compiler once again turned this into a vertical problem, by continuously leveraging the vslidedown instruction to create scenarios that allow for more vertical vector operations (in this case, vadd.vv). Only when left with 2 results, vredsum.vs is called, at a point where we just have to add two registers, rather than our original 16 lane summing.

Zooming out to look at the full Assembly codegen, we got to see a similar-but-different output from the loop-based kernel. The compiler leveraged the hardware well, using the many available vector instructions to its advantage. Due to the simpler nature of the primitive-based kernel, we got slightly easier to read Assembly in the product calculation of the input lanes, and an interesting dive into how our compiler tried to as-much-as-possible turn the horizontal reduction problem, into a vertical one.

Step 7 (bonus): "SIMD" on a CPU without vector support

As you might recall, our Makefile also builds Assembly code for our generic CPU, i.e., a CPU we have previously concluded ships with no additional ISA extensions. So how does the compiler then handle generating SIMD Assembly?

Note: This is more of an exploratory bonus step, and you can feel free to ignore it. I mainly went into it to make it clear just how much "trouble" proper SIMD-supporting silicon saves us, when paired with a language and compiler that knows how to leverage it.

Unlike our vector-enabled-CPU codegen, the output for the generic CPU is verbose. If you want to explore it in-depth, either have a look at the output on your own machine, or dive into the GitHub repo for the post.

But a few headlines worth mentioning from analysing it, would be:

  1. Remember how vwmacc.vv in one instruction performed 16 multiplies, 16 widenings, and 16 adds? Well, our generic CPU doesn't have such luxuries, and instead has to return to our previously discussed __muldi3 library call - and no less than 16 times.
  2. __muldi3's calling conventions, ends up creating a massively cluttered instruction flow, due to needing to constantly move inputs and results around. The compiler anticipates this, and starts the flow out by allocating a 752-byte stack frame, in the addi sp, sp, -752 instruction call. This is further accentuated by the constant sd and ld shuffling, which ensures that new inputs can be passed to __muldi3. All of this chaos was abstracted away by the power of our vector registers and operations.
  3. Remember the pairwise tree-reduction that LLVM elegantly opted for to verticalize the otherwise horizontal result reduction issue? Well, since we have no vectors on this CPU, and thus no lanes, that issue disappears completely: We just end up with a pile of regular add instructions. Completely horizontal, just like the rest of the codegen - which makes sense of course, there are no other options for the compiler to reach for.

These are just a few of the points where the generic CPU struggles to fulfil the requested SIMD behaviour. Is this then an example of suboptimal compiler behaviour, or Mojo falling short? Quite the opposite, actually - we asked for SIMD behaviour to solve our int8 dot product kernel problem, and by God, we got that. We're just paying the price for giving the compiler the worst possible circumstances to fulfil that "request" within. All of this verbosity can more or less be summarized as compiler heroics, and the Mojo + LLVM combo managing to actually solve the issue under these constraints, is incredible.

Conclusion

And that wraps us up somewhat nicely. This post ended up going in a lot of different directions, with varying points of focus. It barely scratched the surface of RISC-V Assembly, nor of vectorization or the Mojo programming language - all of these are entire knowledge domains individually, and this really was just meant to be a fun, exploratory dive into all of them together.

But if I were to draw some "headline" conclusions from this little stroll through the world of Mojo's RISC-V codegen, I think it roughly boils down to:

Dog Bless Compilers

I get that compilers being an industry-defining gift to developers isn't exactly news, but my word, diving into Assembly shows just how much trouble this tool saves us, when combined with a well-made low-to-high-level programming language. Even when given the worst possible circumstances to work under, the LLVM compiler backend always faithfully translated our Mojo source code into matching RISC-V Assembly code.

SIMD is awesome - but be explicit about it

We saw SIMD achieve incredible optimizations when compiled to Assembly. We did however also see that we can never be certain that the compiler will correctly identify a chance to vectorize a scalar implementation. So better safe than sorry: If you have a vectorizable problem, and want SIMD optimizations, explicitly write your code to do so. To echo one of the sentiments from Mitchell Hashimoto's SIMD blog post, it will pretty much always be better to be explicit in performance-critical code. If you rely on assumptions about how the compiler behaves, you're only ever a version change away from that assumption silently falling apart.

So use the SIMD primitives available to you when you can. If you're writing e.g. Mojo code, the language support for this sort of meta programming is fantastic, as we saw in our exercises.

That'll be all from me this time around - thanks for sticking around until now if you did! I hope this might spark some interest in RISC-V, Mojo, or more low-level optimizations. And if not, I hope it was at least somewhat entertaining, and worth your time.