[flang-commits] [flang] [mlir] [OpenMP] [MLIR] [Flang] Replace all uses of variables in ALLOCATE directive to use new value which is created. (PR #212361)

[Sairudra More via flang-commits]

================
@@ -6289,6 +6289,64 @@ static Value getBaseValueForTypeLookup(Value value) {
   return value;
 }
 
+// Determine the LLVM type whose storage size should be allocated for an
+// OpenMP allocate directive list item. Opaque pointers lose element type, so
+// trace through declare wrappers to the underlying global or stack allocation.
+static llvm::Type *
+getAllocatedLlvmTypeForVariable(Value var,
+                                LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::Type *llvmVarTy = moduleTranslation.convertType(var.getType());
+  if (!llvmVarTy->isPointerTy())
+    return llvmVarTy;
+
+  Value baseVar = getBaseValueForTypeLookup(var);
+  if (Operation *globalOp = getGlobalOpFromValue(baseVar))
+    if (auto gop = dyn_cast<LLVM::GlobalOp>(globalOp))
+      return moduleTranslation.convertType(gop.getGlobalType());
+
+  if (auto allocaOp =
+          dyn_cast_if_present<LLVM::AllocaOp>(baseVar.getDefiningOp()))
+    return moduleTranslation.convertType(allocaOp.getElemType());
+
+  if (llvm::Value *baseLlvm = moduleTranslation.lookupValue(baseVar))
+    if (auto *allocaInst = dyn_cast<llvm::AllocaInst>(baseLlvm))
+      return allocaInst->getAllocatedType();
+
+  return llvmVarTy;
+}
+
+// For dynamically-sized stack allocations, compute the allocation size from
+// the alloca's element count at runtime.
+static std::optional<llvm::Value *>
+getDynamicAllocatedSize(Value var, LLVM::ModuleTranslation &moduleTranslation,
+                        llvm::IRBuilderBase &builder,
+                        const llvm::DataLayout &dataLayout) {
+  Value baseVar = getBaseValueForTypeLookup(var);
+  if (auto allocaOp =
+          dyn_cast_if_present<LLVM::AllocaOp>(baseVar.getDefiningOp())) {
+    if (Value arraySize = allocaOp.getArraySize()) {
+      llvm::Type *elemTy =
+          moduleTranslation.convertType(allocaOp.getElemType());
+      llvm::Value *numElems = moduleTranslation.lookupValue(arraySize);
+      uint64_t elemSize = dataLayout.getTypeStoreSize(elemTy).getFixedValue();
+      return builder.CreateMul(numElems, builder.getInt64(elemSize));
----------------
Saieiei wrote:

`LLVM::AllocaOp` permits an `i32` array-size operand. Here, `numElems` remains `i32` while the size constant is `i64`, so `CreateMul` asserts on verifier-valid MLIR. Please normalize the types or diagnose unsupported sizes and add an` i32`-count test.

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