[Scalarizer] InsertElement handling w/ constant insert index

[lebedev.ri]

Summary

As it can be clearly seen from the diff, this results in nicer IR.

[lebedev.ri]

llvm/test/Transforms/Scalarizer/basic.ll

542

Previously we'd end up with https://godbolt.org/z/NyExWc

%val0 = load <4 x i32>, <4 x i32>* %src, align 16
%val1 = insertelement <4 x i32> %val0, i32 %repl, i32 3
%val0.i0 = extractelement <4 x i32> %val0, i32 0
%val2.i0 = shl i32 1, %val0.i0
%val0.i1 = extractelement <4 x i32> %val0, i32 1
%val2.i1 = shl i32 2, %val0.i1
%val0.i2 = extractelement <4 x i32> %val0, i32 2
%val2.i2 = shl i32 3, %val0.i2
%val2.i3 = shl i32 4, %repl
%val2.upto0 = insertelement <4 x i32> undef, i32 %val2.i0, i32 0
%val2.upto1 = insertelement <4 x i32> %val2.upto0, i32 %val2.i1, i32 1
%val2.upto2 = insertelement <4 x i32> %val2.upto1, i32 %val2.i2, i32 2
%val2 = insertelement <4 x i32> %val2.upto2, i32 %val2.i3, i32 3
ret <4 x i32> %val2

[lebedev.ri]

Reordered, precommitted tests.

[jdoerfert]

LGTM.

llvm/lib/Transforms/Scalar/Scalarizer.cpp

public:
   bool visitSelectInst(SelectInst &SI);
   bool visitICmpInst(ICmpInst &ICI);
   bool visitFCmpInst(FCmpInst &FCI);
   bool visitUnaryOperator(UnaryOperator &UO);
   bool visitBinaryOperator(BinaryOperator &BO);
   bool visitGetElementPtrInst(GetElementPtrInst &GEPI);
   bool visitCastInst(CastInst &CI);
   bool visitBitCastInst(BitCastInst &BCI);
+  bool visitInsertElementInst(InsertElementInst &IEI);
   bool visitShuffleVectorInst(ShuffleVectorInst &SVI);
   bool visitPHINode(PHINode &PHI);
   bool visitLoadInst(LoadInst &LI);
   bool visitStoreInst(StoreInst &SI);
   bool visitCallInst(CallInst &ICI);
 
 private:
   Scatterer scatter(Instruction *Point, Value *V);
void ScalarizerVisitor::gather(Instruction *Op, const ValueVector &CV) {
   transferMetadataAndIRFlags(Op, CV);
 
   // If we already have a scattered form of Op (created from ExtractElements
   // of Op itself), replace them with the new form.
   ValueVector &SV = Scattered[Op];
   if (!SV.empty()) {
     for (unsigned I = 0, E = SV.size(); I != E; ++I) {
       Value *V = SV[I];
-      if (V == nullptr)
+      if (V == nullptr || SV[I] == CV[I])
         continue;
 
       Instruction *Old = cast<Instruction>(V);
       CV[I]->takeName(Old);
       Old->replaceAllUsesWith(CV[I]);
       Old->eraseFromParent();
     }
   }
for (unsigned ResI = 0; ResI < DstNumElems; ++ResI) {
       Res[ResI] = Builder.CreateBitCast(V, DstVT->getElementType(),
                                         BCI.getName() + ".i" + Twine(ResI));
     }
   }
   gather(&BCI, Res);
   return true;
 }
 
+bool ScalarizerVisitor::visitInsertElementInst(InsertElementInst &IEI) {
+  VectorType *VT = dyn_cast<VectorType>(IEI.getType());
+  if (!VT)
+    return false;
+
+  unsigned NumElems = VT->getNumElements();
+  IRBuilder<> Builder(&IEI);
+  Scatterer Op0 = scatter(&IEI, IEI.getOperand(0));
+  Value *NewElt = IEI.getOperand(1);
+  Value *InsIdx = IEI.getOperand(2);
+
+  ValueVector Res;
+  Res.resize(NumElems);
+
+  if (auto *CI = dyn_cast<ConstantInt>(InsIdx)) {
+    for (unsigned I = 0; I < NumElems; ++I)
+      Res[I] = CI->getValue().getZExtValue() == I ? NewElt : Op0[I];
+  } else {
+    return false;
+  }
+
+  gather(&IEI, Res);
+  return true;
+}
+
 bool ScalarizerVisitor::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   VectorType *VT = dyn_cast<VectorType>(SVI.getType());
   if (!VT)
     return false;
 
   unsigned NumElems = VT->getNumElements();
   Scatterer Op0 = scatter(&SVI, SVI.getOperand(0));
   Scatterer Op1 = scatter(&SVI, SVI.getOperand(1));

llvm/test/Transforms/Scalarizer/basic.ll

 ; Test GEP with vectors.
 define void @f8(<4 x float *> *%dest, <4 x float *> %ptr0, <4 x i32> %i0,
                 float *%other) {
 ; CHECK-LABEL: @f8(
 ; CHECK: %dest.i0 = bitcast <4 x float*>* %dest to float**
 ; CHECK: %dest.i1 = getelementptr float*, float** %dest.i0, i32 1
 ; CHECK: %dest.i2 = getelementptr float*, float** %dest.i0, i32 2
 ; CHECK: %dest.i3 = getelementptr float*, float** %dest.i0, i32 3
+; CHECK: %ptr0.i0 = extractelement <4 x float*> %ptr0, i32 0
+; CHECK: %ptr0.i2 = extractelement <4 x float*> %ptr0, i32 2
+; CHECK: %ptr0.i3 = extractelement <4 x float*> %ptr0, i32 3
 ; CHECK: %i0.i1 = extractelement <4 x i32> %i0, i32 1
 ; CHECK: %i0.i3 = extractelement <4 x i32> %i0, i32 3
-; CHECK: %ptr0.i0 = extractelement <4 x float*> %ptr0, i32 0
 ; CHECK: %val.i0 = getelementptr float, float* %ptr0.i0, i32 100
 ; CHECK: %val.i1 = getelementptr float, float* %other, i32 %i0.i1
-; CHECK: %ptr0.i2 = extractelement <4 x float*> %ptr0, i32 2
 ; CHECK: %val.i2 = getelementptr float, float* %ptr0.i2, i32 100
-; CHECK: %ptr0.i3 = extractelement <4 x float*> %ptr0, i32 3
 ; CHECK: %val.i3 = getelementptr float, float* %ptr0.i3, i32 %i0.i3
 ; CHECK: store float* %val.i0, float** %dest.i0, align 32
 ; CHECK: store float* %val.i1, float** %dest.i1, align 8
 ; CHECK: store float* %val.i2, float** %dest.i2, align 16
 ; CHECK: store float* %val.i3, float** %dest.i3, align 8
 ; CHECK: ret void
   %i1 = insertelement <4 x i32> %i0, i32 100, i32 0
   %i2 = insertelement <4 x i32> %i1, i32 100, i32 2

 define <2 x float> @f22(<2 x float> %x, <2 x float> %y, <2 x float> %z) {
 ; CHECK-LABEL: @f22(
 ; CHECK: %res.i0 = call fast float @llvm.fma.f32
 ; CHECK: %res.i1 = call fast float @llvm.fma.f32
   %res = call fast <2 x float> @llvm.fma.v2f32(<2 x float> %x, <2 x float> %y, <2 x float> %z)
   ret <2 x float> %res
 }
 
 !0 = !{ !"root" }
 !1 = !{ !"set1", !0 }
 !2 = !{ !"set2", !0 }
 !3 = !{ !3, !{!"llvm.loop.parallel_accesses", !13} }
 !4 = !{ float 4.0 }
 !5 = !{ i64 0, i64 8, null }
 !13 = distinct !{}

llvm/test/Transforms/Scalarizer/constant-insertelement.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -scalarizer -scalarize-load-store -dce -S | FileCheck --check-prefixes=ALL %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 
 ; Test that constant inserts are nicely scalarized
 define <4 x i32> @f1(<4 x i32> *%src, i32 %repl, i32 %index) {
 ; ALL-LABEL: @f1(
 ; ALL-NEXT:    [[SRC_I0:%.*]] = bitcast <4 x i32>* [[SRC:%.*]] to i32*
 ; ALL-NEXT:    [[VAL0_I0:%.*]] = load i32, i32* [[SRC_I0]], align 16
 ; ALL-NEXT:    [[SRC_I1:%.*]] = getelementptr i32, i32* [[SRC_I0]], i32 1
 ; ALL-NEXT:    [[VAL0_I1:%.*]] = load i32, i32* [[SRC_I1]], align 4
 ; ALL-NEXT:    [[SRC_I2:%.*]] = getelementptr i32, i32* [[SRC_I0]], i32 2
 ; ALL-NEXT:    [[VAL0_I2:%.*]] = load i32, i32* [[SRC_I2]], align 8
-; ALL-NEXT:    [[SRC_I3:%.*]] = getelementptr i32, i32* [[SRC_I0]], i32 3
-; ALL-NEXT:    [[VAL0_I3:%.*]] = load i32, i32* [[SRC_I3]], align 4
-; ALL-NEXT:    [[VAL0_UPTO0:%.*]] = insertelement <4 x i32> undef, i32 [[VAL0_I0]], i32 0
-; ALL-NEXT:    [[VAL0_UPTO1:%.*]] = insertelement <4 x i32> [[VAL0_UPTO0]], i32 [[VAL0_I1]], i32 1
-; ALL-NEXT:    [[VAL0_UPTO2:%.*]] = insertelement <4 x i32> [[VAL0_UPTO1]], i32 [[VAL0_I2]], i32 2
-; ALL-NEXT:    [[VAL0:%.*]] = insertelement <4 x i32> [[VAL0_UPTO2]], i32 [[VAL0_I3]], i32 3
-; ALL-NEXT:    [[VAL0_I01:%.*]] = extractelement <4 x i32> [[VAL0]], i32 0
-; ALL-NEXT:    [[VAL2_I0:%.*]] = shl i32 1, [[VAL0_I01]]
-; ALL-NEXT:    [[VAL0_I12:%.*]] = extractelement <4 x i32> [[VAL0]], i32 1
-; ALL-NEXT:    [[VAL2_I1:%.*]] = shl i32 2, [[VAL0_I12]]
-; ALL-NEXT:    [[VAL0_I23:%.*]] = extractelement <4 x i32> [[VAL0]], i32 2
-; ALL-NEXT:    [[VAL2_I2:%.*]] = shl i32 3, [[VAL0_I23]]
+; ALL-NEXT:    [[VAL2_I0:%.*]] = shl i32 1, [[VAL0_I0]]
+; ALL-NEXT:    [[VAL2_I1:%.*]] = shl i32 2, [[VAL0_I1]]
+; ALL-NEXT:    [[VAL2_I2:%.*]] = shl i32 3, [[VAL0_I2]]
 ; ALL-NEXT:    [[VAL2_I3:%.*]] = shl i32 4, [[REPL:%.*]]
 ; ALL-NEXT:    [[VAL2_UPTO0:%.*]] = insertelement <4 x i32> undef, i32 [[VAL2_I0]], i32 0
 ; ALL-NEXT:    [[VAL2_UPTO1:%.*]] = insertelement <4 x i32> [[VAL2_UPTO0]], i32 [[VAL2_I1]], i32 1
 ; ALL-NEXT:    [[VAL2_UPTO2:%.*]] = insertelement <4 x i32> [[VAL2_UPTO1]], i32 [[VAL2_I2]], i32 2
 ; ALL-NEXT:    [[VAL2:%.*]] = insertelement <4 x i32> [[VAL2_UPTO2]], i32 [[VAL2_I3]], i32 3
 ; ALL-NEXT:    ret <4 x i32> [[VAL2]]
 ;
   %val0 = load <4 x i32> , <4 x i32> *%src
   %val1 = insertelement <4 x i32> %val0, i32 %repl, i32 3
   %val2 = shl <4 x i32> <i32 1, i32 2, i32 3, i32 4>, %val1
   ret <4 x i32> %val2
 }