[DAGCombine] Poison-prove scalarizeExtractedVectorLoad.

extractelement is poison if the index is out-of-bounds, so just
scalarizing the load may introduce an out-of-bounds load, which is UB.

To avoid introducing new UB, we can mask the index so it only contains
valid indices.

Fixes PR50382.

Reviewed By: efriedma

Differential Revision: https://reviews.llvm.org/D103077

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 0838e52..5146f0d 100644 (file)
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -18411,26 +18411,19 @@ SDValue DAGCombiner::scalarizeExtractedVectorLoad(SDNode *EVE, EVT InVecVT,
 
   Alignment = NewAlign;
 
-  SDValue NewPtr = OriginalLoad->getBasePtr();
-  SDValue Offset;
-  EVT PtrType = NewPtr.getValueType();
   MachinePointerInfo MPI;
   SDLoc DL(EVE);
   if (auto *ConstEltNo = dyn_cast<ConstantSDNode>(EltNo)) {
     int Elt = ConstEltNo->getZExtValue();
     unsigned PtrOff = VecEltVT.getSizeInBits() * Elt / 8;
-    Offset = DAG.getConstant(PtrOff, DL, PtrType);
     MPI = OriginalLoad->getPointerInfo().getWithOffset(PtrOff);
   } else {
-    Offset = DAG.getZExtOrTrunc(EltNo, DL, PtrType);
-    Offset = DAG.getNode(
-        ISD::MUL, DL, PtrType, Offset,
-        DAG.getConstant(VecEltVT.getStoreSize(), DL, PtrType));
     // Discard the pointer info except the address space because the memory
     // operand can't represent this new access since the offset is variable.
     MPI = MachinePointerInfo(OriginalLoad->getPointerInfo().getAddrSpace());
   }
-  NewPtr = DAG.getMemBasePlusOffset(NewPtr, Offset, DL);
+  SDValue NewPtr = TLI.getVectorElementPointer(DAG, OriginalLoad->getBasePtr(),
+                                               InVecVT, EltNo);
 
   // The replacement we need to do here is a little tricky: we need to
   // replace an extractelement of a load with a load.

diff --git a/llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll b/llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll
index 18651d7..b3381b6 100644 (file)
--- a/llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll
@@ -6372,7 +6372,8 @@ define i16 @load_single_extract_variable_index_i16(<8 x i16>* %A, i32 %idx) {
 
 define i32 @load_single_extract_variable_index_i32(<4 x i32>* %A, i32 %idx) {
 ; CHECK-LABEL: load_single_extract_variable_index_i32
-; CHECK:       ldr w0, [x0, w1, sxtw #2]
+; CHECK:       and [[IDX:.*]], x1, #0x3
+; CHECK-NEXT:  ldr w0, [x0, [[IDX]], lsl #2]
 ; CHECK-NEXT:  ret
 ;
   %lv = load <4 x i32>, <4 x i32>* %A
@@ -6400,8 +6401,12 @@ define i32 @load_single_extract_variable_index_v3i32_small_align(<3 x i32>* %A,
 
 define i32 @load_single_extract_variable_index_v3i32_default_align(<3 x i32>* %A, i32 %idx) {
 ; CHECK-LABEL: load_single_extract_variable_index_v3i32_default_align
-; CHECK:       ldr w0, [x0, w1, sxtw #2]
-; CHECK-NEXT:  ret
+; CHECK:       sxtw  [[IDX:.*]], w1
+; CHECK-NEXT:  cmp [[IDX]], #2
+; CHECK-NEXT:  mov w[[TMP:.*]], #2
+; CHECK-NEXT:  csel    [[IDX]], [[IDX]], x[[TMP]], lo
+; CHECK-NEXT:  ldr w0, [x0, [[IDX]], lsl #2]
+; CHECK-NEXT: ret
 ;
   %lv = load <3 x i32>, <3 x i32>* %A
   %e = extractelement <3 x i32> %lv, i32 %idx

diff --git a/llvm/test/CodeGen/SystemZ/vec-extract-02.ll b/llvm/test/CodeGen/SystemZ/vec-extract-02.ll
index a87b7d5..ccf84e2 100644 (file)
--- a/llvm/test/CodeGen/SystemZ/vec-extract-02.ll
+++ b/llvm/test/CodeGen/SystemZ/vec-extract-02.ll
@@ -6,7 +6,7 @@
 ; The index must be extended from i32 to i64.
 define i32 @f1(<4 x i32> *%ptr, i32 %index) {
 ; CHECK-LABEL: f1:
-; CHECK: risbgn {{%r[0-5]}}, %r3, 30, 189, 2
+; CHECK: risbgn {{%r[0-5]}}, %r3, 60, 189, 2
 ; CHECK: l %r2,
 ; CHECK: br %r14
   %vec = load <4 x i32>, <4 x i32> *%ptr

diff --git a/llvm/test/CodeGen/X86/vecloadextract.ll b/llvm/test/CodeGen/X86/vecloadextract.ll
index c76e81d..84ef9f5 100755 (executable)
--- a/llvm/test/CodeGen/X86/vecloadextract.ll
+++ b/llvm/test/CodeGen/X86/vecloadextract.ll
@@ -19,9 +19,10 @@ define i32 @const_index(<8 x i32>* %v) {
 
 ; CHECK: name: variable_index
 ; CHECK:  bb.0 (%ir-block.0):
-; CHECK:    [[INDEX:%[0-9]+]]:gr32_nosp = MOV32rm %fixed-stack.0, 1, $noreg, 0, $noreg :: (load 4 from %fixed-stack.0)
+; CHECK:    [[INDEX:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.0, 1, $noreg, 0, $noreg :: (load 4 from %fixed-stack.0)
+; CHECK:    [[MASKED_INDEX:%[0-9]+]]:gr32_nosp = AND32ri8 [[INDEX]], 7, implicit-def dead $eflags
 ; CHECK:    [[POINTER:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.1, 1, $noreg, 0, $noreg :: (load 4 from %fixed-stack.1)
-; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[INDEX]], 0, $noreg :: (load 4)
+; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[MASKED_INDEX]], 0, $noreg :: (load 4)
 ; CHECK:    $eax = COPY [[LOAD]]
 ; CHECK:    RET 0, $eax
 define i32 @variable_index(<8 x i32>* %v, i32 %i) {
@@ -32,9 +33,10 @@ define i32 @variable_index(<8 x i32>* %v, i32 %i) {
 
 ; CHECK: name: variable_index_with_addrspace
 ; CHECK:  bb.0 (%ir-block.0):
-; CHECK:    [[INDEX:%[0-9]+]]:gr32_nosp = MOV32rm %fixed-stack.0, 1, $noreg, 0, $noreg :: (load 4 from %fixed-stack.0)
+; CHECK:    [[INDEX:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.0, 1, $noreg, 0, $noreg :: (load 4 from %fixed-stack.0)
+; CHECK:    [[MASKED_INDEX:%[0-9]+]]:gr32_nosp = AND32ri8 [[INDEX]], 7, implicit-def dead $eflags
 ; CHECK:    [[POINTER:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.1, 1, $noreg, 0, $noreg :: (load 4 from %fixed-stack.1)
-; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[INDEX]], 0, $noreg :: (load 4, addrspace 1)
+; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[MASKED_INDEX]], 0, $noreg :: (load 4, addrspace 1)
 ; CHECK:    $eax = COPY [[LOAD]]
 ; CHECK:    RET 0, $eax
 define i32 @variable_index_with_addrspace(<8 x i32> addrspace(1)* %v, i32 %i) {