[X86][AArch64][RISCV] Consider branch lowering for scalar-splat masked.store
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Summary
I noticed a codegen tradeoff around `llvm.masked.store` when the mask is a full splat of a scalar condition.
When vector/predicated store support is enabled, several backends naturally lower this pattern to a native masked or predicated store. This avoids an explicit branch, but it also keeps the store-value computation unconditional.
For a scalar-splat mask, the masked store is effectively one of two cases:
```text
cond == false:
all-false masked.store, no memory store
cond == true:
all-true masked.store, equivalent to an ordinary store
```
If the store value is expensive to compute and the all-false path is common, an explicit branch plus ordinary store can be better because it skips the entire store-value computation on the false path.
This is not a request to blindly canonicalize all scalar-splat masked stores into branches. It is a possible backend/codegen profitability issue.
## Reduced IR
```llvm
declare void @llvm.masked.store.v4f32.p0(<4 x float>, ptr captures(none), <4 x i1>)
define noundef i32 @fp003_minimized_before(
ptr nofree readonly captures(none) %base,
ptr writeonly captures(none) %dst,
i32 %n) local_unnamed_addr {
entry:
%p0 = getelementptr i8, ptr %base, i64 8
%a = load <4 x i32>, ptr %p0, align 8
%p1 = getelementptr i8, ptr %base, i64 24
%b = load <4 x i32>, ptr %p1, align 8
%mask.cmp.scalar = icmp ne i32 %n, 0
%mask.cmp = insertelement <4 x i1> poison, i1 %mask.cmp.scalar, i64 0
%mask = shufflevector <4 x i1> %mask.cmp, <4 x i1> poison,
<4 x i32> zeroinitializer
%p2 = getelementptr i8, ptr %base, i64 40
%other = load <4 x i32>, ptr %p2, align 8
%sum = add <4 x i32> %b, %a
%diff = sub <4 x i32> %other, %sum
%as.float = uitofp <4 x i32> %diff to <4 x float>
tail call void @llvm.masked.store.v4f32.p0(
<4 x float> %as.float,
ptr %dst,
<4 x i1> %mask)
ret i32 0
}
```
The mask is a full splat of `%n != 0`.
For comparison, an explicit branch / ordinary-store form is:
```llvm
define noundef i32 @fp003_minimized_after(
ptr nofree readonly captures(none) %base,
ptr writeonly captures(none) %dst,
i32 %n) local_unnamed_addr {
entry:
%mask.cmp.scalar = icmp ne i32 %n, 0
br i1 %mask.cmp.scalar, label %store.block, label %exit
store.block:
%p0 = getelementptr i8, ptr %base, i64 8
%a = load <4 x i32>, ptr %p0, align 8
%p1 = getelementptr i8, ptr %base, i64 24
%b = load <4 x i32>, ptr %p1, align 8
%p2 = getelementptr i8, ptr %base, i64 40
%other = load <4 x i32>, ptr %p2, align 8
%sum = add <4 x i32> %b, %a
%diff = sub <4 x i32> %other, %sum
%as.float = uitofp <4 x i32> %diff to <4 x float>
store <4 x float> %as.float, ptr %dst, align 8
br label %exit
exit:
ret i32 0
}
```
## x86 AVX-512 codegen
With AVX-512 masked store support, the masked-store form lowers to a native masked store:
```asm
fp003_minimized_before:
xor eax, eax
neg edx
vmovdqu xmm0, xmmword ptr [rdi + 24]
vmovdqu xmm1, xmmword ptr [rdi + 40]
vpaddd xmm0, xmm0, xmmword ptr [rdi + 8]
sbb eax, eax
kmovd k1, eax
vpsubd xmm0, xmm1, xmm0
vcvtudq2ps xmm0, xmm0
vmovups xmmword ptr [rsi] {k1}, xmm0
xor eax, eax
ret
```
The explicit branch form instead skips the computation when `%n == 0`:
```asm
fp003_minimized_after:
test edx, edx
je .LBB0_2
vmovdqu xmm0, xmmword ptr [rdi + 24]
vmovdqu xmm1, xmmword ptr [rdi + 40]
vpaddd xmm0, xmm0, xmmword ptr [rdi + 8]
vpsubd xmm0, xmm1, xmm0
vcvtudq2ps xmm0, xmm0
vmovups xmmword ptr [rsi], xmm0
.LBB0_2:
xor eax, eax
ret
```
## AArch64 SVE codegen
With SVE enabled, the masked-store form lowers to a predicated store:
```asm
fp003_minimized_before:
cmp w2, #0
ldur q1, [x0, #8]
ldur q2, [x0, #24]
cset w8, ne
ptrue p0.s, vl4
dup v0.4h, w8
add v1.4s, v2.4s, v1.4s
ldur q2, [x0, #40]
mov w0, wzr
sub v1.4s, v2.4s, v1.4s
ushll v0.4s, v0.4h, #0
shl v0.4s, v0.4s, #31
cmpne p1.s, p0/z, z0.s, #0
ucvtf v0.4s, v1.4s
st1w { z0.s }, p1, [x1]
ret
```
The explicit branch form is:
```asm
fp003_minimized_after:
cbz w2, .LBB0_2
ldur q0, [x0, #8]
ldur q1, [x0, #24]
add v0.4s, v1.4s, v0.4s
ldur q1, [x0, #40]
sub v0.4s, v1.4s, v0.4s
ucvtf v0.4s, v0.4s
str q0, [x1]
.LBB0_2:
mov w0, wzr
ret
```
## RISC-V RVV codegen
With RVV enabled, the masked-store form lowers to a masked vector store:
```asm
fp003_minimized_before:
addi a3, a0, 8
vsetivli zero, 4, e8, mf4, ta, ma
vle32.v v8, (a3)
addi a3, a0, 24
vle32.v v9, (a3)
snez a2, a2
vmv.v.x v10, a2
addi a0, a0, 40
vsetvli zero, zero, e32, m1, ta, ma
vadd.vv v8, v9, v8
vle32.v v9, (a0)
vsetvli zero, zero, e8, mf4, ta, ma
vmsne.vi v0, v10, 0
vsetvli zero, zero, e32, m1, ta, ma
vsub.vv v8, v9, v8
vfcvt.f.xu.v v8, v8
li a0, 0
vse32.v v8, (a1), v0.t
ret
```
The explicit branch form is:
```asm
fp003_minimized_after:
beqz a2, .LBB0_2
addi a0, a0, 8
vsetivli zero, 4, e32, m1, ta, ma
vle32.v v8, (a0)
addi a2, a0, 16
vle32.v v9, (a2)
addi a0, a0, 32
vadd.vv v8, v9, v8
vle32.v v9, (a0)
vsub.vv v8, v9, v8
vfcvt.f.xu.v v8, v8
vse32.v v8, (a1)
.LBB0_2:
li a0, 0
ret
```
## Observation
The native masked/predicated-store lowering is reasonable when avoiding branches is preferred, especially if the condition is usually true or unpredictable.
However, when the mask is a scalar splat and the store value is expensive to compute, the branch form has an important advantage:
```text
cond == false:
native masked-store form:
still computes the vector value, then performs an inactive masked store
branch + ordinary-store form:
skips all loads, arithmetic, conversions, and the store
```
The branch form may be preferable when the all-false path is common, or when branch probability metadata indicates that the store is unlikely.
## Possible direction
It may be useful for codegen to consider a branch + ordinary-store lowering for scalar-splat `llvm.masked.store` when:
```text
1. the mask is known to be a full splat of a scalar condition,
2. the store is all-false or all-true as a whole,
3. the store value is relatively expensive to compute,
4. the target has no strong reason to prefer native predicated stores, or
5. profile/branch-weight information suggests that the all-false path is common.
```
This is a profitability question, not a correctness bug. Native masked stores are often the right choice, but this case shows that scalar-splat masks can represent coarse-grained control flow, and lowering them unconditionally to predicated stores may miss the opportunity to skip expensive store-value computation.
compiler-explorer sample : https://compiler-explorer.com/z/dnqj9az6q
## Performance comparison
## llvm-mca performance summary
| Target / CPU | Form | Instructions | Cycles | uOps | Block RThroughput | Summary |
|---|---:|---:|---:|---:|---:|---|
| x86-64 / Skylake AVX-512 | before | 1200 | 313 | 1600 | 2.7 | baseline |
| x86-64 / Skylake AVX-512 | after | 1000 | 264 | 1400 | 2.3 | improved |
| AArch64 / Neoverse V1 + SVE | before | 1600 | 314 | 1900 | 3.0 | baseline |
| AArch64 / Neoverse V1 + SVE | after | 1000 | 167 | 1200 | 1.5 | significantly improved |
| RISC-V32 / SiFive E76 + V | before | 1900 | 3603 | 1900 | 13.0 | baseline |
| RISC-V32 / SiFive E76 + V | after | 1400 | 3002 | 1400 | 9.0 | improved |
## Relative change
| Target | Instructions | Cycles | uOps | Block RThroughput |
|---|---:|---:|---:|---:|
| x86-64 / Skylake AVX-512 | -16.7% | -15.7% | -12.5% | -14.8% |
| AArch64 / Neoverse V1 + SVE | -37.5% | -46.8% | -36.8% | -50.0% |
| RISC-V32 / SiFive E76 + V | -26.3% | -16.7% | -26.3% | -30.8% |
Contributor guide
Research direction
Start with the reduced IR and the x86 AVX-512, AArch64 SVE, and RISC-V RVV before/after codegen shown in the issue. Compare the llvm-mca summaries, then determine how scalar-splat masked.store cases should choose between native predication and a branch based on profitability; done means the selected lowering preserves correctness and improves the relevant cases.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100