[ConstProp] Handle insertelement constants

[aeubanks]

Summary

Previously ConstantFoldExtractElementInstruction() would only work with
insertelement instructions, not contants. This properly handles
insertelement constants as well.

[efriedma]

I think I'd prefer to keep this code in ConstantFoldExtractElementInstruction, rather than extend getAggregateElement() this way; I think it's unintuitive to be doing constant folding in getAggregateElement().

llvm/lib/IR/Constants.cpp

421 ↗

(On Diff #285212)

Can we also fold in the case where the indexes are known not-equal?

[aeubanks]

Move code to ConstantFoldExtractElementInstruction()

[aeubanks]

In D85865#2214496, @efriedma wrote:

I think I'd prefer to keep this code in ConstantFoldExtractElementInstruction, rather than extend getAggregateElement() this way; I think it's unintuitive to be doing constant folding in getAggregateElement().

Done.

llvm/lib/IR/Constants.cpp

421 ↗

(On Diff #285212)

Do you mean recursively search the vector operand?

[efriedma]

llvm/lib/IR/ConstantFold.cpp

859

CE->getOperand(1) should return a Constant*; you don't need to dyn_cast<> it.

llvm/lib/IR/Constants.cpp

421 ↗

(On Diff #285212)

Sort of. I was thinking of something like return ConstantExpr::getExtractElement(IE->getOperand(0), CIdx);.

[efriedma]

llvm/lib/IR/ConstantFold.cpp

858

Oh, also, please avoid using getZExtValue() on constants that don't have a known width; it asserts the number is less than 2^64. Unfortunately, there isn't any convenient way to do this comparison at the moment; I guess you could use APSInt::isSameValue.

[aeubanks]

Recurse when index is not equal
Don't cast getOperand() to Constant*
Use APInt::eq()

[aeubanks]

llvm/lib/IR/ConstantFold.cpp

858

Is APInt::eq() good?

[efriedma]

llvm/lib/IR/ConstantFold.cpp

858

IDIdx and CIdx might not have the same bitwidth, I think; the index of an insertelement/extractelement can have any type. APInt::eq asserts if the widths aren't equal.

[aeubanks]

Use APSInt::isSameValue()

llvm/lib/IR/ConstantFold.cpp

858

Ah yes, I changed one of the types in the new test and repro'd. Now using APSInt::isSameValue().

[efriedma]

LGTM

llvm/lib/IR/ConstantFold.cpp

if (CE->getOpcode() == Instruction::GetElementPtr) {
           if (!ScalarOp)
             return nullptr;
           Ops.push_back(ScalarOp);
         } else
           Ops.push_back(Op);
       }
       return CE->getWithOperands(Ops, ValVTy->getElementType(), false,
                                  Ops[0]->getType()->getPointerElementType());
+    } else if (CE->getOpcode() == Instruction::InsertElement) {
+      if (const auto *IEIdx = dyn_cast<ConstantInt>(CE->getOperand(2))) {
+        if (APSInt::isSameValue(APSInt(IEIdx->getValue()),
+                                APSInt(CIdx->getValue()))) {
+          return CE->getOperand(1);
+        } else {
+          return ConstantExpr::getExtractElement(CE->getOperand(0), CIdx);
+        }
+      }
     }
   }
 
   // CAZ of type ScalableVectorType and n < CAZ->getMinNumElements() =>
   //   extractelt CAZ, n -> 0
   if (auto *ValSVTy = dyn_cast<ScalableVectorType>(Val->getType())) {
     if (!CIdx->uge(ValSVTy->getMinNumElements())) {
       if (auto *CAZ = dyn_cast<ConstantAggregateZero>(Val))

llvm/test/Transforms/InstSimplify/vscale.ll

 ; CHECK-LABEL: @insert_extract_element_same_vec_idx_2(
 ; CHECK-NEXT:    ret i32 1
 ;
   %v = insertelement <vscale x 4 x i32> undef, i32 1, i64 4
   %r = extractelement <vscale x 4 x i32> %v, i64 4
   ret i32 %r
 }
 
+define i32 @insert_extract_element_same_vec_idx_3() {
+; CHECK-LABEL: @insert_extract_element_same_vec_idx_3(
+; CHECK-NEXT:    ret i32 1
+;
+  %r = extractelement <vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1, i64 4), i64 4
+  ret i32 %r
+}
+
+define i32 @insert_extract_element_same_vec_idx_4() {
+; CHECK-LABEL: @insert_extract_element_same_vec_idx_4(
+; CHECK-NEXT:    ret i32 1
+;
+  %r = extractelement <vscale x 4 x i32> insertelement (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1, i32 4), i32 2, i64 3), i64 4
+  ret i32 %r
+}
+
 ; more complicated expressions
 
 define <vscale x 2 x i1> @cmp_le_smax_always_true(<vscale x 2 x i64> %x) {
 ; CHECK-LABEL: @cmp_le_smax_always_true(
 ; CHECK-NEXT:    ret <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> undef, i1 true, i32 0), <vscale x 2 x i1> undef, <vscale x 2 x i32> zeroinitializer)
    %cmp = icmp sle <vscale x 2 x i64> %x, shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 9223372036854775807, i32 0), <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer)
    ret <vscale x 2 x i1> %cmp
 }