[llvm] Introduce check-first vectorization for early-exit loops. (PR #227201)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-vectorizers

Author: Arjun H Kumar (arjun-harikumar-amd)

<details>
<summary>Changes</summary>

Check-first is a new strategy for vectorizing loops that have
uncountable early exits and also store to memory. The existing
strategies either require the loop to be read-only, or speculate the
whole body and mask its side effects. Check-first instead evaluates the
exit conditions for an entire vector chunk up front and executes the
body only when no lane exits. When an exit fires, the chunk is abandoned
and replayed by the scalar loop, which resumes from the start of that
chunk rather than from the exiting lane.

The vector loop is restructured into a cascade:

-   vector.check       
-   vector.check1..N-1 
-   vector.body        
-   vector.check.exit 

  
 Example:
 IR before:
```
 loop:
    %iv = phi i64 [ 0, %entry ], [ %iv.next, %latch ]
    %x.ptr = getelementptr inbounds [1024 x i32], ptr @<!-- -->x, i64 0, i64 %iv
    %x.val = load i32, ptr %x.ptr, align 4
    %found = icmp ne i32 %x.val, 0
    br i1 %found, label %exit, label %latch

  latch:
    %a.ptr = getelementptr inbounds [1024 x i32], ptr @<!-- -->a, i64 0, i64 %iv
    store i32 1, ptr %a.ptr, align 4
    %iv.next = add nuw nsw i64 %iv, 1
    %done = icmp eq i64 %iv.next, 1024
    br i1 %done, label %exit, label %loop
```

With this patch the vectorized code is emitted as:

  ```
vector.check:               ; preds = %vector.body, %vector.ph
    %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
    %0 = add i64 %index, 1
    %1 = add i64 %index, 2
    %2 = add i64 %index, 3
    %3 = getelementptr inbounds [1024 x i32], ptr @<!-- -->x, i64 0, i64 %index
    %wide.load = load <4 x i32>, ptr %3, align 4
    %4 = icmp ne <4 x i32> %wide.load, zeroinitializer
    %5 = freeze <4 x i1> %4
    %6 = call i1 @<!-- -->llvm.vector.reduce.or.v4i1(<4 x i1> %5)
    br i1 %6, label %vector.check.exit, label %vector.body

  vector.body:                ; preds = %vector.check
    %7 = getelementptr inbounds [1024 x i32], ptr @<!-- -->a, i64 0, i64 %index
    %8 = getelementptr inbounds [1024 x i32], ptr @<!-- -->a, i64 0, i64 %0
    %9 = getelementptr inbounds [1024 x i32], ptr @<!-- -->a, i64 0, i64 %1
    %10 = getelementptr inbounds [1024 x i32], ptr @<!-- -->a, i64 0, i64 %2
    store i32 1, ptr %7, align 4
    store i32 1, ptr %8, align 4
    store i32 1, ptr %9, align 4
    store i32 1, ptr %10, align 4
    %index.next = add nuw i64 %index, 4
    %11 = icmp eq i64 %index.next, 1024
    br i1 %11, label %middle.block, label %vector.check

  vector.check.exit:          ; preds = %vector.check
    br label %scalar.ph

  scalar.ph:                  ; preds = %vector.check.exit
    br label %loop

  loop:                       ; preds = %scalar.ph, %latch
    %iv = phi i64 [ %index, %scalar.ph ], [ %iv.next, %latch ]
```

---
Full diff: https://github.com/llvm/llvm-project/pull/227201.diff


17 Files Affected:

- (modified) llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h (+23) 
- (modified) llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp (+222-3) 
- (modified) llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h (+2) 
- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+82-12) 
- (modified) llvm/lib/Transforms/Vectorize/VPlan.cpp (+31) 
- (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+69) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp (+15-2) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanLowering.cpp (+3) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp (+2) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp (+12-1) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (+433-1) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.h (+3) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp (+6) 
- (added) llvm/test/Transforms/LoopVectorize/check-first-instruction-reorder.ll (+196) 
- (added) llvm/test/Transforms/LoopVectorize/check-first-multi-exit-cascade.ll (+133) 
- (added) llvm/test/Transforms/LoopVectorize/check-first-single-exit.ll (+136) 
- (added) llvm/test/Transforms/LoopVectorize/check-first-state-update.ll (+109) 


``````````diff
The server is unavailable at this time. Please wait a few minutes before you try again.
``````````

</details>


https://github.com/llvm/llvm-project/pull/227201