[InstCombine] Convert shuffle with single constant vector element to insertelement.

[ABataev]

Summary

If it is known that LHS or RHS of shufflevector instruction is a ConstantVector and only single element from this constant is used in the result, this shufflevector is transformed into an insertelement instruction:
shufflevector V, C, <v1, v2, .., ci, .., vm> -> insertelement V, C[ci],ci-n

[ABataev]

Fixes an issue with shuffling the elements from the other input. In this case we should retain the shuffle.

[RKSimon]

This looks OK to me (especially in context to D24182). Anybody else have any comments?

[spatel]

test/Transforms/InstCombine/insert-const-shuf.ll

16–29

As-is, this patch will cause a regression here (ie, the chain of inserts does not collapse). But D24182 will fix this? If it's not possible to separate these two patches without causing a regression, then they should be merged together. I'm not able to apply D24182 cleanly after applying this patch, so it may be that something is just not updated?

[ABataev]

test/Transforms/InstCombine/insert-const-shuf.ll

16–29

Yes, D24182 is not updated yet. Yes, it is not possible to separate these 2 patches without this kind of regression. Do you really want to join these 2 patches into one?

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

case Instruction::ShuffleVector: {
     if (TmpV) { I->setOperand(0, TmpV); MadeChange = true; }
 
     APInt UndefElts3(LHSVWidth, 0);
     TmpV = SimplifyDemandedVectorElts(I->getOperand(1), RightDemanded,
                                       UndefElts3, Depth + 1);
     if (TmpV) { I->setOperand(1, TmpV); MadeChange = true; }
 
     bool NewUndefElts = false;
+    unsigned LHSIdx = -1u;
+    unsigned RHSIdx = -1u;
+    bool LHSUniform = true;
+    bool RHSUniform = true;
     for (unsigned i = 0; i < VWidth; i++) {
       unsigned MaskVal = Shuffle->getMaskValue(i);
       if (MaskVal == -1u) {
         UndefElts.setBit(i);
       } else if (!DemandedElts[i]) {
         NewUndefElts = true;
         UndefElts.setBit(i);
       } else if (MaskVal < LHSVWidth) {
         if (UndefElts4[MaskVal]) {
           NewUndefElts = true;
           UndefElts.setBit(i);
+        } else {
+          LHSIdx = LHSIdx == -1u ? MaskVal : LHSVWidth;
+          LHSUniform = LHSUniform && (MaskVal == i);
         }
       } else {
         if (UndefElts3[MaskVal - LHSVWidth]) {
           NewUndefElts = true;
           UndefElts.setBit(i);
+        } else {
+          RHSIdx = RHSIdx == -1u ? MaskVal - LHSVWidth : LHSVWidth;
+          RHSUniform = RHSUniform && (MaskVal - LHSVWidth == i);
         }
       }
     }
 
+    // Try to transform shuffle with constant vector and single element from
+    // this constant vector to single insertelement instruction.
+    // shufflevector V, C, <v1, v2, .., ci, .., vm> ->
+    // insertelement V, C[ci], ci-n
+    if (LHSVWidth == Shuffle->getType()->getNumElements()) {
+      Value *Op = nullptr;
+      Constant *Value = nullptr;
+      unsigned Idx = -1u;
+
+      // Find constant vector wigth the single element in shuffle (LHS or RHS).
+      if (LHSIdx < LHSVWidth && RHSUniform) {
+        if (auto *CV = dyn_cast<ConstantVector>(Shuffle->getOperand(0))) {
+          Op = Shuffle->getOperand(1);
+          Value = CV->getOperand(LHSIdx);
+          Idx = LHSIdx;
+        }
+      }
+      if (RHSIdx < LHSVWidth && LHSUniform) {
+        if (auto *CV = dyn_cast<ConstantVector>(Shuffle->getOperand(1))) {
+          Op = Shuffle->getOperand(0);
+          Value = CV->getOperand(RHSIdx);
+          Idx = RHSIdx;
+        }
+      }
+      // Found constant vector with single element - convert to insertelement.
+      if (Op && Value) {
+        Instruction *New = InsertElementInst::Create(
+            Op, Value, ConstantInt::get(Type::getInt32Ty(I->getContext()), Idx),
+            Shuffle->getName());
+        InsertNewInstWith(New, *Shuffle);
+        return New;
+      }
+    }
     if (NewUndefElts) {
       // Add additional discovered undefs.
       SmallVector<Constant*, 16> Elts;
       for (unsigned i = 0; i < VWidth; ++i) {
         if (UndefElts[i])
           Elts.push_back(UndefValue::get(Type::getInt32Ty(I->getContext())));
         else
           Elts.push_back(ConstantInt::get(Type::getInt32Ty(I->getContext()),

test/Transforms/InstCombine/insert-const-shuf.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -instcombine %s | FileCheck %s
 
 ; Eliminate the insertelement.
 
 define <4 x float> @PR29126(<4 x float> %x) {
 ; CHECK-LABEL: @PR29126(
 ; CHECK-NEXT:    [[INS:%.*]] = shufflevector <4 x float> %x, <4 x float> <float undef, float 1.000000e+00, float 2.000000e+00, float 4.200000e+01>, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
 ; CHECK-NEXT:    ret <4 x float> [[INS]]
 ;
   %shuf = shufflevector <4 x float> %x, <4 x float> <float undef, float 1.0, float 2.0, float undef>, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
   %ins = insertelement <4 x float> %shuf, float 42.0, i32 3
   ret <4 x float> %ins
 }
 
 ; A chain of inserts should collapse.
 
 define <4 x float> @twoInserts(<4 x float> %x) {
 ; CHECK-LABEL: @twoInserts(
-; CHECK-NEXT:    [[INS2:%.*]] = shufflevector <4 x float> %x, <4 x float> <float undef, float 0.000000e+00, float 4.200000e+01, float 1.100000e+01>, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[SHUF1:%.*]] = insertelement <4 x float> %x, float 0.000000e+00, i32 1
+; CHECK-NEXT:    [[INS1:%.*]] = insertelement <4 x float> [[SHUF1]], float 4.200000e+01, i32 2
+; CHECK-NEXT:    [[INS2:%.*]] = insertelement <4 x float> [[INS1]], float 1.100000e+01, i32 3
 ; CHECK-NEXT:    ret <4 x float> [[INS2]]
 ;
   %shuf = shufflevector <4 x float> %x, <4 x float> zeroinitializer, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
   %ins1 = insertelement <4 x float> %shuf, float 42.0, i32 2
   %ins2 = insertelement <4 x float> %ins1, float 11.0, i32 3
   ret <4 x float> %ins2
 }
 
+define <4 x i32> @shuffleRetain(<4 x i32> %base) {
+; CHECK-LABEL: @shuffleRetain(
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> %base, <4 x i32> <i32 undef, i32 undef, i32 undef, i32 1>, <4 x i32> <i32 1, i32 2, i32 undef, i32 7>
+; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
+;
+  %shuf = shufflevector <4 x i32> %base, <4 x i32> <i32 4, i32 3, i32 2, i32 1>, <4 x i32> <i32 1, i32 2, i32 undef, i32 7>
+  ret <4 x i32> %shuf
+}
+
 ; TODO: Transform an arbitrary shuffle with constant into a shuffle that is equivalant to a vector select.
 
 define <4 x float> @disguisedSelect(<4 x float> %x) {
 ; CHECK-LABEL: @disguisedSelect(
 ; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x float> %x, <4 x float> <float undef, float 1.000000e+00, float 2.000000e+00, float undef>, <4 x i32> <i32 undef, i32 6, i32 5, i32 3>
 ; CHECK-NEXT:    [[INS:%.*]] = insertelement <4 x float> [[SHUF]], float 4.000000e+00, i32 0
 ; CHECK-NEXT:    ret <4 x float> [[INS]]
 ;

test/Transforms/InstCombine/x86-insertps.ll

;
   ret <4 x float> %res
 
 }
 
 ; The zero mask overrides the insertion lane, so the second input vector is not used.
 
 define <4 x float> @insertps_0xc1(<4 x float> %v1, <4 x float> %v2) {
 ; CHECK-LABEL: @insertps_0xc1(
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x float> %v1, <4 x float> <float 0.000000e+00, float undef, float undef, float undef>, <4 x i32> <i32 4, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x float> %v1, float 0.000000e+00, i32 0
 ; CHECK-NEXT:    ret <4 x float> [[TMP1]]
 ;
   %res = call <4 x float> @llvm.x86.sse41.insertps(<4 x float> %v1, <4 x float> %v2, i8 193)
   ret <4 x float> %res
 
 }
 
 ; If no zero mask bits are set, convert to a shuffle.