[llvm] [SPIR-V] Add store legalization for ptrcast (PR #135369)

[Steven Perron via llvm-commits]

================
@@ -150,6 +150,95 @@ class SPIRVLegalizePointerCast : public FunctionPass {
     DeadInstructions.push_back(LI);
   }
 
+  // Creates an spv_insertelt instruction (equivalent to llvm's insertelement).
+  Value *makeInsertElement(IRBuilder<> &B, Value *Vector, Value *Element,
+                           unsigned Index) {
+    Type *Int32Ty = Type::getInt32Ty(B.getContext());
+    SmallVector<Type *, 4> Types = {Vector->getType(), Vector->getType(),
+                                    Element->getType(), Int32Ty};
+    SmallVector<Value *> Args = {Vector, Element, B.getInt32(Index)};
+    Instruction *NewI =
+        B.CreateIntrinsic(Intrinsic::spv_insertelt, {Types}, {Args});
+    buildAssignType(B, Vector->getType(), NewI);
+    return NewI;
+  }
+
+  // Creates an spv_extractelt instruction (equivalent to llvm's
+  // extractelement).
+  Value *makeExtractElement(IRBuilder<> &B, Type *ElementType, Value *Vector,
+                            unsigned Index) {
+    Type *Int32Ty = Type::getInt32Ty(B.getContext());
+    SmallVector<Type *, 3> Types = {ElementType, Vector->getType(), Int32Ty};
+    SmallVector<Value *> Args = {Vector, B.getInt32(Index)};
+    Instruction *NewI =
+        B.CreateIntrinsic(Intrinsic::spv_extractelt, {Types}, {Args});
+    buildAssignType(B, ElementType, NewI);
+    return NewI;
+  }
+
+  // Stores the given Src vector operand into the Dst vector, adjusting the size
+  // if required.
+  Value *storeVectorFromVector(IRBuilder<> &B, Value *Src, Value *Dst,
+                               Align Alignment) {
+    FixedVectorType *SrcType = cast<FixedVectorType>(Src->getType());
+    FixedVectorType *DstType =
+        cast<FixedVectorType>(GR->findDeducedElementType(Dst));
+    assert(DstType->getNumElements() >= SrcType->getNumElements());
+
+    LoadInst *LI = B.CreateLoad(DstType, Dst);
+    LI->setAlignment(Alignment);
+    Value *OldValues = LI;
+    buildAssignType(B, OldValues->getType(), OldValues);
+    Value *NewValues = Src;
+
+    for (unsigned I = 0; I < SrcType->getNumElements(); ++I) {
+      Value *Element =
+          makeExtractElement(B, SrcType->getElementType(), NewValues, I);
+      OldValues = makeInsertElement(B, OldValues, Element, I);
+    }
+
+    StoreInst *SI = B.CreateStore(OldValues, Dst);
+    SI->setAlignment(Alignment);
+    return SI;
+  }
+
+  // Stores the given Src value into the first entry of the Dst aggregate.
+  Value *storeToFirstValueAggregate(IRBuilder<> &B, Value *Src, Value *Dst,
+                                    Align Alignment) {
+    SmallVector<Type *, 2> Types = {Dst->getType(), Dst->getType()};
+    SmallVector<Value *, 3> Args{/* isInBounds= */ B.getInt1(true), Dst,
+                                 B.getInt32(0), B.getInt32(0)};
+    auto *GEP = B.CreateIntrinsic(Intrinsic::spv_gep, {Types}, {Args});
+    GR->buildAssignPtr(B, Src->getType(), GEP);
+    StoreInst *SI = B.CreateStore(Src, GEP);
+    SI->setAlignment(Alignment);
+    return SI;
+  }
+
+  // Transforms a store instruction (or SPV intrinsic) using a ptrcast as
+  // operand into a valid logical SPIR-V store with no ptrcast.
+  void transformStore(IRBuilder<> &B, Instruction *BadStore, Value *Src,
+                      Value *Dst, Align Alignment) {
+    Type *ToTy = GR->findDeducedElementType(Dst);
+    Type *FromTy = Src->getType();
+
+    auto *SVT = dyn_cast<FixedVectorType>(FromTy);
+    auto *DST = dyn_cast<StructType>(ToTy);
+    auto *DVT = dyn_cast<FixedVectorType>(ToTy);
+
+    B.SetInsertPoint(BadStore);
+    if (DST && DST->getTypeAtIndex(0u) == FromTy)
+      storeToFirstValueAggregate(B, Src, Dst, Alignment);
----------------
s-perron wrote:

Do we need to recurse on this? What if you have a struct in a struct, and you write to the first element of the nested struct?

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