Complete scalarization of insertelement with dynamic index

Author: Icohedron

llvm/lib/Target/DirectX/DXILDataScalarization.cpp

@@ -172,36 +172,41 @@ bool DataScalarizerVisitor::visitStoreInst(StoreInst &SI) {
   return false;
 }
 
-// Allocates and populates an array equivalent to the vector operand Vec.
-// Returns the array and the type of the array.
-static std::pair<Value *, Type *>
-allocaArrayFromVector(IRBuilder<> &Builder, Value *Vec, Type *IdxTy) {
+static bool replaceDynamicInsertElementInst(InsertElementInst &IEI) {
+  IRBuilder<> Builder(&IEI);
+
+  Value *Vec = IEI.getOperand(0);
+  Value *Val = IEI.getOperand(1);
+  Value *Index = IEI.getOperand(2);
+  Type *IndexTy = Index->getType();
+
   Type *ArrTy = equivalentArrayTypeFromVector(Vec->getType());
   Value *ArrAlloca = Builder.CreateAlloca(ArrTy);
-  for (unsigned I = 0; I < ArrTy->getArrayNumElements(); ++I) {
+  const uint64_t ArrNumElems = ArrTy->getArrayNumElements();
+
+  SmallVector<Value *, 4> GEPs(ArrNumElems);
+  for (unsigned I = 0; I < ArrNumElems; ++I) {
     Value *EE = Builder.CreateExtractElement(Vec, I);
     Value *GEP = Builder.CreateInBoundsGEP(
         ArrTy, ArrAlloca,
-        {ConstantInt::get(IdxTy, 0), ConstantInt::get(IdxTy, I)});
+        {ConstantInt::get(IndexTy, 0), ConstantInt::get(IndexTy, I)});
     Builder.CreateStore(EE, GEP);
+    GEPs[I] = GEP;
   }
-  return std::make_pair(ArrAlloca, ArrTy);
-}
 
-static bool replaceDynamicInsertElementInst(InsertElementInst &IEI) {
-  IRBuilder<> Builder(&IEI);
+  Value *GEPForStore = Builder.CreateInBoundsGEP(
+      ArrTy, ArrAlloca, {ConstantInt::get(IndexTy, 0), Index});
+  Builder.CreateStore(Val, GEPForStore);
 
-  Value *Vec = IEI.getOperand(0);
-  Value *Val = IEI.getOperand(1);
-  Value *Index = IEI.getOperand(2);
-  Type *IndexTy = Index->getType();
+  Value *NewIEI = PoisonValue::get(Vec->getType());
+  for (unsigned I = 0; I < ArrNumElems; ++I) {
+    Value *GEP = GEPs[I];
+    Value *Load = Builder.CreateLoad(ArrTy->getArrayElementType(), GEP);
+    NewIEI =
+        Builder.CreateInsertElement(NewIEI, Load, ConstantInt::get(IndexTy, I));
+  }
 
-  std::pair<Value *, Type *> Arr = allocaArrayFromVector(Builder, Vec, IndexTy);
-  Value *ArrAlloca = Arr.first;
-  Type *ArrTy = Arr.second;
-  Value *GEP = Builder.CreateInBoundsGEP(ArrTy, ArrAlloca,
-                                         {ConstantInt::get(IndexTy, 0), Index});
-  Builder.CreateStore(Val, GEP);
+  IEI.replaceAllUsesWith(NewIEI);
   IEI.eraseFromParent();
   return true;
 }
@@ -220,10 +225,15 @@ static bool replaceDynamicExtractElementInst(ExtractElementInst &EEI) {
   Value *Index = EEI.getIndexOperand();
   Type *IndexTy = Index->getType();
 
-  std::pair<Value *, Type *> Arr =
-      allocaArrayFromVector(Builder, EEI.getVectorOperand(), IndexTy);
-  Value *ArrAlloca = Arr.first;
-  Type *ArrTy = Arr.second;
+  Type *ArrTy = equivalentArrayTypeFromVector(EEI.getVectorOperandType());
+  Value *ArrAlloca = Builder.CreateAlloca(ArrTy);
+  for (unsigned I = 0; I < ArrTy->getArrayNumElements(); ++I) {
+    Value *EE = Builder.CreateExtractElement(EEI.getVectorOperand(), I);
+    Value *GEP = Builder.CreateInBoundsGEP(
+        ArrTy, ArrAlloca,
+        {ConstantInt::get(IndexTy, 0), ConstantInt::get(IndexTy, I)});
+    Builder.CreateStore(EE, GEP);
+  }
 
   Value *GEP = Builder.CreateInBoundsGEP(ArrTy, ArrAlloca,
                                          {ConstantInt::get(IndexTy, 0), Index});

llvm/test/CodeGen/DirectX/scalarize-dynamic-vector-index.ll

@@ -25,8 +25,8 @@ define float @extract_float_vec_dynamic(<4 x float> %v, i32 %i) {
   ret float %ee
 }
 
-define void @insert_i32_vec_dynamic(<3 x i32> %v, i32 %a, i32 %i) {
-; CHECK-LABEL: define void @insert_i32_vec_dynamic(
+define <3 x i32> @insert_i32_vec_dynamic(<3 x i32> %v, i32 %a, i32 %i) {
+; CHECK-LABEL: define <3 x i32> @insert_i32_vec_dynamic(
 ; CHECK-SAME: <3 x i32> [[V:%.*]], i32 [[A:%.*]], i32 [[I:%.*]]) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = alloca [3 x i32], align 4
 ; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <3 x i32> [[V]], i64 0
@@ -40,10 +40,16 @@ define void @insert_i32_vec_dynamic(<3 x i32> %v, i32 %a, i32 %i) {
 ; CHECK-NEXT:    store i32 [[TMP6]], ptr [[TMP7]], align 4
 ; CHECK-NEXT:    [[TMP8:%.*]] = getelementptr inbounds [3 x i32], ptr [[TMP1]], i32 0, i32 [[I]]
 ; CHECK-NEXT:    store i32 [[A]], ptr [[TMP8]], align 4
-; CHECK-NEXT:    ret void
+; CHECK-NEXT:    [[TMP9:%.*]] = load i32, ptr [[TMP3]], align 4
+; CHECK-NEXT:    [[TMP10:%.*]] = insertelement <3 x i32> poison, i32 [[TMP9]], i32 0
+; CHECK-NEXT:    [[TMP11:%.*]] = load i32, ptr [[TMP5]], align 4
+; CHECK-NEXT:    [[TMP12:%.*]] = insertelement <3 x i32> [[TMP10]], i32 [[TMP11]], i32 1
+; CHECK-NEXT:    [[TMP13:%.*]] = load i32, ptr [[TMP7]], align 4
+; CHECK-NEXT:    [[TMP14:%.*]] = insertelement <3 x i32> [[TMP12]], i32 [[TMP13]], i32 2
+; CHECK-NEXT:    ret <3 x i32> [[TMP14]]
 ;
-  insertelement <3 x i32> %v, i32 %a, i32 %i
-  ret void
+  %ie = insertelement <3 x i32> %v, i32 %a, i32 %i
+  ret <3 x i32> %ie
 }
 
 ; An extractelement with a constant index should not be converted to array form
@@ -58,13 +64,13 @@ define i16 @extract_i16_vec_constant(<4 x i16> %v) {
 }
 
 ; An insertelement with a constant index should not be converted to array form
-define void @insert_half_vec_constant(<2 x half> %v, half %a) {
-; CHECK-LABEL: define void @insert_half_vec_constant(
+define <2 x half> @insert_half_vec_constant(<2 x half> %v, half %a) {
+; CHECK-LABEL: define <2 x half> @insert_half_vec_constant(
 ; CHECK-SAME: <2 x half> [[V:%.*]], half [[A:%.*]]) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x half> [[V]], half [[A]], i32 1
-; CHECK-NEXT:    ret void
+; CHECK-NEXT:    ret <2 x half> [[TMP1]]
 ;
-  insertelement <2 x half> %v, half %a, i32 1
-  ret void
+  %ie = insertelement <2 x half> %v, half %a, i32 1
+  ret <2 x half> %ie
 }