From 126f90b252509486eab5bfbe06894805b26c8da2 Mon Sep 17 00:00:00 2001 From: Florian Hahn Date: Sat, 29 May 2021 18:50:14 +0100 Subject: [PATCH] [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 --- llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 11 ++--------- .../test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll | 11 ++++++++--- llvm/test/CodeGen/SystemZ/vec-extract-02.ll | 2 +- llvm/test/CodeGen/X86/vecloadextract.ll | 10 ++++++---- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp index 0838e52a0ecb..5146f0dd3d85 100644 --- 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(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 18651d70aa66..b3381b69a3a8 100644 --- 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 a87b7d52771b..ccf84e26ba4b 100644 --- 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 c76e81ddb647..84ef9f59dd40 100755 --- 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) {