What is loop unrolling? How you can speed up Mojo🔥 code with @unroll

Open any introductory programming book and you’ll find several pages dedicated to structured programming concepts, i.e. making use of loops, conditions and functions extensively for better clarity and maintainability for your code. It helps you express your ideas and solutions neatly and elegantly in code. However, these benefits come at a cost: performance overhead. 

Consider a simple for-loop: each iteration involves condition checks and structural setup, adding to the computational workload. This is especially true for loops with a tiny body. While these overheads usually outweigh the ease-of-use and clarity of code benefits for most use-cases, in some performance-critical situations we might want to trade-off the structural elegance to gain faster performance. Loop unrolling is one such procedure and in this blog post we’ll discuss how you can unroll Mojo🔥 loops to speed up code.

What is loop unrolling?

Loop unrolling is a loop transformation technique that aims at improving the performance of a program at the expense of code size. When a loop is unrolled, instructions that control the loop are eliminated (full unroll) or reduced (by unroll factor) and replaced with multiple copies of the loop body, thereby reducing the computational overhead but increasing the size of the code. A programmer can manually unroll a loop by hand coding the unrolled loop or provide special instructions in code so the Mojo compiler can do it for you. In this blog post, we’ll explain how loop unrolling works and how to do it in Mojo.

Understanding loop unrolling

Typical computer programs include both straight-line code and loops for repeated execution of code until a specific condition is met. For any non-trivial program, the percentage of time that is spent executing loops is longer than it is to execute the straight line code. While straight line code is typically executed only once, loops may have to be executed tens, hundreds, or thousands of times or more. 

There are many compiler optimizations techniques to speed up execution of loops and loop unrolling is one such compiler optimization technique. Let’s take an example of a simple Mojo loop that prints values from 0 to 4:

When compiled, the CPU sees the following pseudo assembly:

In the code above you can see that for every loop iteration there is arithmetic operation i+=1 one comparison operation i<5 and one conditional jump operation jump loop_body. This is wasteful, since the ratio between useful instructions: print(i) and instructions needed just to control the loop is high.

What if we can adjust that ratio? At compile time, the compiler knows the loop bounds, it could replace the entire loop with unrolled code (this is called full unrolling). In Mojo you can use the @unroll decorator to instruct the compiler to unroll it the loop as follows:

Would generate the following pseudo assembly:

Mojo unrolls the loop and replaces 2 lines of loop code with 5 lines of repeated loop body instruction, increasing the overall code size but removing loop control code since there is no loop overhead anymore. Let’s dive deeper into Mojo APIs for loop unrolling.

Mojo🔥 APIs for loop unrolling

Mojo🔥 offers first class support for manual loop unrolling. In the earlier example, we saw Mojo’s @unroll decorator which lets you unroll loops without having to manually write the unrolled version. In Mojo, loop unrolling syntax is part of the core language. We made the very conscious decision that programmers should have more control on the compiler’s behavior. 

Traditional compilers such as clang, gcc and others also support automatic loop unrolling that aims to expand the loop and sacrifice code size for more branch free instruction generation. But these are compiler options and outside the control of the programmer. It is often considered “compiler magic” since it is part of black-box compiler optimizations. As with any out-of-the box optimization, they work very well for specific cases and don’t work for others and the programmer has no control over it. See Chris Lattner’s CGO2012’s keynote on "Compiler Optimization: Increasing Research Impact” for a deeper discussion on this topic.

Mojo’s design philosophy is to give mojo programmers more explicit and predictable control on what the compiler does. For example Mojo’s @unroll and @unroll(n) are guaranteed syntax that if compilation passes, the decorated loops are guaranteed to be unrolled; if the compiler cannot unroll the loop, compilation fails and tells the user why the loop couldn’t be unrolled. The unroll syntax is not just a hint to the compiler, but gives the programmer fine grained control of this compiler behavior. Mojo is still young and evolving rapidly, and over time unroll will provide more control over different optimization levels. If you have feature requests we’re all ears, please file requests on GitHub.

Mojo🔥 supports 3 language features for loop unrolling @unroll, @unroll(n) and unroll[]() function. Let’s take a look at an example for each.

@unroll: full loop unrolling

The simplest way to unroll is to use the @unroll decorator which fully unrolls the loop. This is easy and straightforward to use, but increases program code size based on the total number of iterations in the loop. This can be undesirable, particularly for applications where you’re trying to keep binary sizes low.

@unroll(n): with specified unroll factor n

In the previous example, we did a full loop unroll, i.e. we eliminated the loop completely. In practice, loop bounds can be too large for full unroll to provide performance gains, so it’s efficient to specify an unroll factor instead. Small unrolling factors like 2 or 4 typically see speed improvements and Mojo lets you specify the unroll factor n using @unroll(n) loop decorator. Specifying an unroll factor affords a good balance between the benefit of loop unrolling and potential code size explosion.

The factor by which to unroll can be specified based on potential speed gains against code size increase. For example an unroll by factor 4 looks like this:

Unlike the full unroll in the previous example, this won’t eliminate the loop completely, but would loop fewer times and update the loop body with code repetitions based on unroll factor.

The CPU would see pseudo assembly that looks like this:

Currently, @unroll and @unroll(n) can only be added to loops with a constant number of iterations that the compiler can infer statically. There is also a restriction on the loops that they need to be “simple loops”, those of which don’t have multiple exits in the loop body. We will add more sophisticated support for loop unrolling in the near future and we always welcome feature requests on GitHub.

unroll(): repeatedly evaluate a function

In addition to @unroll decorators, Mojo also offers an unroll() function in the standard library. Rather than unroll a loop, unroll[]() repeatedly evaluates a function parameterized by 1, 2 or 3 integer parameters. 

This is equivalent to 1, 2 or 3 nested loop with the body of the loop defined in a function func, passed as a compile time parameter to unroll()

Peek under the hood: Performance vs. executable size trade-off

Let’s see another simple example of loop unrolling, to analyze the trade-offs between program performance and its executable size.

The above program loops over different iteration values of [0, 1, 2, 3] and calls foo with these values in each iteration. The output of the program is:

We can also write this program below without using a loop:

Comparing these two programs:

• no-loop.mojo is functionally the same as loop.mojo.
• no-loop.mojo uses 12 lines of code to accomplish the same thing as loop.mojo which only uses 3 lines of code. If the loop body contains many lines of code L, unrolling the loop will create N * L lines of code if we unroll the loop N times.
• loop.mojo has a control flow of for in the program , while no-loop.mojo does not have control flow. Control flow requires extra work during execution time, e.g. for needs to check the condition whether the iteration is still in the range or not, and choose either entering the loop body or going to exit the loop.

Manually unrolling a loop is quite tedious when the loop body is large and it can also lead to significant lines of code increase in the program. Unrolling the loop during compilation can avoid tedious manual editing of the program with no source code size explosion.

Now, let’s take a look at what happens if we do loop unrolling in the compiler i.e. no need for the programmer to manually expand the source code of the program.

Here is the LLVM IR for the loop version: 

In this loop version:

• There are multiple basic blocks [2] (3:, 6:, 7:, 8:, :12) and control flow (br label)
• Branching can occur at an extra cost during execution in many computer architectures 😞
• The IR size will stay (almost) the same if loop iteration increases so that the generated binary executable size stays the same. 😌

Where as LLVM IR for the unrolled the loop:

If the compiler unrolls the loop automatically, in this case, the generated IR:

• Does not have control flow anymore with only one basic block.
• Avoid branching cost during execution with a single basic block. Open more possibilities to other compiler optimizations with straight-line code. 😌
• The IR size will grow proportional to the number of loop iterations which leads to binary executable size increase. 😞

Conclusion

Loop unrolling is a simple technique used to optimize a program's execution speed by reducing or eliminating loop control instructions and loop tests. Mojo gives the programmer explicit and fine-grained control on how unrolling should be done on loops. This empowers the programmer to determine the most suitable unrolling factor for a given situation, based on their performance vs. increased code size trade off appetite. 

We hope you enjoyed this explainer, download Mojo, try it out and share your feedback with us! Here are some additional resources to get started:

• Get started with downloading Mojo
• Head over to the docs to read the programming manual and learn about APIs
• Explore the examples on GitHub
• Join our Discord community
• Contribute to  discussions on the Mojo GitHub
• Read and subscribe to Modverse Newsletter 
• Read Mojo blog posts, watch developer videos and past live streams
• Report feedback, including issues on our GitHub tracker

Until next time🔥!