[llvm] 4d2ae88 - [InstCombine] Fix invalid scalarization of div

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2024-07-03T11:05:33+02:00
New Revision: 4d2ae88d1617a910ec3a1436ce53579523ac2f97

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

LOG: [InstCombine] Fix invalid scalarization of div

If the binop is not speculatable, and the extract index is out of
range, then scalarizing will perform the operation on a poison
operand, resulting in immediate UB, instead of the previous
poison result.

Fixes https://github.com/llvm/llvm-project/issues/97053.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
    llvm/test/Transforms/InstCombine/scalarization.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 3de56a4038039..753ed55523c84 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -419,6 +419,7 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
   // If extracting a specified index from the vector, see if we can recursively
   // find a previously computed scalar that was inserted into the vector.
   auto *IndexC = dyn_cast<ConstantInt>(Index);
+  bool HasKnownValidIndex = false;
   if (IndexC) {
     // Canonicalize type of constant indices to i64 to simplify CSE
     if (auto *NewIdx = getPreferredVectorIndex(IndexC))
@@ -426,6 +427,7 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
 
     ElementCount EC = EI.getVectorOperandType()->getElementCount();
     unsigned NumElts = EC.getKnownMinValue();
+    HasKnownValidIndex = IndexC->getValue().ult(NumElts);
 
     if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(SrcVec)) {
       Intrinsic::ID IID = II->getIntrinsicID();
@@ -471,8 +473,11 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
     return UnaryOperator::CreateWithCopiedFlags(UO->getOpcode(), E, UO);
   }
 
+  // If the binop is not speculatable, we cannot hoist the extractelement if
+  // it may make the operand poison.
   BinaryOperator *BO;
-  if (match(SrcVec, m_BinOp(BO)) && cheapToScalarize(SrcVec, Index)) {
+  if (match(SrcVec, m_BinOp(BO)) && cheapToScalarize(SrcVec, Index) &&
+      (HasKnownValidIndex || isSafeToSpeculativelyExecute(BO))) {
     // extelt (binop X, Y), Index --> binop (extelt X, Index), (extelt Y, Index)
     Value *X = BO->getOperand(0), *Y = BO->getOperand(1);
     Value *E0 = Builder.CreateExtractElement(X, Index);

diff  --git a/llvm/test/Transforms/InstCombine/scalarization.ll b/llvm/test/Transforms/InstCombine/scalarization.ll
index 781b5bfb7d6ff..2f539ece88320 100644
--- a/llvm/test/Transforms/InstCombine/scalarization.ll
+++ b/llvm/test/Transforms/InstCombine/scalarization.ll
@@ -158,12 +158,11 @@ define i8 @extract_element_binop_splat_variable_index(<4 x i8> %x, i32 %y) {
 
 ; We cannot move the extractelement before the sdiv here, because %z may be
 ; out of range, making the divisor poison and resulting in immediate UB.
-; FIXME: This is a miscompile.
 define i8 @extract_element_binop_splat_variable_index_may_trap(<4 x i8> %x, <4 x i8> %y, i32 %z) {
 ;
 ; CHECK-LABEL: @extract_element_binop_splat_variable_index_may_trap(
-; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x i8> [[Y:%.*]], i32 [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = sdiv i8 42, [[TMP1]]
+; CHECK-NEXT:    [[B:%.*]] = sdiv <4 x i8> <i8 42, i8 42, i8 42, i8 42>, [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = extractelement <4 x i8> [[B]], i32 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %b = sdiv <4 x i8> splat (i8 42), %y