[llvm] d1b402b - [InstCombine] Avoid overflow in `foldVecExtTruncToExtElt` (#180414)

[via llvm-commits]

Author: Yingwei Zheng
Date: 2026-02-09T21:09:48+08:00
New Revision: d1b402b61298d7f28578650e452874e51ea818c4

URL: https://github.com/llvm/llvm-project/commit/d1b402b61298d7f28578650e452874e51ea818c4
DIFF: https://github.com/llvm/llvm-project/commit/d1b402b61298d7f28578650e452874e51ea818c4.diff

LOG: [InstCombine] Avoid overflow in `foldVecExtTruncToExtElt` (#180414)

This weird pattern was introduced by LoopVectorize. But it was placed in
an unreachable path, so we cannot assert that the indices are always
valid in InstCombine.
Closes https://github.com/llvm/llvm-project/issues/180233.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
    llvm/test/Transforms/InstCombine/trunc-extractelement.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 0cd2c09726a2d..a75c24bb91219 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -698,6 +698,10 @@ static Instruction *foldVecExtTruncToExtElt(TruncInst &Trunc,
   auto VecElts = VecOpTy->getElementCount();
 
   uint64_t BitCastNumElts = VecElts.getKnownMinValue() * TruncRatio;
+  // Make sure we don't overflow in the calculation of the new index.
+  // (VecOpIdx + 1) * TruncRatio should not overflow.
+  if (Cst->uge(std::numeric_limits<uint64_t>::max() / TruncRatio))
+    return nullptr;
   uint64_t VecOpIdx = Cst->getZExtValue();
   uint64_t NewIdx = IC.getDataLayout().isBigEndian()
                         ? (VecOpIdx + 1) * TruncRatio - 1
@@ -711,17 +715,21 @@ static Instruction *foldVecExtTruncToExtElt(TruncInst &Trunc,
       return nullptr;
 
     uint64_t IdxOfs = ShiftAmount->udiv(DstBits).getZExtValue();
+    // IdxOfs is guaranteed to be less than TruncRatio, so we won't overflow in
+    // the adjustment.
+    assert(IdxOfs < TruncRatio &&
+           "IdxOfs is expected to be less than TruncRatio.");
     NewIdx = IC.getDataLayout().isBigEndian() ? (NewIdx - IdxOfs)
                                               : (NewIdx + IdxOfs);
   }
 
   assert(BitCastNumElts <= std::numeric_limits<uint32_t>::max() &&
-         NewIdx <= std::numeric_limits<uint32_t>::max() && "overflow 32-bits");
+         "overflow 32-bits");
 
   auto *BitCastTo =
       VectorType::get(DstType, BitCastNumElts, VecElts.isScalable());
   Value *BitCast = IC.Builder.CreateBitCast(VecOp, BitCastTo);
-  return ExtractElementInst::Create(BitCast, IC.Builder.getInt32(NewIdx));
+  return ExtractElementInst::Create(BitCast, IC.Builder.getInt64(NewIdx));
 }
 
 /// Funnel/Rotate left/right may occur in a wider type than necessary because of

diff  --git a/llvm/test/Transforms/InstCombine/trunc-extractelement.ll b/llvm/test/Transforms/InstCombine/trunc-extractelement.ll
index ba2d07009d9c7..074b5b35d7f2d 100644
--- a/llvm/test/Transforms/InstCombine/trunc-extractelement.ll
+++ b/llvm/test/Transforms/InstCombine/trunc-extractelement.ll
@@ -296,3 +296,37 @@ define <4 x i64> @PR45314(<4 x i64> %x) {
   %b = bitcast <8 x i32> %s to <4 x i64>
   ret <4 x i64> %b
 }
+
+; Make sure we don't overflow when computing the new index.
+define i16 @test_overflow_idx1(<vscale x 16 x i32> %vec) {
+; ANY-LABEL: @test_overflow_idx1(
+; ANY-NEXT:  entry:
+; ANY-NEXT:    [[EXT:%.*]] = extractelement <vscale x 16 x i32> [[VEC:%.*]], i64 -1
+; ANY-NEXT:    [[TRUNC:%.*]] = trunc i32 [[EXT]] to i16
+; ANY-NEXT:    ret i16 [[TRUNC]]
+;
+entry:
+  %ext = extractelement <vscale x 16 x i32> %vec, i64 -1
+  %trunc = trunc i32 %ext to i16
+  ret i16 %trunc
+}
+
+; Make sure we don't overflow when computing the new index.
+define i16 @test_overflow_idx2(<vscale x 16 x i32> %vec) {
+; LE-LABEL: @test_overflow_idx2(
+; LE-NEXT:  entry:
+; LE-NEXT:    [[TMP0:%.*]] = bitcast <vscale x 16 x i32> [[VEC:%.*]] to <vscale x 32 x i16>
+; LE-NEXT:    [[TRUNC:%.*]] = extractelement <vscale x 32 x i16> [[TMP0]], i64 4294967296
+; LE-NEXT:    ret i16 [[TRUNC]]
+;
+; BE-LABEL: @test_overflow_idx2(
+; BE-NEXT:  entry:
+; BE-NEXT:    [[TMP0:%.*]] = bitcast <vscale x 16 x i32> [[VEC:%.*]] to <vscale x 32 x i16>
+; BE-NEXT:    [[TRUNC:%.*]] = extractelement <vscale x 32 x i16> [[TMP0]], i64 4294967297
+; BE-NEXT:    ret i16 [[TRUNC]]
+;
+entry:
+  %ext = extractelement <vscale x 16 x i32> %vec, i32 2147483648
+  %trunc = trunc i32 %ext to i16
+  ret i16 %trunc
+}