[llvm] [LoopVectorize] Fix an integer narrowing conversion in `getPredBlockCostDivisor(...)` (PR #187605)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Alan Zhao (alanzhao1)

<details>
<summary>Changes</summary>

`LoopVectorizationCostModel::getPredBlockCostDivisor(...)` may return large `uint64_t` values that get coerced to an `unsigned` by `VPCostContext::getPredBlockCostDivisor(...)`, which can cause division by zero.

Fixes #<!-- -->187584

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


3 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanHelpers.h (+1-1) 
- (added) llvm/test/Transforms/LoopVectorize/cost-divisor-overflow.ll (+68) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index c4b5b40578a6e..229b410992c67 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6885,7 +6885,7 @@ bool VPCostContext::skipCostComputation(Instruction *UI, bool IsVector) const {
          SkipCostComputation.contains(UI);
 }
 
-unsigned VPCostContext::getPredBlockCostDivisor(BasicBlock *BB) const {
+uint64_t VPCostContext::getPredBlockCostDivisor(BasicBlock *BB) const {
   return CM.getPredBlockCostDivisor(CostKind, BB);
 }
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
index 113ca8c4d0f7c..0d4a842c700f6 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
@@ -366,7 +366,7 @@ struct VPCostContext {
 
   /// \returns how much the cost of a predicated block should be divided by.
   /// Forwards to LoopVectorizationCostModel::getPredBlockCostDivisor.
-  unsigned getPredBlockCostDivisor(BasicBlock *BB) const;
+  uint64_t getPredBlockCostDivisor(BasicBlock *BB) const;
 
   /// Returns the OperandInfo for \p V, if it is a live-in.
   TargetTransformInfo::OperandValueInfo getOperandInfo(VPValue *V) const;
diff --git a/llvm/test/Transforms/LoopVectorize/cost-divisor-overflow.ll b/llvm/test/Transforms/LoopVectorize/cost-divisor-overflow.ll
new file mode 100644
index 0000000000000..1e836ac07974f
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/cost-divisor-overflow.ll
@@ -0,0 +1,68 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 6
+; RUN: opt < %s -passes=loop-vectorize -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+; Tests that the computed predicate block cost divisor doesn't get coerced to
+; zero and cause a division by zero error.
+; See https://github.com/llvm/llvm-project/issues/187584
+define void @wombat() {
+; CHECK-LABEL: define void @wombat() {
+; CHECK-NEXT:  [[BB:.*]]:
+; CHECK-NEXT:    br label %[[BB2:.*]]
+; CHECK:       [[BB1:.*]]:
+; CHECK-NEXT:    ret void
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i64 [ [[ADD:%.*]], %[[BB6:.*]] ], [ 0, %[[BB]] ]
+; CHECK-NEXT:    [[PHI3:%.*]] = phi i64 [ [[PHI7:%.*]], %[[BB6]] ], [ 0, %[[BB]] ]
+; CHECK-NEXT:    br i1 true, label %[[BB4:.*]], label %[[BB6]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    br i1 false, label %[[BB6]], label %[[BB5:.*]], !prof [[PROF0:![0-9]+]]
+; CHECK:       [[BB5]]:
+; CHECK-NEXT:    [[LOAD:%.*]] = load i64, ptr null, align 8
+; CHECK-NEXT:    [[CALL:%.*]] = call i64 @llvm.smin.i64(i64 [[LOAD]], i64 [[PHI3]])
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[PHI7]] = phi i64 [ [[PHI3]], %[[BB4]] ], [ [[CALL]], %[[BB5]] ], [ [[PHI3]], %[[BB2]] ]
+; CHECK-NEXT:    [[ADD]] = add i64 [[PHI]], 1
+; CHECK-NEXT:    [[ICMP:%.*]] = icmp eq i64 [[PHI]], 1
+; CHECK-NEXT:    br i1 [[ICMP]], label %[[BB1]], label %[[BB2]]
+;
+bb:
+  br label %bb2
+
+bb1:                                              ; preds = %bb6
+  ret void
+
+bb2:                                              ; preds = %bb6, %bb
+  %phi = phi i64 [ %add, %bb6 ], [ 0, %bb ]
+  %phi3 = phi i64 [ %phi7, %bb6 ], [ 0, %bb ]
+  br i1 true, label %bb4, label %bb6
+
+bb4:                                              ; preds = %bb2
+  br i1 false, label %bb6, label %bb5, !prof !0
+
+bb5:                                              ; preds = %bb4
+  %load = load i64, ptr null, align 8
+  %call = call i64 @llvm.smin.i64(i64 %load, i64 %phi3)
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb4, %bb2
+  %phi7 = phi i64 [ %phi3, %bb4 ], [ %call, %bb5 ], [ %phi3, %bb2 ]
+  %add = add i64 %phi, 1
+  %icmp = icmp eq i64 %phi, 1
+  br i1 %icmp, label %bb1, label %bb2
+}
+
+; Function Attrs: nocallback nocreateundeforpoison nofree nosync nounwind speculatable willreturn memory(none)
+declare i64 @llvm.smin.i64(i64, i64) #0
+
+attributes #0 = { nocallback nocreateundeforpoison nofree nosync nounwind speculatable willreturn memory(none) }
+
+!0 = !{!"branch_weights", i32 -2147483648, i32 0}
+;.
+; CHECK: attributes #[[ATTR0:[0-9]+]] = { nocallback nocreateundeforpoison nofree nosync nounwind speculatable willreturn memory(none) }
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 -2147483648, i32 0}
+;.

``````````

</details>


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