[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

18423

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

18423

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

18423

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,
   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 {
+    // Unless the index is known to be valid, the extract may be poison. Emit a
+    // mask to ensure the index is valid such cases, so no new out-of-bound loads
+    // are introduced.
+    KnownBits EltNoBits = DAG.computeKnownBits(EltNo);
+    unsigned NumElts = InVecVT.getVectorElementCount().getKnownMinValue();
+    if (!EltNoBits.getMaxValue().ult(NumElts))
+      EltNo = DAG.getNode(ISD::AND, DL, EltNo.getValueType(), EltNo,
+                          DAG.getConstant((1 << Log2_64(NumElts)) - 1, DL,
+                                          EltNo.getValueType()));
+
     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());
   }

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_masked_i32(<4 x i32>* %A, i32 %idx) {

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:    [[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:    [[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:    [[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:    [[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:    [[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:    [[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
 }