[DAGCombine] Poison-prove scalarizeExtractedVectorLoad.

[fhahn]

Summary

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.

[fhahn]

Adjust failing SystemZ test. It would be great if someone familiar with SystemZ could take a look to double-check this makes sense. Perhaps @jonpa?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

18425

Is there any chance of NumElts not being a pow2 here?

[jonpa]

In D103077#2779949, @fhahn wrote:

Adjust failing SystemZ test. It would be great if someone familiar with SystemZ could take a look to double-check this makes sense. Perhaps @jonpa?

Yes, this looks like a correct test update to me: the changed immediate operand reflects that now only two shifted bits are inserted into %r1 (instead of all 32).

I wonder what the rationale behind this patch is: if the instruction is poison to begin with, then the program is broken and I don't understand how changing the argument value can change that..?

Adding @uweigand as well just in case...

[fhahn]

In D103077#2780921, @jonpa wrote:

In D103077#2779949, @fhahn wrote:

Adjust failing SystemZ test. It would be great if someone familiar with SystemZ could take a look to double-check this makes sense. Perhaps @jonpa?

Yes, this looks like a correct test update to me: the changed immediate operand reflects that now only two shifted bits are inserted into %r1 (instead of all 32).

I wonder what the rationale behind this patch is: if the instruction is poison to begin with, then the program is broken and I don't understand how changing the argument value can change that..?

Whether the whole program is broken (has UB) depends on how the returned value is used in the callers. Returning poison is not UB, unless the function has the noundef attribute.

[uweigand]

In D103077#2780921, @jonpa wrote:

In D103077#2779949, @fhahn wrote:

Adjust failing SystemZ test. It would be great if someone familiar with SystemZ could take a look to double-check this makes sense. Perhaps @jonpa?

Yes, this looks like a correct test update to me: the changed immediate operand reflects that now only two shifted bits are inserted into %r1 (instead of all 32).

I wonder what the rationale behind this patch is: if the instruction is poison to begin with, then the program is broken and I don't understand how changing the argument value can change that..?

Adding @uweigand as well just in case...

This looks correct to me as well, it simply implements the mask to valid index (which can be merged into the existing risbgn instruction).

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

18425

I can't think of anything that would enforce it being a power of two. The usage of getKnownMinValue() here is also wrong: for a scalable vector, we need to clamp the index relative to the actual number of lanes at runtime, not the minimum number of lanes.

TargetLowering::getVectorElementPointer() has all the logic to do this correctly; can we use it here?

[fhahn]

Updated to use TLI.getVectorElementPointer, thanks!

[fhahn]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

18425

Yes, it looks like we can still have non-power-of-2 vectors here. We bail out earlier for scalable vectors.

I adjusted the code to use TargetLowering::getVectorElementPointer() as suggested by Eli! Seems like a great way to simplify the existing logic which I was not aware of, thanks!

[efriedma]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::scalarizeExtractedVectorLoad(SDNode *EVE, EVT InVecVT,
 
   ISD::LoadExtType ExtTy = ResultVT.bitsGT(VecEltVT) ?
     ISD::NON_EXTLOAD : ISD::EXTLOAD;
   if (!TLI.shouldReduceLoadWidth(OriginalLoad, ExtTy, VecEltVT))
     return SDValue();
 
   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.
   // Use ReplaceAllUsesOfValuesWith to do the replacement.
   // Note that this replacement assumes that the extractvalue is the only
   // use of the load; that's okay because we don't want to perform this
   // transformation in other cases anyway.
   SDValue Load;

llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll

 ;
   %lv = load <8 x i16>, <8 x i16>* %A
   %e = extractelement <8 x i16> %lv, i32 %idx
   ret i16 %e
 }
 
 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
   %e = extractelement <4 x i32> %lv, i32 %idx
   ret i32 %e
 }
 
 define i32 @load_single_extract_variable_index_v3i32_small_align(<3 x i32>* %A, i32 %idx) {

 ;
   %lv = load <3 x i32>, <3 x i32>* %A, align 2
   %e = extractelement <3 x i32> %lv, i32 %idx
   ret i32 %e
 }
 
 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
   ret i32 %e
 }
 
 define i32 @load_single_extract_valid_const_index_v3i32(<3 x i32>* %A, i32 %idx) {
 ; CHECK-LABEL: load_single_extract_valid_const_index_v3i32

llvm/test/CodeGen/SystemZ/vec-extract-02.ll

 ; Verify ReplaceExtractVectorEltOfLoadWithNarrowedLoad fixes
 ;
 ; RUN: llc < %s -mtriple=s390x-linux-gnu -mcpu=z13 | FileCheck %s
 
 ; Test a case where a vector extraction can be simplified to a scalar load.
 ; 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
   %res = extractelement <4 x i32> %vec, i32 %index
   ret i32 %res
 }

llvm/test/CodeGen/X86/vecloadextract.ll

 define i32 @const_index(<8 x i32>* %v) {
   %a = load <8 x i32>, <8 x i32>* %v
   %b = extractelement <8 x i32> %a, i32 1
   ret i32 %b
 }
 
 ; 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) {
   %a = load <8 x i32>, <8 x i32>* %v
   %b = extractelement <8 x i32> %a, i32 %i
   ret i32 %b
 }
 
 ; 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) {
   %a = load <8 x i32>, <8 x i32> addrspace(1)* %v
   %b = extractelement <8 x i32> %a, i32 %i
   ret i32 %b
 }