[RFC] `@llvm.parallel_alloca` intrinsic for vectorizable stack allocations

IR & Optimizations Loop Optimizations
1

This RFC proposes the addition of a new LLVM intrinsic function, @llvm.parallel_alloca. Similarly to alloca, this intrinsic allocates stack memory. However, it is not hoistable and can be vectorized such that each vector lane receives a unique pointer.

Motivation

alloca instructions are typically assumed to be hoistable into the entry block of the function, either directly when the frontend is emitting them (e.g. Clang) or during inlining. If you declare a variable in a loop, the related alloca is placed in the entry block of the function instead of inside the loop.

Usually, this behavior is fine. The lifetime of the allocation doesn’t overlap between loop iterations when the loop is executed serially. As multiple iterations don’t simultaneously access the allocation, the same memory can be shared between all iterations.

However, if the loop is vectorized, multiple iterations may run in parallel. The hoisted alloca may now introduce a memory dependency between iterations that did not need to exist, and can prevent vectorization. Worse yet, if the loop is annotated with @llvm.loop.parallel_accesses (e.g. via #pragma clang loop vectorize(assume_safety)), this dependency is ignored and the program behaves incorrectly.

The latter behavior is particularly insidious when combined with inlining:

void helperFunc(int *A, int i) {
    int t[16];
    // Fill entire local array with 'i'
    t[0] = i; t[1] = i; t[2] = i; t[3] = i;
    t[4] = i; t[5] = i; t[6] = i; t[7] = i;
    t[8] = i; t[9] = i; t[10] = i; t[11] = i;
    t[12] = i; t[13] = i; t[14] = i; t[15] = i;
    // Just read one of them, it has to be 'i'. Right? Right???
    A[i] = t[i];
}

void foo(int *A, int n) {
    #pragma clang loop vectorize(assume_safety)
    for (int i = 0; i < n; ++i) {
        // This call gets inlined.
        helperFunc(A, i);
    }
}

As the alloca for the t array is actually hoisted into the entry block of foo during inlining, the above example causes big problems. Multiple iterations are simultaneously writing into the same allocation for t. If foo is called with n = 16, A is 7 7 7 7 7 7 7 7 15 15 15 15 15 15 15 15 on my machine.

Worse yet, it’s not only vectorize(assume_safety) that uses the @llvm.loop.parallel_accesses metadata. #pragma omp simd is also affected. I believe this breaks the semantics of OpenMP.

The reason why this behavior doesn’t always break nearly every program using those pragmas is that the Mem2Reg pass usually promotes allocas into registers, which do get vectorized safely. However, if that doesn’t happen for one reason or another, like here, we can be in big trouble.

Method

In general, we can’t demand passes to be aware of metadata for correctness, so we can’t demand that alloca can never be hoisted from parallel loops. Loops are marked as parallel using the llvm.loop.parallel_accesses metadata, and LLVM passes shouldn’t have to be aware of metadata to produce correct code.

Instead, we propose adding a new intrinsic for emitting parallel-safe stack allocations:

<N x ptr> @llvm.parallel_alloca(i64 immarg size, i64 immarg align)
<N x ptr> @llvm.parallel_alloca(i32 immarg size, i32 immarg align)

This intrinsic is mostly a drop-in replacement for alloca, but with certain semantic differences that make it behave well in parallel loops.

The user can do e.g. call ptr @llvm.parallel_alloca(64, 8) inside a parallel loop. This can be treated similarly to a regular alloca [64 x i8], align 8, except that:

  • It must be treated as if it is not speculatable and has side effects.
  • It can be vectorized to return multiple pointers and allocate the same amount of memory for each vector lane.

If this intrinsic is called outside of a parallel loop, it can be treated as a normal alloca and lowered as such. If the intrinsic is lowered into an alloca in a parallel loop, the parallel loop must be demoted into a sequential loop. However, the intended behavior is that the intrinsic is not lowered until after loop vectorization has run and potentially vectorized the intrinsic.

For a single pointer return value, ptr @llvm.parallel_alloca(s, a) lowers into an alloca [s x i8], align a in the entry block. The address space of the alloca is determined from the address space of the returned pointer type of @llvm.parallel_alloca.

For a vector of pointers, <N x ptr> @llvm.parallel_alloca(s, a) is lowered as below:

; In entry block:
; K = roundToNextMultipleOf(s, a)
%1 = alloca [(N * K) x i8], align a
%2 = getelementptr inbounds i8, ptr %1, <0, K, 2K, 3K, ...>
; replace all uses of the original @llvm.parallel_alloca call with %2.

This effectively “privatizes” the allocation for each vector lane; each lane gets a unique pointer with at least the given amount of bytes with the specified alignment.

Lifetime annotations are not extended to accept pointers from @llvm.parallel_alloca to limit the impact of this RFC. The allocated memory is freed similarly to alloca with no lifetime annotations; it is only freed when the function that called the parallel_alloca intrinsic returns.

Impact

Areas of the compiler utilizing the parallel metadata may be affected. That is to say, if a pass is aware of llvm.loop.parallel_accesses and llvm.access.group, it also has to be aware of @llvm.parallel_alloca in the future.

The general rule is that code behaves correctly if it is either fully unaware of parallel metadata, or fully aware. This new intrinsic is an extension to the existing parallel metadata.

There are very few passes that currently utilize parallel metadata, so we expect the impact to be minor.

Implementation

  • Add a utility function lowerParallelAlloca for lowering @llvm.parallel_alloca into a normal alloca in the caller’s entry block, as specified in the Method section above.

  • InlineFunction should upgrade alloca into @llvm.parallel_alloca when inlining a function call with @llvm.access.group metadata.

    • @llvm.access.group is part of the “parallel metadata”, so if a pass wants to propagate that metadata into instructions that didn’t have it before, it must also be aware of parallel_alloca.

  • Clang should emit @llvm.parallel_alloca instead of alloca when creating temporary variables inside a parallel loop.

    • Otherwise, previous behavior persists with alloca of local variables being hoisted outside of parallel loops.

  • Mem2Reg should allow promoting @llvm.parallel_alloca

    • Not being aware of @llvm.parallel_alloca would produce correct but slow code
    • Can be implemented in a simple manner: if the @llvm.parallel_alloca was turned into an alloca in the entry block, would it be promotable? If so, do that and promote just like every other alloca.

  • LoopVectorize should call lowerParallelAlloca immediately after vectorizing it and strip parallel metadata from the loop.

    • It can readily vectorize the intrinsic with VPWidenIntrinsicRecipe
    • It could also not lower it there, but that could prevent other optimization passes from operating on the resulting alloca until the intrinsic is finally lowered

  • PreISelIntrinsicLowering needs to lower @llvm.parallel_alloca if no prior pass lowered it yet

    • We should be able to disregard parallel metadata here, but can also remove it to be safe.

Authors: @juliusikkala (me), @pjaaskel (Pekka Jääskeläinen).

There is a work-in-progress implementation in my branch.

Prior related discussions:

Attached is a slide deck version of this RFC:

parallel_alloca.pdf (114.6 KB)

2

When is the memory allocated by llvm.parallel_alloca freed?

(See also incorrect transformation around icmp: unclear semantics of "lifetime" intrinsics leads to miscompilation · Issue #45725 · llvm/llvm-project · GitHub )

3

Thanks for linking that thread, it somewhat changed my outlook on what to do with lifetimes here.

It seems that specifying such a parallel_alloca_free requires more effort than I have capacity for right now, and we’re mostly concerned with fixing this miscompilation with parallel loops that also impacts pocl and prevents Slang’s LLVM target from using parallel metadata. So for now, memory would only be freed when the surrounding function returns.

Based on that discussion, I think we’ll want to make sure that we don’t tie too closely to the llvm.lifetime.* intrinsics here so that there’s enough wiggle room to add something like parallel_alloca_free in a separate RFC later and turn this parallel_alloca into a more general “alloca but fixed” intrinsic. I was planning on leaving lifetime intrinsics out anyway, but mostly to reduce the impact radius of this RFC. (will edit original post to mention intentional avoidance of lifetime intrinsics)

4

“When the function returns” is problematic: you have N allocations, where N is the number of loop iterations, and you have to keep them all around until the function returns. So you need dynamic allocation, which is probably not what you want.

What you actually want is “when the function returns or the @llvm.parallel_alloca call runs again”? That’s tricky to specify, though, because it depends on the identity of the alloca, but instructions don’t have an identity: they can be cloned/merged/etc.

I guess you can add an @llvm.parallel_alloca_anchor intrinsic in the entry block, which returns a token, and use that token as an argument to @llvm.parallel_alloca.

5

Yeah, that’s not what I want. So far, I’ve mostly just reasoned about this by how it lowers: in the end, each call simply generates an alloca in the entry block.

Yes, that is the behavior I was expecting indeed. Thanks, I just realized that this does indeed lead into a practical case of allocating a bit too much: when used in a vectorizable loop, the call gets cloned into both the vector body and the scalar epilogue, so the allocation occurs twice. Also, unrolling a loop with @llvm.parallel_alloca could be quite a tragedy.

All of this gets even more annoying when some of the calls to the “same” @llvm.parallel_alloca get vectorized and some don’t, because they lower differently. I suppose when lowering, one would need to collect all calls with the same token and find the widest one.

At that point it would seem cleaner to consider the parallel_alloca_free style approach, although I’m not fully sure how to proceed with that. I’ll have to take some more time thinking that through, because I’d love to avoid having to do any major overhauls to the stack allocation infrastructure in LLVM.

Edit: Nevermind, I’ll just try the token approach and see if it causes any trouble. If it’s fine, I’ll update the RFC to add that.