What does @unroll and @unroll(n) do?

[cjohnson318]

What does @unroll and @unroll(n) do? Can someone explain where to use this and why? (Sorry, I'm a dolt.)

I see this in the docs, but I'm not sure what to do with it.

[weiweichen]

These are two loop decorators to tell the compiler to unroll a loop (https://en.wikipedia.org/wiki/Loop_unrolling). This doesn't impact the functionality of the loops, but potentially can help for better performance since it opens possibility for further compiler optimizations.

The syntax to use them is something like

# Fully unroll the loop's 10 iterations into 10 `do_something` calls and remove the for-loop.
@unroll 
for i in range(10):
  do_something(i)

or

# Unroll every 2 iterations and loop over 5 times.
@unroll(2)
for i in range (10):
  do_something(i)

This decorator can be attached to while statement too.
Note that currently the compiler can only unroll a loop, when:

• Its lower bound, upper bound and induction variable step every iteration are compile time constants.
• There is no early exits in the loop body that makes the loop trip count dynamic during runtime.

Otherwise, compilation fails if a loop is decorated with @unroll.

Here is a brief description of these two decorators in Mojo changelog on 2023-07-26.

[iamtimdavis]

@scottamain might be useful to have an example in the docs

[cjohnson318]

@ntrel okay, that made sense to me, thanks!

[leandrolcampos]

Hi @weiweichen,

I hope this message finds you well.

I have a question: why can't the unrolling factor n be a compile-time value, but only a literal integer? Or is this a temporary limitation?

This code works

@unroll(3)
for i in range(5):
    print(i)

but this fails

alias n: Int = 3
@unroll(n)
for i in range(5):
    print(i)

with the message

error: Expression [37]:2:8: unsupported decorator on 'for' statement
@unroll(n)
~~~~~~^~~

expression failed to parse (no further compiler diagnostics)

[gryznar]

IMHO this could be reported as a bug (at least in terms of wrong error message)

[weiweichen]

Yes, unrolling factor being compile-time value support will be added in the future.

[weiweichen]

this in the docs performs the same loop unrolling functionality as a library function. There are a few differences between using library function of unroll and decorator @unroll:

• Library function call requires the induction variable to be a parameter while the decorator uses the induction variable as a dynamic variable.
• Library function call unroll the loop so that the program the compiler starts to compile is with unrolled code. This can potentially increase the amount of code to compile depends on the amount being unrolled.
• Decorator unrolling happens at later stage of compilation which prevents program size explosion too early.