LLVM should split some potentionally infinite canonical loops into branch and two loops (infinite and finite).
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
Consider the following Rust function:
```rust
pub fn sum(start: u64, n: u32)->u64 {
let mut res = start;
let mut i = 0;
while i <= n {
res += i as u64;
i += 1;
}
res
}
```
The loop there is guaranteed to progress if `n` has any value except `u32::MAX`. And when we can assume that loop would finish, LLVM can nicely remove it and replace by arithmetic progression sum.
Unfortunately, LLVM cannot be sure if loop finishes or runs endlessly so it ends up keeping it, and worse, unrolling and generating lots of unneccessary code:
```llvm-ir
define noundef i64 @example::sum::h44644981a6eb5cf7(i64 noundef %0, i32 noundef %n) unnamed_addr #0 !dbg !7 {
start:
#dbg_value(i64 %0, !15, !DIExpression(), !22)
#dbg_value(i64 %0, !17, !DIExpression(), !23)
#dbg_value(i32 %n, !16, !DIExpression(), !22)
#dbg_value(i32 0, !19, !DIExpression(), !24)
%1 = add i32 %n, 1, !dbg !25
%n.off = add i32 %n, -7, !dbg !25
%switch = icmp ult i32 %n.off, -8, !dbg !25
br i1 %switch, label %vector.ph, label %bb2.preheader, !dbg !25
bb2.preheader:
%start1.sroa.0.06.ph = phi i64 [ %0, %start ], [ %8, %middle.block ]
%i.sroa.0.05.ph = phi i32 [ 0, %start ], [ %n.vec, %middle.block ]
br label %bb2, !dbg !25
vector.ph:
%n.vec = and i32 %1, -4
%2 = insertelement <2 x i64> , i64 %0, i64 0
br label %vector.body, !dbg !26
vector.body:
%index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ], !dbg !26
%vec.phi = phi <2 x i64> [ %2, %vector.ph ], [ %5, %vector.body ]
%vec.phi7 = phi <2 x i64> [ zeroinitializer, %vector.ph ], [ %6, %vector.body ]
%vec.ind = phi <2 x i32> [ , %vector.ph ], [ %vec.ind.next, %vector.body ]
%step.add = add <2 x i32> %vec.ind, splat (i32 2)
%3 = zext <2 x i32> %vec.ind to <2 x i64>, !dbg !27
%4 = zext <2 x i32> %step.add to <2 x i64>, !dbg !27
%5 = add <2 x i64> %vec.phi, %3, !dbg !28
%6 = add <2 x i64> %vec.phi7, %4, !dbg !28
%index.next = add nuw i32 %index, 4, !dbg !26
%vec.ind.next = add <2 x i32> %vec.ind, splat (i32 4)
%7 = icmp eq i32 %index.next, %n.vec, !dbg !25
br i1 %7, label %middle.block, label %vector.body, !dbg !25
middle.block:
%bin.rdx = add <2 x i64> %6, %5, !dbg !25
%8 = tail call i64 @llvm.vector.reduce.add.v2i64(<2 x i64> %bin.rdx), !dbg !25
%cmp.n = icmp eq i32 %1, %n.vec, !dbg !25
br i1 %cmp.n, label %bb3, label %bb2.preheader, !dbg !25
bb3:
%.lcssa = phi i64 [ %8, %middle.block ], [ %9, %bb2 ], !dbg !28
ret i64 %.lcssa, !dbg !32
bb2:
%start1.sroa.0.06 = phi i64 [ %9, %bb2 ], [ %start1.sroa.0.06.ph, %bb2.preheader ]
%i.sroa.0.05 = phi i32 [ %10, %bb2 ], [ %i.sroa.0.05.ph, %bb2.preheader ]
#dbg_value(i64 %start1.sroa.0.06, !17, !DIExpression(), !23)
#dbg_value(i32 %i.sroa.0.05, !19, !DIExpression(), !24)
%_6 = zext i32 %i.sroa.0.05 to i64, !dbg !27
%9 = add i64 %start1.sroa.0.06, %_6, !dbg !28
#dbg_value(i64 %9, !15, !DIExpression(), !22)
#dbg_value(i64 %9, !17, !DIExpression(), !23)
%10 = add i32 %i.sroa.0.05, 1, !dbg !26
#dbg_value(i32 %10, !19, !DIExpression(), !24)
%_4.not = icmp ugt i32 %10, %n, !dbg !25
br i1 %_4.not, label %bb3, label %bb2, !dbg !25
}
declare i64 @llvm.vector.reduce.add.v2i64(<2 x i64>) #1
attributes #0 = { nofree noinline norecurse nosync nounwind nonlazybind memory(none) uwtable "probe-stack"="inline-asm" "target-cpu"="x86-64" }
attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
```
If the compiler can check if the value of `n` can cause wrapping and split loops to the non-wrapping and the always wrapping, it can optimize both of them better:
```rust
#[inline(never)]
pub fn sum(start: u64, n: u32)->u64 {
let mut res = start;
let mut i = 0;
if n == u32::MAX {
while i <= u32::MAX {
res += i as u64;
i += 1;
}
}
else {
while i < n + 1 {
res += i as u64;
i += 1;
}
}
res
}
```
Generated code (way shorter):
```llvm-ir
define noundef i64 @example::sum::h44644981a6eb5cf7(i64 noundef %0, i32 noundef %n) unnamed_addr #0 !dbg !7 {
start:
#dbg_value(i64 %0, !15, !DIExpression(), !22)
#dbg_value(i64 %0, !17, !DIExpression(), !23)
#dbg_value(i32 %n, !16, !DIExpression(), !22)
#dbg_value(i32 0, !19, !DIExpression(), !24)
%1 = icmp eq i32 %n, -1, !dbg !25
br i1 %1, label %bb2, label %bb7, !dbg !25
bb2:
#dbg_value(!DIArgList(i64 poison, i32 poison), !15, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_arg, 1, DW_OP_LLVM_convert, 32, DW_ATE_unsigned, DW_OP_LLVM_convert, 64, DW_ATE_unsigned, DW_OP_plus, DW_OP_stack_value), !22)
#dbg_value(!DIArgList(i64 poison, i32 poison), !17, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_arg, 1, DW_OP_LLVM_convert, 32, DW_ATE_unsigned, DW_OP_LLVM_convert, 64, DW_ATE_unsigned, DW_OP_plus, DW_OP_stack_value), !23)
#dbg_value(i32 poison, !19, !DIExpression(DW_OP_plus_uconst, 1, DW_OP_stack_value), !24)
br label %bb2, !dbg !26
bb7:
#dbg_value(i64 poison, !17, !DIExpression(), !23)
#dbg_value(i64 poison, !19, !DIExpression(), !24)
#dbg_value(!DIArgList(i64 poison, i64 poison), !15, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_arg, 1, DW_OP_plus, DW_OP_stack_value), !22)
%2 = zext i32 %n to i64, !dbg !27
%3 = add i64 %0, %2, !dbg !27
%4 = zext i32 %n to i64, !dbg !27
%5 = add i32 %n, -1, !dbg !27
%6 = zext i32 %5 to i64, !dbg !27
%7 = mul nuw i64 %4, %6, !dbg !27
%8 = lshr i64 %7, 1, !dbg !27
%9 = add i64 %3, %8, !dbg !27
#dbg_value(i64 %9, !17, !DIExpression(), !23)
#dbg_value(i64 %9, !15, !DIExpression(), !22)
ret i64 %9, !dbg !28
}
attributes #0 = { nofree noinline norecurse nosync nounwind nonlazybind memory(none) uwtable "probe-stack"="inline-asm" "target-cpu"="x86-64" }
```
As can be seen, in such case LLVM optimizes both loops way better by removing one loop and replacing them by sequence of operations and another loop by removing all its contents (because it can be sure that the loop never finishes and any changes cannot be observed).
You can see both versions of loops in godbolt:
https://godbolt.org/z/W4KhccbM3
This optimization should check not only `incremented i <= max_of_type_i` but also cases of `decremented i >= min_of_type_i`.
Also, to avoid penalizing code that can use exceptions or other ways to terminate deliberate endless loops, this optimization should apply only if the body of the loop always progresses (so all called functions in loop are `mustprogress nosync nounwind memory(read)`, for example).
Contributor guide
Assessment
This issue has not been assessed yet.