[llvm] [DirectX] Scalarize `extractelement` and `insertelement` with dynamic indices (PR #141676)

[Deric C. via llvm-commits]

================
@@ -172,38 +183,61 @@ bool DataScalarizerVisitor::visitStoreInst(StoreInst &SI) {
   return false;
 }
 
-static bool replaceDynamicInsertElementInst(InsertElementInst &IEI) {
-  IRBuilder<> Builder(&IEI);
+DataScalarizerVisitor::AllocaAndGEPs
+DataScalarizerVisitor::createArrayFromVector(IRBuilder<> &Builder, Value *Vec,
+                                             const Twine &Name = "") {
+  // If there is already an alloca for this vector, return it
+  auto VA = VectorAllocaMap.find(Vec);
+  if (VA != VectorAllocaMap.end())
+    return VA->second;
 
-  Value *Vec = IEI.getOperand(0);
-  Value *Val = IEI.getOperand(1);
-  Value *Index = IEI.getOperand(2);
-  Type *IndexTy = Index->getType();
+  auto InsertPoint = Builder.GetInsertPoint();
+  Builder.SetInsertPointPastAllocas(Builder.GetInsertBlock()->getParent());
 
   Type *ArrTy = equivalentArrayTypeFromVector(Vec->getType());
-  Value *ArrAlloca = Builder.CreateAlloca(ArrTy);
+  AllocaInst *ArrAlloca =
+      Builder.CreateAlloca(ArrTy, nullptr, Name + ".alloca");
   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(IndexTy, 0), ConstantInt::get(IndexTy, I)});
-    Builder.CreateStore(EE, GEP);
-    GEPs[I] = GEP;
+    Value *EE = Builder.CreateExtractElement(Vec, I, Name + ".extract");
+    GEPs[I] = Builder.CreateInBoundsGEP(
+        ArrTy, ArrAlloca, {Builder.getInt32(0), Builder.getInt32(I)},
+        Name + ".index");
+    Builder.CreateStore(EE, GEPs[I]);
   }
 
-  Value *GEPForStore = Builder.CreateInBoundsGEP(
-      ArrTy, ArrAlloca, {ConstantInt::get(IndexTy, 0), Index});
+  VectorAllocaMap.insert({Vec, {ArrAlloca, GEPs}});
+  Builder.SetInsertPoint(InsertPoint);
+  return {ArrAlloca, GEPs};
+}
+
+bool DataScalarizerVisitor::replaceDynamicInsertElementInst(
+    InsertElementInst &IEI) {
+  IRBuilder<> Builder(&IEI);
+
+  Value *Vec = IEI.getOperand(0);
+  Value *Val = IEI.getOperand(1);
+  Value *Index = IEI.getOperand(2);
+
+  AllocaAndGEPs ArrAllocaAndGEPs =
+      createArrayFromVector(Builder, Vec, IEI.getName());
+  AllocaInst *ArrAlloca = ArrAllocaAndGEPs.first;
+  SmallVector<Value *, 4> &ArrGEPs = ArrAllocaAndGEPs.second;
+
+  Type *ArrTy = ArrAlloca->getAllocatedType();
+  Value *GEPForStore =
+      Builder.CreateInBoundsGEP(ArrTy, ArrAlloca, {Builder.getInt32(0), Index},
+                                IEI.getName() + ".dynindex");
----------------
Icohedron wrote:

Makes sense, but also feels kind of weird to not return the Alloca, given the name of the function. I would need to think of a new name. The createArrayFromVector function would also need a new arg for the index.


https://github.com/llvm/llvm-project/pull/141676