[instcombine] (extelt (inselt Vec, Value, Index), Index) -> Value

[dsanders]

Summary

When Index is variable but still trivially known to be equal we can use Value
from before the insertion, possibly eliminating the vector.

Reverts a functional change from:
Author: Philip Reames <[e-mail]>
Date: Wed Dec 8 12:21:10 2021 -0800

[instcombine] A couple style tweaks to visitExtractElementInst [nfc]

Thanks to Michele Scandale for identifying the bug

[reames]

LGTM

And thank you for adding test coverage.

[paulwalker-arm]

Sorry for the post commit question but I recall similar work being blocked when we wanted to say extractelt(splat(x), idx) -> x but the consensus was it could be an invalid transformation because if idx is out of range then the result is poison. Does this transformation fall into the same trap whereby we're incorrectly removing a source of poison? I'm not specifically seeking a revert I just want to understand the rules because for scalable vectors there are several instances where we are being hampered by such poison preservation.

[nikic]

In D135625#3858141, @paulwalker-arm wrote:

Sorry for the post commit question but I recall similar work being blocked when we wanted to say extractelt(splat(x), idx) -> x but the consensus was it could be an invalid transformation because if idx is out of range then the result is poison. Does this transformation fall into the same trap whereby we're incorrectly removing a source of poison? I'm not specifically seeking a revert I just want to understand the rules because for scalable vectors there are several instances where we are being hampered by such poison preservation.

It's fine to remove a source of poison, you just can't add one.

While the transform implemented here is correct, it is part of the wrong pass. This should be part of InstSimplify in simplifyExtractElementInst(). Can you please move it there?

[paulwalker-arm]

In D135625#3858157, @nikic wrote:

It's fine to remove a source of poison, you just can't add one.

Thanks for the clarification @nikic.

[dsanders]

While the transform implemented here is correct, it is part of the wrong pass. This should be part of InstSimplify in simplifyExtractElementInst(). Can you please move it there?

Sure, I've done that in https://reviews.llvm.org/D136099

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

if (match(SrcVec, m_Cmp(Pred, m_Value(X), m_Value(Y))) &&
     // extelt (cmp X, Y), Index --> cmp (extelt X, Index), (extelt Y, Index)
     Value *E0 = Builder.CreateExtractElement(X, Index);
     Value *E1 = Builder.CreateExtractElement(Y, Index);
     return CmpInst::Create(cast<CmpInst>(SrcVec)->getOpcode(), Pred, E0, E1);
   }
 
   if (auto *I = dyn_cast<Instruction>(SrcVec)) {
     if (auto *IE = dyn_cast<InsertElementInst>(I)) {
+      // If the possibly-variable indices are trivially known to be equal
+      // (because they are the same operand) then use the value that was
+      // inserted directly.
+      if (IE->getOperand(2) == Index)
+        return replaceInstUsesWith(EI, IE->getOperand(1));
       // instsimplify already handled the case where the indices are constants
       // and equal by value, if both are constants, they must not be the same
       // value, extract from the pre-inserted value instead.
       if (isa<Constant>(IE->getOperand(2)) && IndexC)
         return replaceOperand(EI, 0, IE->getOperand(0));
     } else if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
       auto *VecType = cast<VectorType>(GEP->getType());
       ElementCount EC = VecType->getElementCount();

llvm/test/Transforms/InstCombine/vector-inseltextelt.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+; extracting a just-inserted element should yield the original value
+define float @constant_index(<4 x float> %x, float %val) {
+; CHECK-LABEL: @constant_index(
+; CHECK-NEXT:    ret float [[VAL:%.*]]
+;
+  %ins = insertelement <4 x float> %x, float %val, i32 1
+  %ext = extractelement <4 x float> %ins, i32 1
+  ret float %ext
+}
+
+; extracting a just-inserted element should yield the original value even if
+; the index is dynamic
+define float @dynamic_index(<4 x float> %x, float %val, i32 %idx) {
+; CHECK-LABEL: @dynamic_index(
+; CHECK-NEXT:    ret float [[VAL:%.*]]
+;
+  %ins = insertelement <4 x float> %x, float %val, i32 %idx
+  %ext = extractelement <4 x float> %ins, i32 %idx
+  ret float %ext
+}