[PATCH] D137934: [InstCombine] Fold extractelt with select of constants

[Thomas Symalla via Phabricator via llvm-commits]

tsymalla created this revision.
Herald added a subscriber: hiraditya.
Herald added a project: All.
tsymalla requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

An extractelt with a constant index which extracts an element from the
two vector operands of a select can be directly folded into a select.

extractelt (select %cond, <vec1>, <vec2>), %c1, %c2 ->
select %cond, <vec1>[c1], <vec2>[c2]


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D137934

Files:
  llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
  llvm/test/Transforms/InstCombine/extractelement.ll


Index: llvm/test/Transforms/InstCombine/extractelement.ll
===================================================================
--- llvm/test/Transforms/InstCombine/extractelement.ll
+++ llvm/test/Transforms/InstCombine/extractelement.ll
@@ -785,8 +785,7 @@
 
 define i32 @extelt_select_const_operand_vector(i1 %c) {
 ; ANY-LABEL: @extelt_select_const_operand_vector(
-; ANY-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], <3 x i32> <i32 poison, i32 poison, i32 4>, <3 x i32> <i32 poison, i32 poison, i32 7>
-; ANY-NEXT:    [[R:%.*]] = extractelement <3 x i32> [[S]], i64 2
+; ANY-NEXT:    [[R:%.*]] = select i1 [[C:%.*]], i32 4, i32 7
 ; ANY-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
Index: llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
===================================================================
--- llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -391,6 +391,35 @@
                           IndexC->getValue().zextOrTrunc(64));
 }
 
+static SelectInst *foldExtractElementConstSelect(ExtractElementInst &EI,
+                                                 Value *Vec, ConstantInt *Idx) {
+  if (!EI.hasOneUse())
+    return nullptr;
+
+  if (auto *Select = dyn_cast<SelectInst>(Vec);
+      Select != nullptr && Select->hasOneUse()) {
+    auto *TrueVal = dyn_cast<ConstantVector>(Select->getTrueValue());
+    auto *FalseVal = dyn_cast<ConstantVector>(Select->getFalseValue());
+    if (TrueVal && FalseVal) {
+      const uint64_t IndexVal = Idx->getValue().getZExtValue();
+      const uint64_t NumElements = TrueVal->getType()->getNumElements();
+
+      if (IndexVal < NumElements) {
+        auto *TrueConst =
+            dyn_cast<ConstantInt>(TrueVal->getAggregateElement(IndexVal));
+        auto *FalseConst =
+            dyn_cast<ConstantInt>(FalseVal->getAggregateElement(IndexVal));
+
+        if (TrueConst && FalseConst)
+          return SelectInst::Create(Select->getCondition(), TrueConst,
+                                    FalseConst);
+      }
+    }
+  }
+
+  return nullptr;
+}
+
 Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
   Value *SrcVec = EI.getVectorOperand();
   Value *Index = EI.getIndexOperand();
@@ -402,6 +431,11 @@
   // find a previously computed scalar that was inserted into the vector.
   auto *IndexC = dyn_cast<ConstantInt>(Index);
   if (IndexC) {
+    // extractelt (select %cond, <vec1>, <vec2>), %c1, %c2 ->
+    // select %cond, <vec1>[c1], <vec2>[c2]
+    if (auto *Select = foldExtractElementConstSelect(EI, SrcVec, IndexC))
+      return Select;
+
     // Canonicalize type of constant indices to i64 to simplify CSE
     if (auto *NewIdx = getPreferredVectorIndex(IndexC))
       return replaceOperand(EI, 1, NewIdx);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D137934.475102.patch
Type: text/x-patch
Size: 2869 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20221114/5d1d4d8a/attachment.bin>