From 0dfc0d8402f0b90ce0918fab90d72c07f7228e14 Mon Sep 17 00:00:00 2001
From: Arina Neshlyaeva <arina.neshlyaeva@intel.com>
Date: Mon, 19 Sep 2022 15:52:03 -0700
Subject: [PATCH] Use AddressInfo in AddElementOffset

---
 src/ctx.cpp  | 45 ++++++++++++++++++++++-----------------------
 src/ctx.h    | 14 ++++++++------
 src/expr.cpp | 17 +++++++++--------
 src/func.cpp |  8 ++++----
 src/stmt.cpp | 15 ++++++++-------
 5 files changed, 51 insertions(+), 48 deletions(-)

diff --git a/src/ctx.cpp b/src/ctx.cpp
index 229b9cf4f..fa0f758c1 100644
--- a/src/ctx.cpp
+++ b/src/ctx.cpp
@@ -2057,7 +2057,6 @@ const PointerType *FunctionEmitContext::RegularizePointer(const Type *ptrRefType
     else {
         ptrType = CastType<PointerType>(ptrRefType);
     }
-    AssertPos(currentPos, ptrType != NULL);
     return ptrType;
 }
 
@@ -2177,25 +2176,23 @@ llvm::Value *FunctionEmitContext::GetElementPtrInst(llvm::Value *basePtr, llvm::
     }
 }
 
-llvm::Value *FunctionEmitContext::AddElementOffset(llvm::Value *fullBasePtr, int elementNum, const Type *ptrRefType,
+llvm::Value *FunctionEmitContext::AddElementOffset(AddressInfo *fullBasePtrInfo, int elementNum,
                                                    const llvm::Twine &name, const PointerType **resultPtrType) {
+    llvm::Value *fullBasePtr = fullBasePtrInfo->getPointer();
     if (resultPtrType != NULL)
-        AssertPos(currentPos, ptrRefType != NULL);
+        AssertPos(currentPos, fullBasePtrInfo->getISPCType() != NULL);
 
-    llvm::PointerType *llvmPtrType = llvm::dyn_cast<llvm::PointerType>(fullBasePtr->getType());
-    if (llvmPtrType != NULL) {
-        llvm::StructType *llvmStructType = llvm::dyn_cast<llvm::StructType>(llvmPtrType->PTR_ELT_TYPE());
-        if (llvmStructType != NULL && llvmStructType->isSized() == false) {
-            AssertPos(currentPos, m->errorCount > 0);
-            return NULL;
-        }
+    llvm::StructType *llvmStructType = llvm::dyn_cast<llvm::StructType>(fullBasePtrInfo->getElementType());
+    if (llvmStructType != NULL && llvmStructType->isSized() == false) {
+        AssertPos(currentPos, m->errorCount > 0);
+        return NULL;
     }
 
-    // (Unfortunately) it's not required to pass a non-NULL ptrRefType, but
+    // (Unfortunately) it's not required to have a non-NULL ispcType in fullBasePtrInfo, but
     // if we have one, regularize into a pointer type.
     const PointerType *ptrType = NULL;
-    if (ptrRefType != NULL) {
-        ptrType = RegularizePointer(ptrRefType);
+    if (fullBasePtrInfo->getISPCType() != NULL) {
+        ptrType = RegularizePointer(fullBasePtrInfo->getISPCType());
     }
 
     // Similarly, we have to see if the pointer type is a struct to see if
@@ -2230,7 +2227,7 @@ llvm::Value *FunctionEmitContext::AddElementOffset(llvm::Value *fullBasePtr, int
         // If the pointer is uniform, we can use the regular LLVM GEP.
         llvm::Value *offsets[2] = {LLVMInt32(0), LLVMInt32(elementNum)};
         llvm::ArrayRef<llvm::Value *> arrayRef(&offsets[0], &offsets[2]);
-        resultPtr = llvm::GetElementPtrInst::Create(PTYPE(basePtr), basePtr, arrayRef,
+        resultPtr = llvm::GetElementPtrInst::Create(fullBasePtrInfo->getElementType(), basePtr, arrayRef,
                                                     name.isTriviallyEmpty() ? "struct_offset" : name, bblock);
     } else {
         // Otherwise do the math to find the offset and add it to the given
@@ -2380,7 +2377,7 @@ llvm::Value *FunctionEmitContext::loadUniformFromSOA(llvm::Value *ptr, llvm::Val
 
         for (int i = 0; i < ct->GetElementCount(); ++i) {
             const PointerType *eltPtrType;
-            llvm::Value *eltPtr = AddElementOffset(ptr, i, ptrType, "elt_offset", &eltPtrType);
+            llvm::Value *eltPtr = AddElementOffset(new AddressInfo(ptr, ptrType), i, "elt_offset", &eltPtrType);
             llvm::Value *eltValue = LoadInst(eltPtr, mask, eltPtrType, name);
             retValue = InsertInst(retValue, eltValue, i, "set_value");
         }
@@ -2513,7 +2510,7 @@ llvm::Value *FunctionEmitContext::gather(llvm::Value *ptr, const PointerType *pt
 
         for (int i = 0; i < collectionType->GetElementCount(); ++i) {
             const PointerType *eltPtrType;
-            llvm::Value *eltPtr = AddElementOffset(ptr, i, ptrType, "gather_elt_ptr", &eltPtrType);
+            llvm::Value *eltPtr = AddElementOffset(new AddressInfo(ptr, ptrType), i, "gather_elt_ptr", &eltPtrType);
 
             eltPtr = addVaryingOffsetsIfNeeded(eltPtr, eltPtrType);
 
@@ -2776,7 +2773,7 @@ void FunctionEmitContext::maskedStore(llvm::Value *value, llvm::Value *ptr, cons
                 continue;
             }
             llvm::Value *eltValue = ExtractInst(value, i, "value_member");
-            llvm::Value *eltPtr = AddElementOffset(ptr, i, ptrType, "struct_ptr_ptr");
+            llvm::Value *eltPtr = AddElementOffset(new AddressInfo(ptr, ptrType), i, "struct_ptr_ptr");
             const Type *eltPtrType = PointerType::GetUniform(eltType);
             StoreInst(eltValue, eltPtr, mask, eltType, eltPtrType);
         }
@@ -2805,13 +2802,14 @@ void FunctionEmitContext::maskedStore(llvm::Value *value, llvm::Value *ptr, cons
 
             // First, extract the pointer from the slice struct and masked
             // store that.
+            AddressInfo *ptrInfo = new AddressInfo(ptr, ptrType);
             llvm::Value *v0 = ExtractInst(value, 0);
-            llvm::Value *p0 = AddElementOffset(ptr, 0, ptrType);
+            llvm::Value *p0 = AddElementOffset(ptrInfo, 0);
             maskedStore(v0, p0, PointerType::GetUniform(pt->GetAsNonSlice()), mask);
 
             // And then do same for the integer offset
             llvm::Value *v1 = ExtractInst(value, 1);
-            llvm::Value *p1 = AddElementOffset(ptr, 1, ptrType);
+            llvm::Value *p1 = AddElementOffset(ptrInfo, 1);
             const Type *offsetType = AtomicType::VaryingInt32;
             maskedStore(v1, p1, PointerType::GetUniform(offsetType), mask);
 
@@ -2903,7 +2901,7 @@ void FunctionEmitContext::scatter(llvm::Value *value, llvm::Value *ptr, const Ty
 
             // Get the (varying) pointer to the i'th element of the target
             // collection
-            llvm::Value *eltPtr = AddElementOffset(ptr, i, ptrType);
+            llvm::Value *eltPtr = AddElementOffset(new AddressInfo(ptr, ptrType), i);
 
             // The destination element type may be uniform (e.g. if we're
             // scattering to an array of uniform structs).  Thus, we need
@@ -3063,7 +3061,7 @@ void FunctionEmitContext::storeUniformToSOA(llvm::Value *value, llvm::Value *ptr
             llvm::Value *eltValue = ExtractInst(value, i);
             const Type *eltType = ct->GetElementType(i);
             const PointerType *dstEltPtrType;
-            llvm::Value *dstEltPtr = AddElementOffset(ptr, i, ptrType, "slice_offset", &dstEltPtrType);
+            llvm::Value *dstEltPtr = AddElementOffset(new AddressInfo(ptr, ptrType), i, "slice_offset", &dstEltPtrType);
             StoreInst(eltValue, dstEltPtr, mask, eltType, dstEltPtrType);
         }
     } else {
@@ -3655,8 +3653,9 @@ llvm::Value *FunctionEmitContext::LaunchInst(llvm::Value *callee, std::vector<ll
 
     // Copy the values of the parameters into the appropriate place in
     // the argument block
+    AddressInfo *argmemInfo = new AddressInfo(argmem, argStructType);
     for (unsigned int i = 0; i < argVals.size(); ++i) {
-        llvm::Value *ptr = AddElementOffset(argmem, i, NULL, "funarg");
+        llvm::Value *ptr = AddElementOffset(argmemInfo, i, "funarg");
         // don't need to do masked store here, I think
         StoreInst(argVals[i], new AddressInfo(ptr, llvmArgTypes[i]));
     }
@@ -3664,7 +3663,7 @@ llvm::Value *FunctionEmitContext::LaunchInst(llvm::Value *callee, std::vector<ll
     if (argStructType->getNumElements() == argVals.size() + 1) {
         // copy in the mask
         llvm::Value *mask = GetFullMask();
-        llvm::Value *ptr = AddElementOffset(argmem, argVals.size(), NULL, "funarg_mask");
+        llvm::Value *ptr = AddElementOffset(argmemInfo, argVals.size(), "funarg_mask");
         StoreInst(mask, new AddressInfo(ptr, LLVMTypes::MaskType));
     }
 
diff --git a/src/ctx.h b/src/ctx.h
index 1a7375dc9..ecd1839ac 100644
--- a/src/ctx.h
+++ b/src/ctx.h
@@ -67,6 +67,9 @@ class AddressInfo {
     // Return the type of the values stored in this address.
     llvm::Type *getElementType() const { return elementType; }
 
+    // Return the ISPC type. May be NULL.
+    const Type *getISPCType() const { return ispcType; }
+
     // Return the address space that this address resides in.
     unsigned getAddressSpace() const { return getType()->getAddressSpace(); }
 
@@ -449,7 +452,8 @@ class FunctionEmitContext {
         and an integer offset to a slice within that type. */
     llvm::Value *MakeSlicePointer(llvm::Value *ptr, llvm::Value *offset);
 
-    /* Regularize to a standard pointer type. */
+    /* Regularize to a standard pointer type.
+       May return NULL if type is not PointerType or ReferenceType */
     const PointerType *RegularizePointer(const Type *ptrRefType);
 
     /** These GEP methods are generalizations of the standard ones in LLVM;
@@ -467,11 +471,9 @@ class FunctionEmitContext {
     /** This method returns a new pointer that represents offsetting the
         given base pointer to point at the given element number of the
         structure type that the base pointer points to.  (The provided
-        pointer must be a pointer to a structure type.  The ptrType gives
-        the type of the pointer, though it may be NULL if the base pointer
-        is uniform. */
-    llvm::Value *AddElementOffset(llvm::Value *basePtr, int elementNum, const Type *ptrType,
-                                  const llvm::Twine &name = "", const PointerType **resultPtrType = NULL);
+        pointer in AddressInfo must be a pointer to a structure type.) */
+    llvm::Value *AddElementOffset(AddressInfo *basePtrInfo, int elementNum, const llvm::Twine &name = "",
+                                  const PointerType **resultPtrType = NULL);
 
     /** Bool is stored as i8 and <WIDTH x i8> but represented in IR as i1 and
      * <WIDTH x MASK>. This is a helper function to match bool size at storage
diff --git a/src/expr.cpp b/src/expr.cpp
index de10c2cda..68ea8dc50 100644
--- a/src/expr.cpp
+++ b/src/expr.cpp
@@ -759,7 +759,8 @@ void ispc::InitSymbol(AddressInfo *ptrInfo, const Type *symType, Expr *initExpr,
 
                 llvm::Value *ep;
                 if (CastType<StructType>(symType) != NULL)
-                    ep = ctx->AddElementOffset(ptrInfo->getPointer(), i, NULL, "element");
+                    ep = ctx->AddElementOffset(new AddressInfo(ptrInfo->getPointer(), CastType<StructType>(symType)), i,
+                                               "element");
                 else {
                     ep = ctx->GetElementPtrInst(ptrInfo->getPointer(), LLVMInt32(0), LLVMInt32(i),
                                                 /* Type of aggregate structure */ PointerType::GetUniform(symType),
@@ -5065,7 +5066,6 @@ llvm::Value *VectorMemberExpr::GetValue(FunctionEmitContext *ctx) const {
 
         // Allocate temporary memory to store the result
         AddressInfo *resultPtrInfo = ctx->AllocaInst(memberType, "vector_tmp");
-
         if (resultPtrInfo == NULL) {
             AssertPos(pos, m->errorCount > 0);
             return NULL;
@@ -5082,12 +5082,12 @@ llvm::Value *VectorMemberExpr::GetValue(FunctionEmitContext *ctx) const {
         ctx->SetDebugPos(pos);
         for (size_t i = 0; i < identifier.size(); ++i) {
             char idStr[2] = {identifier[i], '\0'};
-            llvm::Value *elementPtr =
-                ctx->AddElementOffset(basePtr, indices[i], basePtrType, llvm::Twine(basePtr->getName()) + idStr);
+            llvm::Value *elementPtr = ctx->AddElementOffset(new AddressInfo(basePtr, basePtrType), indices[i],
+                                                            llvm::Twine(basePtr->getName()) + idStr);
             llvm::Value *elementValue = ctx->LoadInst(elementPtr, elementMask, elementPtrType);
 
-            llvm::Value *ptmp = ctx->AddElementOffset(resultPtrInfo->getPointer(), i, NULL,
-                                                      llvm::Twine(resultPtrInfo->getPointer()->getName()) + idStr);
+            llvm::Value *ptmp =
+                ctx->AddElementOffset(resultPtrInfo, i, llvm::Twine(resultPtrInfo->getPointer()->getName()) + idStr);
             // TODO: when we have swizzle on bool type, it breaks here on StoreInst.
             // The condition in StoreInst checking that SwitchBoolSize is needed doesn't detect it.
             ctx->StoreInst(elementValue, new AddressInfo(ptmp, exprVectorType->GetElementType()), elementPtrType);
@@ -5218,7 +5218,7 @@ llvm::Value *MemberExpr::GetValue(FunctionEmitContext *ctx) const {
         if (elementNumber == -1)
             return NULL;
 
-        lvalue = ctx->AddElementOffset(ptrInfo->getPointer(), elementNumber, PointerType::GetUniform(exprType));
+        lvalue = ctx->AddElementOffset(ptrInfo, elementNumber);
         lvalueType = PointerType::GetUniform(GetType());
         mask = LLVMMaskAllOn;
     } else {
@@ -5254,7 +5254,8 @@ llvm::Value *MemberExpr::GetLValue(FunctionEmitContext *ctx) const {
 
     const Type *exprLValueType = dereferenceExpr ? exprType : expr->GetLValueType();
     ctx->SetDebugPos(pos);
-    llvm::Value *ptr = ctx->AddElementOffset(basePtr, elementNumber, exprLValueType, basePtr->getName().str().c_str());
+    llvm::Value *ptr = ctx->AddElementOffset(new AddressInfo(basePtr, exprLValueType), elementNumber,
+                                             basePtr->getName().str().c_str());
     if (ptr == NULL) {
         AssertPos(pos, m->errorCount > 0);
         return NULL;
diff --git a/src/func.cpp b/src/func.cpp
index 67338efa0..a9108a055 100644
--- a/src/func.cpp
+++ b/src/func.cpp
@@ -260,7 +260,7 @@ static void lCopyInTaskParameter(int i, AddressInfo *structArgPtrInfo, const std
     sym->storageInfo = ctx->AllocaInst(sym->type, sym->name.c_str());
 
     // get a pointer to the value in the struct
-    llvm::Value *ptr = ctx->AddElementOffset(structArgPtrInfo->getPointer(), i, NULL, sym->name.c_str());
+    llvm::Value *ptr = ctx->AddElementOffset(structArgPtrInfo, i, sym->name.c_str());
 
     // and copy the value from the struct and into the local alloca'ed
     // memory
@@ -317,17 +317,17 @@ void Function::emitCode(FunctionEmitContext *ctx, llvm::Function *function, Sour
 
         std::vector<llvm::Type *> llvmArgTypes = type->LLVMFunctionArgTypes(g->ctx);
         llvm::Type *st = llvm::StructType::get(*g->ctx, llvmArgTypes);
-
+        AddressInfo *stInfo = new AddressInfo(structParamPtr, st);
         // Copy the function parameter values from the structure into local
         // storage
         for (unsigned int i = 0; i < args.size(); ++i)
-            lCopyInTaskParameter(i, new AddressInfo(structParamPtr, st), args, ctx);
+            lCopyInTaskParameter(i, stInfo, args, ctx);
 
         if (type->isUnmasked == false) {
             // Copy in the mask as well.
             int nArgs = (int)args.size();
             // The mask is the last parameter in the argument structure
-            llvm::Value *ptr = ctx->AddElementOffset(structParamPtr, nArgs, NULL, "task_struct_mask");
+            llvm::Value *ptr = ctx->AddElementOffset(stInfo, nArgs, "task_struct_mask");
             llvm::Value *ptrval = ctx->LoadInst(new AddressInfo(ptr, LLVMTypes::MaskType), NULL, "mask");
             ctx->SetFunctionMask(ptrval);
         }
diff --git a/src/stmt.cpp b/src/stmt.cpp
index 106ff7a7d..67a0bbb84 100644
--- a/src/stmt.cpp
+++ b/src/stmt.cpp
@@ -3508,7 +3508,7 @@ static AddressInfo *lEmitPrintArgCode(Expr *expr, FunctionEmitContext *ctx) {
     return new AddressInfo(ptr, LLVMTypes::VoidPointerType);
 }
 
-static bool lProcessPrintArg(Expr *expr, FunctionEmitContext *ctx, llvm::Value *argPtrArray, int offset,
+static bool lProcessPrintArg(Expr *expr, FunctionEmitContext *ctx, AddressInfo *argPtrArray, int offset,
                              std::string &argTypes) {
     if (!expr)
         return false;
@@ -3521,7 +3521,7 @@ static bool lProcessPrintArg(Expr *expr, FunctionEmitContext *ctx, llvm::Value *
     AddressInfo *ptrInfo = lEmitPrintArgCode(expr, ctx);
     if (!ptrInfo)
         return false;
-    llvm::Value *arrayPtr = ctx->AddElementOffset(argPtrArray, offset, NULL);
+    llvm::Value *arrayPtr = ctx->AddElementOffset(argPtrArray, offset);
     ctx->StoreInst(ptrInfo->getPointer(), new AddressInfo(arrayPtr, ptrInfo->getElementType()));
     return true;
 }
@@ -3546,23 +3546,24 @@ std::vector<llvm::Value *> PrintStmt::getDoPrintArgs(FunctionEmitContext *ctx) c
         checkFormatString(format, nArgs, pos);
         // Allocate space for the array of pointers to values to be printed
         llvm::Type *argPtrArrayType = llvm::ArrayType::get(LLVMTypes::VoidPointerType, nArgs);
-        llvm::Value *argPtrArray = ctx->AllocaInst(argPtrArrayType, "print_arg_ptrs")->getPointer();
+        AddressInfo *argPtrArrayInfo = ctx->AllocaInst(argPtrArrayType, "print_arg_ptrs");
         // Store the array pointer as a void **, which is what __do_print()
         // expects
-        doPrintArgs[ARGS_IDX] = ctx->BitCastInst(argPtrArray, llvm::PointerType::get(LLVMTypes::VoidPointerType, 0));
+        doPrintArgs[ARGS_IDX] =
+            ctx->BitCastInst(argPtrArrayInfo->getPointer(), llvm::PointerType::get(LLVMTypes::VoidPointerType, 0));
 
         // Now, for each of the arguments, emit code to evaluate its value
         // and store the value into alloca's storage.  Then store the
-        // pointer to the alloca's storage into argPtrArray.
+        // pointer to the alloca's storage into argPtrArrayInfo.
         if (elist) {
             for (unsigned int i = 0; i < elist->exprs.size(); ++i) {
                 Expr *expr = elist->exprs[i];
-                if (!lProcessPrintArg(expr, ctx, argPtrArray, i, argTypes)) {
+                if (!lProcessPrintArg(expr, ctx, argPtrArrayInfo, i, argTypes)) {
                     return {};
                 }
             }
         } else {
-            if (lProcessPrintArg(values, ctx, argPtrArray, 0, argTypes)) {
+            if (lProcessPrintArg(values, ctx, argPtrArrayInfo, 0, argTypes)) {
                 return {};
             }
         }
