[VectorCombine] Freeze index unless it is known to be non-poison.

[fhahn]

Summary

If the index itself is already poison, the poison propagates through
instructions clamping the index to a valid range. This still causes
introducing a load of poison, as flagged by Alive2 and pointed out
at 575e2aff5574.

This patch updates the code to freeze the index, unless it is proven to
not be poison.

[nlopes]

LGTM, thanks!

[efriedma]

llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll

153

This is still not right. See https://alive2.llvm.org/ce/z/dqg8HF

[fhahn]

llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll

153

Thanks Eli, I thought I ran Alive2 on the test after the changes but something went wrong. I limited the transform to non-poison indices only for now in 96ca03493ae5.

Not sure if there's much else we can do to make this save using freeze, if we cannot rely on the constrained range.

[efriedma]

llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll

153

The following is legal, in general:

define i32 @src(<4 x i32>* %x, i4 %idx) {
entry:
  %lv = load <4 x i32>, <4 x i32>* %x
  %r = extractelement <4 x i32> %lv, i4 %idx
  ret i32 %r
}

define i32 @tgt(<4 x i32>* %x, i4 %idx) {
entry:
  %idx.freeze = freeze i4 %idx
  %idx.clamped = and i4 %idx.freeze, 3
  %gep = getelementptr inbounds <4 x i32>, <4 x i32>* %x, i32 0, i4 %idx.clamped
  %r = load i32, i32* %gep
  ret i32 %r
}

I haven't thought about when it's profitable.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
 
   if (ScalarizedCost >= OriginalCost)
     return false;
 
   // Replace extracts with narrow scalar loads.
   for (User *U : LI->users()) {
     auto *EI = cast<ExtractElementInst>(U);
     Builder.SetInsertPoint(EI);
-    Value *GEP = Builder.CreateInBoundsGEP(
-        FixedVT, Ptr, {Builder.getInt32(0), EI->getOperand(1)});
+
+    Value *Idx = EI->getOperand(1);
+    if (!isGuaranteedNotToBePoison(Idx, &AC, LI, &DT))
+      Idx = Builder.CreateFreeze(Idx);
+
+    Value *GEP =
+        Builder.CreateInBoundsGEP(FixedVT, Ptr, {Builder.getInt32(0), Idx});
     auto *NewLoad = cast<LoadInst>(Builder.CreateLoad(
         FixedVT->getElementType(), GEP, EI->getName() + ".scalar"));
 
     // Set the alignment for the new load. For index 0, we can use the original
     // alignment. Otherwise choose the common alignment of the load's align and
     // the alignment for the scalar type.
     auto *ConstIdx = dyn_cast<ConstantInt>(EI->getOperand(1));
     if (ConstIdx && ConstIdx->isNullValue())

llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll

 }
 
 declare void @llvm.assume(i1)
 
 define i32 @load_extract_idx_var_i64_known_valid_by_and(<4 x i32>* %x, i64 %idx) {
 ; CHECK-LABEL: @load_extract_idx_var_i64_known_valid_by_and(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = and i64 [[IDX:%.*]], 3
-; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i32 0, i64 [[IDX_CLAMPED]]
-; CHECK-NEXT:    [[R:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i64 [[IDX_CLAMPED]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i32 0, i64 [[TMP0]]
+; CHECK-NEXT:    [[R:%.*]] = load i32, i32* [[TMP1]], align 4
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
 entry:
   %idx.clamped = and i64 %idx, 3
   %lv = load <4 x i32>, <4 x i32>* %x
   %r = extractelement <4 x i32> %lv, i64 %idx.clamped
   ret i32 %r
 }
 
entry:
   %r = extractelement <4 x i32> %lv, i64 %idx.clamped
   ret i32 %r
 }
 
 define i32 @load_extract_idx_var_i64_known_valid_by_urem(<4 x i32>* %x, i64 %idx) {
 ; CHECK-LABEL: @load_extract_idx_var_i64_known_valid_by_urem(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = urem i64 [[IDX:%.*]], 4
-; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i32 0, i64 [[IDX_CLAMPED]]
-; CHECK-NEXT:    [[R:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i64 [[IDX_CLAMPED]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i32 0, i64 [[TMP0]]
+; CHECK-NEXT:    [[R:%.*]] = load i32, i32* [[TMP1]], align 4
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
 entry:
   %idx.clamped = urem i64 %idx, 4
   %lv = load <4 x i32>, <4 x i32>* %x
   %r = extractelement <4 x i32> %lv, i64 %idx.clamped
   ret i32 %r
 }
;
   %res = add i32 %e.0, %e.1
   ret i32 %res
 }
 
 define i32 @load_multiple_extracts_with_variable_indices_large_vector_all_valid_by_and(<16 x i32>* %x, i64 %idx.0, i64 %idx.1) {
 ; CHECK-LABEL: @load_multiple_extracts_with_variable_indices_large_vector_all_valid_by_and(
 ; CHECK-NEXT:    [[IDX_0_CLAMPED:%.*]] = and i64 [[IDX_0:%.*]], 15
 ; CHECK-NEXT:    [[IDX_1_CLAMPED:%.*]] = and i64 [[IDX_1:%.*]], 15
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X:%.*]], i32 0, i64 [[IDX_0_CLAMPED]]
-; CHECK-NEXT:    [[E_0:%.*]] = load i32, i32* [[TMP1]], align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X]], i32 0, i64 [[IDX_1_CLAMPED]]
-; CHECK-NEXT:    [[E_1:%.*]] = load i32, i32* [[TMP2]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i64 [[IDX_0_CLAMPED]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X:%.*]], i32 0, i64 [[TMP1]]
+; CHECK-NEXT:    [[E_0:%.*]] = load i32, i32* [[TMP2]], align 4
+; CHECK-NEXT:    [[TMP3:%.*]] = freeze i64 [[IDX_1_CLAMPED]]
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X]], i32 0, i64 [[TMP3]]
+; CHECK-NEXT:    [[E_1:%.*]] = load i32, i32* [[TMP4]], align 4
 ; CHECK-NEXT:    [[RES:%.*]] = add i32 [[E_0]], [[E_1]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %idx.0.clamped = and i64 %idx.0, 15
   %idx.1.clamped = and i64 %idx.1, 15
 
   %lv = load <16 x i32>, <16 x i32>* %x
   %e.0 = extractelement <16 x i32> %lv, i64 %idx.0.clamped
   %e.1 = extractelement <16 x i32> %lv, i64 %idx.1.clamped
   %res = add i32 %e.0, %e.1
   ret i32 %res
 }
 
 define i32 @load_multiple_extracts_with_variable_indices_large_vector_all_valid_by_and_some_noundef(<16 x i32>* %x, i64 %idx.0, i64 noundef %idx.1) {
 ; CHECK-LABEL: @load_multiple_extracts_with_variable_indices_large_vector_all_valid_by_and_some_noundef(
 ; CHECK-NEXT:    [[IDX_0_CLAMPED:%.*]] = and i64 [[IDX_0:%.*]], 15
 ; CHECK-NEXT:    [[IDX_1_CLAMPED:%.*]] = and i64 [[IDX_1:%.*]], 15
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X:%.*]], i32 0, i64 [[IDX_0_CLAMPED]]
-; CHECK-NEXT:    [[E_0:%.*]] = load i32, i32* [[TMP1]], align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X]], i32 0, i64 [[IDX_1_CLAMPED]]
-; CHECK-NEXT:    [[E_1:%.*]] = load i32, i32* [[TMP2]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i64 [[IDX_0_CLAMPED]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X:%.*]], i32 0, i64 [[TMP1]]
+; CHECK-NEXT:    [[E_0:%.*]] = load i32, i32* [[TMP2]], align 4
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds <16 x i32>, <16 x i32>* [[X]], i32 0, i64 [[IDX_1_CLAMPED]]
+; CHECK-NEXT:    [[E_1:%.*]] = load i32, i32* [[TMP3]], align 4
 ; CHECK-NEXT:    [[RES:%.*]] = add i32 [[E_0]], [[E_1]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %idx.0.clamped = and i64 %idx.0, 15
   %idx.1.clamped = and i64 %idx.1, 15
 
   %lv = load <16 x i32>, <16 x i32>* %x
   %e.0 = extractelement <16 x i32> %lv, i64 %idx.0.clamped
   %e.1 = extractelement <16 x i32> %lv, i64 %idx.1.clamped
   %res = add i32 %e.0, %e.1
   ret i32 %res
 }