[RFC] Add elementwise modifier to atomic loads and stores

IR & Optimizations
1

Overview

This RFC proposes adding an elementwise modifier to atomic loads and stores:

%result = load atomic elementwise <2 x float>, ptr %p syncscope("agent") monotonic, align 4
store atomic elementwise <2 x float> %value, ptr %p syncscope("agent") monotonic, align 4

For vector atomic loads and stores, the modifier gives the instruction per-element atomic semantics. A vector operation with elementwise behaves as if it were expanded into one scalar atomic operation per fixed-vector element, with no additional ordering between the per-element operations beyond the ordering specified for each scalar operation. This is analogous to the existing proposal: RFC] Add elementwise modifier to atomicrmw.

A draft implementation for load atomic elementwise is available at https://github.com/llvm/llvm-project/pull/204556.

Motivation

Currently, LoadStoreVectorizer does not vectorize atomic loads or stores. It only considers simple loads and stores, and skips instructions for which LoadInst::isSimple() or StoreInst::isSimple() is false. As a result, adjacent scalar atomic loads cannot currently be combined by this pass.
For example, the following adjacent scalar atomic loads remain scalar today:

define void @test(i64 %idx, ptr addrspace(1) %out) {
  %p = inttoptr i64 %idx to ptr addrspace(1)
  %num1 = load atomic float, ptr addrspace(1) %p syncscope("agent") monotonic, align 4
  %p1 = getelementptr inbounds nuw i8, ptr addrspace(1) %p, i64 4
  %num2 = load atomic float, ptr addrspace(1) %p1 syncscope("agent") monotonic, align 4
  %res = fadd float %num1, %num2
  store float %res, ptr addrspace(1) %out, align 4
  ret void
}

Without a new IR modifier, vectorizing the example into the following IR would change the semantics:

define void @test(i64 %idx, ptr addrspace(1) %out) #0 {
  %p = inttoptr i64 %idx to ptr addrspace(1)
  %1 = load atomic <2 x float>, ptr addrspace(1) %p syncscope("agent") monotonic, align 4
  %num11 = extractelement <2 x float> %1, i32 0
  %num22 = extractelement <2 x float> %1, i32 1
  %res = fadd float %num11, %num22
  store float %res, ptr addrspace(1) %out, align 4
  ret void
}

A vector atomic load without elementwise is a whole-value atomic operation. The original scalar loads are independently atomic, so replacing them with a whole-vector atomic load would change the IR semantics. The same issue applies to stores: replacing independent scalar atomic stores with a vector atomic store would turn independently atomic element stores into one whole-value atomic store.

AMDGPU reports setSupportsUnalignedAtomics(false), so SelectionDAG rejects an atomic load or store whose IR alignment is smaller than the memory value type size. As a result, a whole-vector atomic load or store of <2 x float> with align 4 is rejected as an under-aligned 64-bit atomic memory access:

load atomic <2 x float>, ptr addrspace(1) %p syncscope("agent") monotonic, align 4

However, the intended operation is not a single 64-bit atomic memory access. It is two independent 32-bit atomic element accesses from adjacent addresses. The proposed elementwise modifier makes that distinction explicit in IR:

load atomic elementwise <2 x float>, ptr addrspace(1) %p syncscope("agent") monotonic, align 4

This also gives targets a way to reason about the alignment using the scalar element type rather than the whole vector type. On AMDGPU, ordinary global memory operations can use dword-aligned wide accesses: SITargetLowering::allowsMisalignedMemoryAccessesImpl() allows extended global address space accesses when the alignment is at least 4 bytes, even if the access size is wider. With elementwise, AMDGPU can preserve the per-element atomic semantics while still selecting an appropriate wide global memory operation, such as a 64-bit global load or store, when that is valid for the target.

Proposed Syntax

<result> = load atomic [volatile] [elementwise] <ty>, ptr <pointer> [syncscope("<target-scope>")] <ordering>, align <alignment> [, !invariant.group !<empty_node>]
store atomic [volatile] [elementwise] <ty> <value>, ptr <pointer> [syncscope("<target-scope>")] <ordering>, align <alignment>

Lowering Strategy

Targets may lower an elementwise atomic vector load or store using any instruction sequence that preserves the per-element atomic semantics. For AMDGPU, this may allow adjacent dword-aligned elementwise atomic operations to be represented using an appropriate wide global memory operation when the target can prove that doing so preserves the required semantics.

Implementation Plan

If this idea is approved, I have the following plans:

  1. Add IR support for load atomic elementwise and store atomic elementwise.

  2. Teach AMDGPU lowering to preserve per-element atomic semantics for elementwise atomic loads and stores.

  3. Extend LoadStoreVectorizer to combine eligible adjacent scalar atomic loads and stores into elementwise vector atomic operations.

1 Like
[RFC] Add elementwise modifier to atomicrmw
2

CC @arsenm @danilaml @YonahGoldberg @preames @shiltian @jayfoad

3

I wonder why? Isn’t it legal to vectorize non-elementwise atomic loads/stores? I guess you would need TLI checks to see if the target supports that larger atomic. Maybe easier to just always vectorize to the elementwise version and have AtomicExpand handle.

But we might want TLI checks any way? One requirement from my previous RFC was that by default elementwise atomics should be treated the same as non-elementwise and targets can opt-in to expanding. So we still don’t want to vectorize to an elementwise atomic that is too large for a target?

Is the point here that load atomic <2 x float> with align 4 can likely not be lowered as a 64-bit atomic because it’s under aligned, but we can break this up into 2 load atomic float with align 4 and lower it this way? In PTX the .v2b32 load still needs to be 8-byte aligned.

I have a patch to handle this kind of expanding for atomicrmw in AtomicExpandPass. We can easily extend to load/store atomic.

I originally implemented the elementwise on atomicrmw because you couldn’t lower vector atomicrmw to PTX vector atom because in PTX all vector ops are elementwise atomic. But with load/store atomic we have vector versions, but we just legalize to the scalar version to get around the fact that it’s not legal to weaken the atomic. IDK how much it would buy us to change this lowering.

So anyway I guess this is nice because it gives more flexibility. Front-ends can specify arbitrarily large vector elementwise atomics and we can break it up for them.

+1 in general, I just have some questions about the motivations.

4

Actually I just saw we have ld.relaxed.v8b32 that we can target with this that we’re not able to target currently because there is no ld.relaxed.b256

5

Actually, I think atomic loads/stores with unordered or monotonic ordering can be vectorized, but we still need some TLI checks to make sure the target supports the resulting larger atomic operation.

Yes, that is the point.

I think atomic loads/stores are a bit different from atomicrmw instructions here. On AMDGPU, atomicrmw instructions are lowered to real atomic RMW instructions. For example, atomicrmw fadd lowered to global_atomic_add_f64.

However, atomic loads/stores are usually lowered to normal load/store instructions according to the address space. For example, load atomic i64, ptr addrspace(1) lowered to global_load_b64, which is similar to a normal global load.

The problem is that AMDGPU currently assumes natural alignment for these vector atomic loads/stores. So load atomic <2 x float> with align 4 cannot be safely lowered as a single 64-bit load, because that would require 8-byte alignment. What I want to add with elementwise is a way to tell the backend that the vector components are independent atomic operations. We can use element-size alignment.

Thanks! Let’s take the first step and add elementwise to atomic loads/stores.

6

It’s not quite clear if merging atomics to wider atomics is legal in LLVM IR, there are differing opinions in the active discussions here [RFC] Semantics of partially overlapping atomic accesses and here [RFC][LangRef] Specify that the accessed bytes of concurrent atomics must be either disjoint or the same by ritter-x2a · Pull Request #204329 · llvm/llvm-project · GitHub .
There are at least cases where the current libcall lowering on the CPU would break if arbitrary adjacent atomics are merged.

3 Likes
7

This proposal seems fine

1 Like
8

+1 to the proposal.