Adding non-inlined for loops to Jasmin

[haselwarter]

Adding non-inlined for loops to Jasmin

This RFC proposes the addition of non-inlined for loops to Jasmin. Currently, Jasmin supports while loops, but for loops are always inlined, i.e. unfolded at compile-time. In the remainder of the proposal, when we refer to for loops, we mean non-inlined for loops unless explicitely stated otherwise.

We aim to provide a simple extension to Jasmin, even if it does not directly cover all possible use cases. Some extensions are discussed at the end of the proposal.

Motivation

For Jasmin programmers

More idiomatic code closer to specifications. For instance, hacspec's for loops could be more directly translated.

For tool developers

• easier termination / losslessness analysis: EasyCrypt, SSProve, possibly other backends
• once the extensions outlined below are considered: more predictably structured assembly code (e.g. spilling always occurs in a certain way)

Proposal

We propose to translate for loops to while loops at an early stage of compilation.

A loop of the form

FOR x=var EQ exp_begin=pexpr TO exp_end=pexpr (COMMA step=pexpr)? {
    body=pblock
}

gets translated to

x = exp_begin ;
while (x < exp_end) {
    body
    x += step ;
}

A loop of the form

FOR x=var EQ exp_begin=pexpr DOWNTO exp_end=pexpr (COMMA step=pexpr)? {
    body=pblock
}

gets translated to

x = exp_begin ;
while (x > exp_end) {
    body
    x -= step ;
}

If the step=pexpr annotation is omitted, the step size defaults to 1.

Considerations

Ensuring termination

Care has to be taken to ensure termination, which is one of the motivations for the introduction of structured for loops. Two potential problems arise. First, the step size needs to be positive. Second, the loop counter (and the loop bound) must not be modified in the loop body.

The easiest way to satisfy the positive-step-size requirement is to require the step size to be known at compile time.

The enforcement of the constraint that the loop counter (and bound) must not be modified by the loop body would have to be implemented as a syntactical check.

Implicit or explicit unrolling

We can discriminate whether or not a loop should be unrolled by (1) unrolling if the loop counter is declared inline, or (2) requiring a user annotation (inline is already a keyword). Option (2) has the disadvantage that it introduces a way for the user to write programs that must be rejected, by declaring a loop not inline but using an inline counter. On the other hand, this problem already exists in other uses of inline variables, say in the condition of a while loop.

Limitations and possible extensions

In discussions on the formosa zulip and at HACS, a number of potential issues with for loops were brought up.

No implicit spilling

If the loop counter is stored in a register, that register cannot be overwritten during the entirety of the loop body. This can lead to problems with register allocation. The problem is usually solved by "spilling" the loop counter to the stack at the beginning of the loop body, and restoring it from the stack to the counter register at the end of the loop body.

In keeping with the general philosophy of producing predictable assembly code, no implicit spilling of the loop counter to the stack is introduced. The issue is not exclusive to for loops, but harder to circumvent than in other situations, as we shall now see.

The problem that arises is that spilling for for loops cannot be done manually in the way it is currently done in while loops. Consider the following example.

reg u64 i;
for i = 0 to n {
    ...body...
}

We might want to spill i to the stack as follows:

reg u64 i;
stack u64 is;
for i = 0 to n {
    is = i;
    ...body...
    i = is;
}

Unfortunately, we cannot allow the user to write this spilling by hand, as it would violate the constraint that the loop counter must not be modified in the loop body.

A future proposal could extend non-inlined for loops with explicit spilling annotations. We suspect that, even then, the choice of spilling location should still be left to the programmer rather than being mandated by the Jasmin compiler. Below we declare an explicit spilling to the stack.

reg u64 i;
stack u64 is;
for i = 0 to n [spill=is] {
    ...body...
}

The [spill=is] annotation would then produce the above loop with explicit spilling of i into is. Again, in order to ensure termination, a syntactic check would be required to ensure that no writes to is occur in body.

If stack space is not considered a scarce resource, an explicit variable declaration for the spilling could be omitted and inserted by the compiler if the annatiton is simply [spill] rather than [spill=is].

No loops over register arrays

Registers are not dynamically addressable at runtime, so it will not be possible to replace existing unrolled for loops operating on register arrays with the new non-inlined for loops. This is not a limitation of the proposal per se but rather a property of register arrays.

For-each iterators

The current proposal only considers for loops with simple fixed-step iterations. If the current proposal can be implemented successfully, an extension to for-each style iteration could be considered.

[bgregoir]

Hi, thank for the proposal. For the moment this feature does not seem to be a blocking point for users since they can write code using normal while loop. Since we have a lot of features that are blocking for users I prefers that we focus first on those. Naturally, I can understand that this feature can be useful for you and for other. So if you or someone else has the time to develop it, we will certainly integrate it. Furthermore, I'm open to discussion if someone need help to implement it. For the "inline" this can be used as #inline before the for loop. I think it is clearer and more explicite than looking for the loop index.