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

Raghu Maddhipatla via flang-commits flang-commits at lists.llvm.org
Fri Aug 14 06:53:52 PDT 2026


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

>From 65370264971f29234b15e9ca0d89aa5e431756d0 Mon Sep 17 00:00:00 2001
From: Raghu Maddhipatla <Raghu.Maddhipatla at amd.com>
Date: Thu, 9 Jul 2026 09:39:47 -0500
Subject: [PATCH 1/6] [OpenMP] [MLIR] [Flang] Replace all uses of variables in
 ALLOCATE directive to use new value which is allocated.

---
 .../mlir/Target/LLVMIR/ModuleTranslation.h    |  3 +++
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 11 ++++++++++
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 10 +++++++++
 .../LLVMIR/openmp-allocate-directive.mlir     | 22 +++++++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index b23f78ff0f98a..f7e66bd9576d9 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -102,6 +102,9 @@ class ModuleTranslation {
     return valueMapping.lookup(value);
   }
 
+  /// Remap old value with new value.
+  void remapAllValuesWith(llvm::Value *oldValue, llvm::Value *newValue);
+
   /// Looks up remapped a list of remapped values.
   SmallVector<llvm::Value *> lookupValues(ValueRange values);
 
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1a861767b5512..ff1f12c4011af 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -10250,6 +10250,17 @@ convertAllocateDirOp(Operation &opInst, llvm::IRBuilderBase &builder,
     }
     // Record the alloc pointer keyed by the MLIR variable value.
     ompIface.registerAllocatedPtr(var, allocCall);
+
+    Value baseVar = getBaseValueForTypeLookup(var);
+    if (llvm::Value *baseLlvm = moduleTranslation.lookupValue(baseVar)) {
+      llvm::Value *boundPtr = builder.CreatePointerBitCastOrAddrSpaceCast(
+          allocCall, baseLlvm->getType());
+      moduleTranslation.remapAllValuesWith(baseLlvm, boundPtr);
+    } else if (llvm::Value *varLlvm = moduleTranslation.lookupValue(var)) {
+      llvm::Value *boundPtr = builder.CreatePointerBitCastOrAddrSpaceCast(
+          allocCall, varLlvm->getType());
+      moduleTranslation.remapAllValuesWith(varLlvm, boundPtr);
+    }
   }
 
   return success();
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 536327e64bb9f..befa111c9015c 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -2486,6 +2486,16 @@ SmallVector<llvm::Value *> ModuleTranslation::lookupValues(ValueRange values) {
   return remapped;
 }
 
+void ModuleTranslation::remapAllValuesWith(llvm::Value *oldValue,
+                                           llvm::Value *newValue) {
+  if (oldValue == newValue)
+    return;
+  oldValue->replaceAllUsesWith(newValue);
+  for (auto &entry : valueMapping)
+    if (entry.second == oldValue)
+      entry.second = newValue;
+}
+
 llvm::OpenMPIRBuilder *ModuleTranslation::getOpenMPBuilder() {
   if (!ompBuilder) {
     ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
diff --git a/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir b/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
index d8975eb512abe..69ab9b6c00616 100644
--- a/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
@@ -115,3 +115,25 @@ llvm.func @test_allocate_array_global() {
   omp.allocate_free (%z : !llvm.ptr) allocator(%alloc6 : i32)
   llvm.return
 }
+
+// -----
+
+// Verifies that loads and stores after omp.allocate_dir use the OMP-allocated
+// pointer rather than the original storage.
+//
+// CHECK-LABEL: define void @test_allocate_use
+// CHECK:   %[[TID:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   %[[ALLOC:.*]] = call ptr @__kmpc_alloc(i32 %[[TID]], i64 8, ptr null)
+// CHECK:   store i32 42, ptr %[[ALLOC]]
+// CHECK:   %[[VAL:.*]] = load i32, ptr %[[ALLOC]]
+// CHECK:   %[[TID_FREE:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   call void @__kmpc_free(i32 %[[TID_FREE]], ptr %[[ALLOC]], ptr null)
+// CHECK:   ret void
+llvm.func @test_allocate_use(%arg0: !llvm.ptr) {
+  omp.allocate_dir (%arg0 : !llvm.ptr)
+  %c42 = llvm.mlir.constant(42 : i32) : i32
+  llvm.store %c42, %arg0 : i32, !llvm.ptr
+  %v = llvm.load %arg0 : !llvm.ptr -> i32
+  omp.allocate_free (%arg0 : !llvm.ptr)
+  llvm.return
+}

>From 42939d1db08c565ac896e6fe57a39bd409099128 Mon Sep 17 00:00:00 2001
From: Raghu Maddhipatla <Raghu.Maddhipatla at amd.com>
Date: Wed, 15 Jul 2026 23:38:15 -0500
Subject: [PATCH 2/6] Addressed implementation support for more test cases.

---
 .../mlir/Target/LLVMIR/ModuleTranslation.h    |  6 +-
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 81 +++++++++++++++----
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 54 ++++++++++++-
 .../LLVMIR/openmp-allocate-directive.mlir     | 52 ++++++++++++
 4 files changed, 172 insertions(+), 21 deletions(-)

diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index f7e66bd9576d9..7c5e941ed5ca5 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -102,8 +102,10 @@ class ModuleTranslation {
     return valueMapping.lookup(value);
   }
 
-  /// Remap old value with new value.
-  void remapAllValuesWith(llvm::Value *oldValue, llvm::Value *newValue);
+  /// Remap old value with new value. When oldValue is a LLVM constant, builder
+  /// is used to materialize any constant-expression users as instructions.
+  void remapAllValuesWith(llvm::Value *oldValue, llvm::Value *newValue,
+                          llvm::IRBuilderBase *builder = nullptr);
 
   /// Looks up remapped a list of remapped values.
   SmallVector<llvm::Value *> lookupValues(ValueRange values);
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index ff1f12c4011af..bba4628a5dbd6 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6920,6 +6920,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));
+    }
+  }
+  if (llvm::Value *baseLlvm = moduleTranslation.lookupValue(baseVar)) {
+    if (auto *allocaInst = dyn_cast<llvm::AllocaInst>(baseLlvm)) {
+      if (allocaInst->isArrayAllocation() &&
+          !llvm::isa<llvm::ArrayType>(allocaInst->getAllocatedType())) {
+        uint64_t elemSize =
+            dataLayout.getTypeStoreSize(allocaInst->getAllocatedType())
+                .getFixedValue();
+        return builder.CreateMul(allocaInst->getArraySize(),
+                                 builder.getInt64(elemSize));
+      }
+    }
+  }
+  return std::nullopt;
+}
+
 static llvm::SmallString<64>
 getDeclareTargetRefPtrSuffix(LLVM::GlobalOp globalOp,
                              llvm::OpenMPIRBuilder &ompBuilder,
@@ -10198,21 +10256,14 @@ convertAllocateDirOp(Operation &opInst, llvm::IRBuilderBase &builder,
   }
 
   for (Value var : vars) {
-    llvm::Type *llvmVarTy = moduleTranslation.convertType(var.getType());
-
-    // Opaque pointers lose element type. Trace to GlobalOp for type
-    // Falls back to llvmVarTy when not from a global.
-    llvm::Type *typeToInspect = llvmVarTy;
-    if (llvmVarTy->isPointerTy()) {
-      Value baseVar = getBaseValueForTypeLookup(var);
-      if (Operation *globalOp = getGlobalOpFromValue(baseVar)) {
-        if (auto gop = dyn_cast<LLVM::GlobalOp>(globalOp))
-          typeToInspect = moduleTranslation.convertType(gop.getGlobalType());
-      }
-    }
+    llvm::Type *typeToInspect =
+        getAllocatedLlvmTypeForVariable(var, moduleTranslation);
 
     llvm::Value *size;
-    if (auto arrTy = llvm::dyn_cast<llvm::ArrayType>(typeToInspect)) {
+    if (std::optional<llvm::Value *> dynamicSize = getDynamicAllocatedSize(
+            var, moduleTranslation, builder, dataLayout)) {
+      size = *dynamicSize;
+    } else if (auto arrTy = llvm::dyn_cast<llvm::ArrayType>(typeToInspect)) {
       llvm::Value *elementCount = builder.getInt64(1);
       llvm::Type *currentType = arrTy;
       while (auto nestedArrTy = llvm::dyn_cast<llvm::ArrayType>(currentType)) {
@@ -10255,11 +10306,11 @@ convertAllocateDirOp(Operation &opInst, llvm::IRBuilderBase &builder,
     if (llvm::Value *baseLlvm = moduleTranslation.lookupValue(baseVar)) {
       llvm::Value *boundPtr = builder.CreatePointerBitCastOrAddrSpaceCast(
           allocCall, baseLlvm->getType());
-      moduleTranslation.remapAllValuesWith(baseLlvm, boundPtr);
+      moduleTranslation.remapAllValuesWith(baseLlvm, boundPtr, &builder);
     } else if (llvm::Value *varLlvm = moduleTranslation.lookupValue(var)) {
       llvm::Value *boundPtr = builder.CreatePointerBitCastOrAddrSpaceCast(
           allocCall, varLlvm->getType());
-      moduleTranslation.remapAllValuesWith(varLlvm, boundPtr);
+      moduleTranslation.remapAllValuesWith(varLlvm, boundPtr, &builder);
     }
   }
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index befa111c9015c..ba69c44cf6caa 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -43,9 +43,11 @@
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/Operator.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -2486,14 +2488,58 @@ SmallVector<llvm::Value *> ModuleTranslation::lookupValues(ValueRange values) {
   return remapped;
 }
 
+static void remapConstantPointerUses(
+    llvm::Constant *oldPtr, llvm::Value *newPtr, llvm::IRBuilderBase &builder,
+    llvm::DenseMap<llvm::Constant *, llvm::Value *> &replacements) {
+  for (llvm::Use &use : llvm::make_early_inc_range(oldPtr->uses())) {
+    if (auto *constantExpr =
+            llvm::dyn_cast<llvm::ConstantExpr>(use.getUser())) {
+      if (constantExpr->getOpcode() == llvm::Instruction::GetElementPtr) {
+        auto *gep = llvm::cast<llvm::GEPOperator>(constantExpr);
+        llvm::SmallVector<llvm::Value *, 4> indices;
+        for (unsigned i = 1, e = constantExpr->getNumOperands(); i < e; ++i)
+          indices.push_back(constantExpr->getOperand(i));
+        llvm::Value *newGEP =
+            builder.CreateGEP(gep->getSourceElementType(), newPtr, indices);
+        replacements[constantExpr] = newGEP;
+        constantExpr->replaceAllUsesWith(newGEP);
+        continue;
+      }
+      llvm::Instruction *newInst = constantExpr->getAsInstruction();
+      builder.Insert(newInst);
+      replacements[constantExpr] = newInst;
+      constantExpr->replaceAllUsesWith(newInst);
+      continue;
+    }
+    use.set(newPtr);
+  }
+}
+
 void ModuleTranslation::remapAllValuesWith(llvm::Value *oldValue,
-                                           llvm::Value *newValue) {
+                                           llvm::Value *newValue,
+                                           llvm::IRBuilderBase *builder) {
   if (oldValue == newValue)
     return;
-  oldValue->replaceAllUsesWith(newValue);
-  for (auto &entry : valueMapping)
-    if (entry.second == oldValue)
+
+  llvm::DenseMap<llvm::Constant *, llvm::Value *> constantReplacements;
+  if (auto *constant = llvm::dyn_cast<llvm::Constant>(oldValue)) {
+    assert(builder &&
+           "IRBuilder required when remapping constant storage pointers");
+    remapConstantPointerUses(constant, newValue, *builder,
+                             constantReplacements);
+  } else {
+    oldValue->replaceAllUsesWith(newValue);
+  }
+
+  for (auto &entry : valueMapping) {
+    if (entry.second == oldValue) {
       entry.second = newValue;
+      continue;
+    }
+    if (auto *constant = llvm::dyn_cast<llvm::Constant>(entry.second))
+      if (llvm::Value *replacement = constantReplacements.lookup(constant))
+        entry.second = replacement;
+  }
 }
 
 llvm::OpenMPIRBuilder *ModuleTranslation::getOpenMPBuilder() {
diff --git a/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir b/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
index 69ab9b6c00616..9cf2894daf8a0 100644
--- a/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
@@ -118,6 +118,25 @@ llvm.func @test_allocate_array_global() {
 
 // -----
 
+// Verifies that array size is correctly calculated from a stack alloca:
+// [10 x i32] = 40 bytes, rounded up to alignment 64 => 64 bytes.
+//
+// CHECK-LABEL: define void @test_allocate_array_stack
+// CHECK:   %[[TID:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   %[[ALLOC:.*]] = call ptr @__kmpc_aligned_alloc(i32 %[[TID]], i64 64, i64 64, ptr null)
+// CHECK:   %[[TID_FREE:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   call void @__kmpc_free(i32 %[[TID_FREE]], ptr %[[ALLOC]], ptr null)
+// CHECK:   ret void
+llvm.func @test_allocate_array_stack() {
+  %one = llvm.mlir.constant(1 : i64) : i64
+  %arr = llvm.alloca %one x !llvm.array<10 x i32> : (i64) -> !llvm.ptr
+  omp.allocate_dir (%arr : !llvm.ptr) align(64)
+  omp.allocate_free (%arr : !llvm.ptr)
+  llvm.return
+}
+
+// -----
+
 // Verifies that loads and stores after omp.allocate_dir use the OMP-allocated
 // pointer rather than the original storage.
 //
@@ -137,3 +156,36 @@ llvm.func @test_allocate_use(%arg0: !llvm.ptr) {
   omp.allocate_free (%arg0 : !llvm.ptr)
   llvm.return
 }
+
+// -----
+
+// Verifies remapping when a global has multiple GEP users (COMMON block shape).
+//
+// CHECK-LABEL: define void @test_allocate_global_gep_users
+// CHECK:   %[[TID:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   %[[ALLOC:.*]] = call ptr @__kmpc_alloc(i32 %[[TID]], i64 8, ptr null)
+// CHECK:   %[[GEP4:.*]] = getelementptr i8, ptr %[[ALLOC]], i64 4
+// CHECK:   store i32 1, ptr %[[ALLOC]], align 4
+// CHECK:   store i32 2, ptr %[[GEP4]], align 4
+// CHECK:   %[[TID_FREE:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   call void @__kmpc_free(i32 %[[TID_FREE]], ptr %[[ALLOC]], ptr null)
+// CHECK:   ret void
+llvm.mlir.global internal @common_like() : !llvm.array<2 x i32> {
+  %0 = llvm.mlir.zero : !llvm.array<2 x i32>
+  llvm.return %0 : !llvm.array<2 x i32>
+}
+
+llvm.func @test_allocate_global_gep_users() {
+  %base = llvm.mlir.addressof @common_like : !llvm.ptr
+  %c0 = llvm.mlir.constant(0 : i64) : i64
+  %c1 = llvm.mlir.constant(1 : i64) : i64
+  %m0 = llvm.getelementptr %base[%c0, %c0] : (!llvm.ptr, i64, i64) -> !llvm.ptr, !llvm.array<2 x i32>
+  %m1 = llvm.getelementptr %base[%c0, %c1] : (!llvm.ptr, i64, i64) -> !llvm.ptr, !llvm.array<2 x i32>
+  omp.allocate_dir (%base : !llvm.ptr)
+  %one = llvm.mlir.constant(1 : i32) : i32
+  %two = llvm.mlir.constant(2 : i32) : i32
+  llvm.store %one, %m0 : i32, !llvm.ptr
+  llvm.store %two, %m1 : i32, !llvm.ptr
+  omp.allocate_free (%base : !llvm.ptr)
+  llvm.return
+}

>From 0ff8840662a6a4aed6437c4bbcbe70a21c54a8a1 Mon Sep 17 00:00:00 2001
From: Raghu Maddhipatla <Raghu.Maddhipatla at amd.com>
Date: Mon, 27 Jul 2026 16:48:21 -0500
Subject: [PATCH 3/6] Add a TODO warning about lowering of SAVE attribute or
 named COMMON block variables are not completely supported on ALLOCATE
 directive.

---
 flang/lib/Lower/OpenMP/OpenMP.cpp             | 35 ++++++++++++++-----
 .../omp-declarative-allocate-save-warning.f90 | 34 ++++++++++++++++++
 2 files changed, 61 insertions(+), 8 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 52edfdfdd738f..a470a9165db99 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2761,19 +2761,40 @@ static void genWsloopClauses(
 //===----------------------------------------------------------------------===//
 // Code generation functions for leaf constructs
 //===----------------------------------------------------------------------===//
-static mlir::omp::AllocateDirOp genAllocateDirOp(
+
+static void genAllocateDirOp(
     lower::AbstractConverter &converter, semantics::SemanticsContext &semaCtx,
     lower::StatementContext &stmtCtx, lower::pft::Evaluation &eval,
     mlir::Location loc, const ObjectList &objects, const ConstructQueue &queue,
     ConstructQueue::const_iterator item) {
+  ObjectList supportedObjects;
+  supportedObjects.reserve(objects.size());
+  for (const Object &object : objects) {
+    const semantics::Symbol *sym = object.sym();
+    assert(sym && "Expected Symbol");
+    const semantics::Symbol &ultimate = sym->GetUltimate();
+    if (semantics::omp::IsCommonBlock(ultimate) ||
+        ultimate.attrs().test(semantics::Attr::SAVE)) {
+      mlir::emitWarning(loc,
+                "TODO : OpenMP declarative ALLOCATE on SAVE variables or "
+                "COMMON blocks is not yet supported, ignoring the ALLOCATE "
+                "directive for '" + sym->name().ToString() + "'");
+      continue;
+    }
+    supportedObjects.push_back(object);
+  }
+
+  if (supportedObjects.empty())
+    return;
+
   llvm::SmallVector<mlir::Value> operandRange;
   mlir::omp::AllocateDirOperands clauseOps;
-  genAllocateClauses(converter, semaCtx, stmtCtx, objects, item->clauses, loc,
-                     operandRange, clauseOps);
+  genAllocateClauses(converter, semaCtx, stmtCtx, supportedObjects,
+                     item->clauses, loc, operandRange, clauseOps);
 
-  auto allocDirOp = mlir::omp::AllocateDirOp::create(
-      converter.getFirOpBuilder(), loc, operandRange, clauseOps.align,
-      clauseOps.allocator);
+  mlir::omp::AllocateDirOp::create(converter.getFirOpBuilder(), loc,
+                                   operandRange, clauseOps.align,
+                                   clauseOps.allocator);
 
   // Register a cleanup at the Fortran scope exit.
   fir::FirOpBuilder *builder = &converter.getFirOpBuilder();
@@ -2782,8 +2803,6 @@ static mlir::omp::AllocateDirOp genAllocateDirOp(
                                        allocator]() {
     mlir::omp::AllocateFreeOp::create(*builder, loc, operandRange, allocator);
   });
-
-  return allocDirOp;
 }
 
 static mlir::omp::BarrierOp
diff --git a/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90 b/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
new file mode 100644
index 0000000000000..f9bd360d78c36
--- /dev/null
+++ b/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
@@ -0,0 +1,34 @@
+! Verify that declarative ALLOCATE on SAVE variables or named COMMON blocks
+! emits a lowering warning and does not generate omp.allocate_dir /
+! omp.allocate_free for those variables.
+
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=51 %s -o - 2>&1 | FileCheck %s
+
+subroutine save_allocate_warning
+  use omp_lib
+  implicit none
+  integer, save :: counter = 100
+
+  !$omp allocate(counter) allocator(omp_default_mem_alloc)
+end subroutine save_allocate_warning
+
+subroutine common_allocate_warning
+  use omp_lib
+  implicit none
+  real :: cb_a, cb_b
+  common /myblock/ cb_a, cb_b
+
+  !$omp allocate(/myblock/) allocator(omp_default_mem_alloc)
+end subroutine common_allocate_warning
+
+! Warnings are emitted during lowering before HLFIR is printed.
+! CHECK: warning: {{.*}}TODO : OpenMP declarative ALLOCATE on SAVE variables or COMMON blocks is not yet supported, ignoring the ALLOCATE directive for 'counter'
+! CHECK: warning: {{.*}}TODO : OpenMP declarative ALLOCATE on SAVE variables or COMMON blocks is not yet supported, ignoring the ALLOCATE directive for 'myblock'
+
+! CHECK-LABEL: func.func @_QPsave_allocate_warning
+! CHECK-NOT: omp.allocate_dir
+! CHECK-NOT: omp.allocate_free
+
+! CHECK-LABEL: func.func @_QPcommon_allocate_warning
+! CHECK-NOT: omp.allocate_dir
+! CHECK-NOT: omp.allocate_free

>From 20117830d1c65addc26acf2733b45a745528b066 Mon Sep 17 00:00:00 2001
From: Raghu Maddhipatla <Raghu.Maddhipatla at amd.com>
Date: Mon, 27 Jul 2026 16:52:12 -0500
Subject: [PATCH 4/6] Apply clang-format changes.

---
 flang/lib/Lower/OpenMP/OpenMP.cpp | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index a470a9165db99..8153533c543bf 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2762,11 +2762,13 @@ static void genWsloopClauses(
 // Code generation functions for leaf constructs
 //===----------------------------------------------------------------------===//
 
-static void genAllocateDirOp(
-    lower::AbstractConverter &converter, semantics::SemanticsContext &semaCtx,
-    lower::StatementContext &stmtCtx, lower::pft::Evaluation &eval,
-    mlir::Location loc, const ObjectList &objects, const ConstructQueue &queue,
-    ConstructQueue::const_iterator item) {
+static void genAllocateDirOp(lower::AbstractConverter &converter,
+                             semantics::SemanticsContext &semaCtx,
+                             lower::StatementContext &stmtCtx,
+                             lower::pft::Evaluation &eval, mlir::Location loc,
+                             const ObjectList &objects,
+                             const ConstructQueue &queue,
+                             ConstructQueue::const_iterator item) {
   ObjectList supportedObjects;
   supportedObjects.reserve(objects.size());
   for (const Object &object : objects) {
@@ -2775,10 +2777,11 @@ static void genAllocateDirOp(
     const semantics::Symbol &ultimate = sym->GetUltimate();
     if (semantics::omp::IsCommonBlock(ultimate) ||
         ultimate.attrs().test(semantics::Attr::SAVE)) {
-      mlir::emitWarning(loc,
-                "TODO : OpenMP declarative ALLOCATE on SAVE variables or "
-                "COMMON blocks is not yet supported, ignoring the ALLOCATE "
-                "directive for '" + sym->name().ToString() + "'");
+      mlir::emitWarning(
+          loc, "TODO : OpenMP declarative ALLOCATE on SAVE variables or "
+               "COMMON blocks is not yet supported, ignoring the ALLOCATE "
+               "directive for '" +
+                   sym->name().ToString() + "'");
       continue;
     }
     supportedObjects.push_back(object);

>From 31420aa982b8c1720970f3884dc271b117ba3fc9 Mon Sep 17 00:00:00 2001
From: Raghu Maddhipatla <Raghu.Maddhipatla at amd.com>
Date: Mon, 27 Jul 2026 23:19:38 -0500
Subject: [PATCH 5/6] Fix windows build test failure.

---
 .../OpenMP/omp-declarative-allocate-save-warning.f90      | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90 b/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
index f9bd360d78c36..a3b145c429395 100644
--- a/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
+++ b/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
@@ -2,23 +2,21 @@
 ! emits a lowering warning and does not generate omp.allocate_dir /
 ! omp.allocate_free for those variables.
 
-! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=51 %s -o - 2>&1 | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir %openmp_flags %s -o - 2>&1 | FileCheck %s
 
 subroutine save_allocate_warning
-  use omp_lib
   implicit none
   integer, save :: counter = 100
 
-  !$omp allocate(counter) allocator(omp_default_mem_alloc)
+  !$omp allocate(counter) allocator(1)
 end subroutine save_allocate_warning
 
 subroutine common_allocate_warning
-  use omp_lib
   implicit none
   real :: cb_a, cb_b
   common /myblock/ cb_a, cb_b
 
-  !$omp allocate(/myblock/) allocator(omp_default_mem_alloc)
+  !$omp allocate(/myblock/) allocator(1)
 end subroutine common_allocate_warning
 
 ! Warnings are emitted during lowering before HLFIR is printed.

>From 3242f7758e02fb2406ccf9005817f50417fcc5d3 Mon Sep 17 00:00:00 2001
From: Raghu Maddhipatla <Raghu.Maddhipatla at amd.com>
Date: Thu, 13 Aug 2026 17:01:21 -0500
Subject: [PATCH 6/6] Addressed review comments.

---
 flang/lib/Lower/OpenMP/OpenMP.cpp             |  2 +-
 .../OpenMP/omp-declarative-allocate-align.f90 | 10 +--
 .../omp-declarative-allocate-save-warning.f90 | 12 ++++
 .../mlir/Target/LLVMIR/ModuleTranslation.h    |  8 +--
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 23 +++----
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 52 +-------------
 .../LLVMIR/openmp-allocate-directive.mlir     | 67 +++++++++++++------
 7 files changed, 83 insertions(+), 91 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 8153533c543bf..fa9643e188170 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2776,7 +2776,7 @@ static void genAllocateDirOp(lower::AbstractConverter &converter,
     assert(sym && "Expected Symbol");
     const semantics::Symbol &ultimate = sym->GetUltimate();
     if (semantics::omp::IsCommonBlock(ultimate) ||
-        ultimate.attrs().test(semantics::Attr::SAVE)) {
+        semantics::IsSaved(ultimate)) {
       mlir::emitWarning(
           loc, "TODO : OpenMP declarative ALLOCATE on SAVE variables or "
                "COMMON blocks is not yet supported, ignoring the ALLOCATE "
diff --git a/flang/test/Lower/OpenMP/omp-declarative-allocate-align.f90 b/flang/test/Lower/OpenMP/omp-declarative-allocate-align.f90
index fdcc4ac1fef20..95e07f23adc20 100644
--- a/flang/test/Lower/OpenMP/omp-declarative-allocate-align.f90
+++ b/flang/test/Lower/OpenMP/omp-declarative-allocate-align.f90
@@ -10,7 +10,7 @@
 
 program main
   integer :: x, y
-  integer :: z(10)
+  integer :: z(5)
   character c
   real :: r
   complex :: cmplx
@@ -35,17 +35,17 @@ program main
 ! CHECK: %[[X_DECL:.*]]:2 = hlfir.declare %[[X_ALLOC]] {uniq_name = "_QFEx"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
 ! CHECK: %[[Y_ALLOC:.*]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFEy"}
 ! CHECK: %[[Y_DECL:.*]]:2 = hlfir.declare %[[Y_ALLOC]] {uniq_name = "_QFEy"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
-! CHECK: %[[Z_REF:.*]] = fir.address_of(@_QFEz) : !fir.ref<!fir.array<10xi32>>
-! CHECK: %[[Z_DECL:.*]]:2 = hlfir.declare %[[Z_REF]]({{.*}}) {uniq_name = "_QFEz"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
+! CHECK: %[[Z_ALLOC:.*]] = fir.alloca !fir.array<5xi32> {bindc_name = "z", uniq_name = "_QFEz"}
+! CHECK: %[[Z_DECL:.*]]:2 = hlfir.declare %[[Z_ALLOC]]({{.*}}) {uniq_name = "_QFEz"} : (!fir.ref<!fir.array<5xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<5xi32>>, !fir.ref<!fir.array<5xi32>>)
 ! CHECK: omp.allocate_dir(%[[X_DECL]]#0 : !fir.ref<i32>) align(16)
 ! CHECK: %[[ALLOC1:.*]] = arith.constant 1 : i32
 ! CHECK: omp.allocate_dir(%[[Y_DECL]]#0 : !fir.ref<i32>) allocator(%[[ALLOC1]] : i32)
 ! CHECK: %[[ALLOC6:.*]] = arith.constant 6 : i32
-! CHECK: omp.allocate_dir(%[[Z_DECL]]#0 : !fir.ref<!fir.array<10xi32>>) align(64) allocator(%[[ALLOC6]] : i32)
+! CHECK: omp.allocate_dir(%[[Z_DECL]]#0 : !fir.ref<!fir.array<5xi32>>) align(64) allocator(%[[ALLOC6]] : i32)
 ! CHECK: %[[ALLOC3:.*]] = arith.constant 3 : i32
 ! CHECK: omp.allocate_dir(%[[C_DECL]]#0, %[[R_DECL]]#0, %[[CMPLX_DECL]]#0 : !fir.ref<!fir.char<1>>, !fir.ref<f32>, !fir.ref<complex<f32>>) align(32) allocator(%[[ALLOC3]] : i32)
 ! CHECK: omp.allocate_free(%[[C_DECL]]#0, %[[R_DECL]]#0, %[[CMPLX_DECL]]#0 : !fir.ref<!fir.char<1>>, !fir.ref<f32>, !fir.ref<complex<f32>>) allocator(%[[ALLOC3]] : i32)
-! CHECK: omp.allocate_free(%[[Z_DECL]]#0 : !fir.ref<!fir.array<10xi32>>) allocator(%[[ALLOC6]] : i32)
+! CHECK: omp.allocate_free(%[[Z_DECL]]#0 : !fir.ref<!fir.array<5xi32>>) allocator(%[[ALLOC6]] : i32)
 ! CHECK: omp.allocate_free(%[[Y_DECL]]#0 : !fir.ref<i32>) allocator(%[[ALLOC1]] : i32)
 ! CHECK: omp.allocate_free(%[[X_DECL]]#0 : !fir.ref<i32>)
 ! CHECK: return
diff --git a/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90 b/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
index a3b145c429395..f70cae6ebeefb 100644
--- a/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
+++ b/flang/test/Lower/OpenMP/omp-declarative-allocate-save-warning.f90
@@ -11,6 +11,13 @@ subroutine save_allocate_warning
   !$omp allocate(counter) allocator(1)
 end subroutine save_allocate_warning
 
+subroutine implicit_save_allocate_warning
+  implicit none
+  integer :: implicit_counter = 100
+
+  !$omp allocate(implicit_counter) allocator(1)
+end subroutine implicit_save_allocate_warning
+
 subroutine common_allocate_warning
   implicit none
   real :: cb_a, cb_b
@@ -21,12 +28,17 @@ end subroutine common_allocate_warning
 
 ! Warnings are emitted during lowering before HLFIR is printed.
 ! CHECK: warning: {{.*}}TODO : OpenMP declarative ALLOCATE on SAVE variables or COMMON blocks is not yet supported, ignoring the ALLOCATE directive for 'counter'
+! CHECK: warning: {{.*}}TODO : OpenMP declarative ALLOCATE on SAVE variables or COMMON blocks is not yet supported, ignoring the ALLOCATE directive for 'implicit_counter'
 ! CHECK: warning: {{.*}}TODO : OpenMP declarative ALLOCATE on SAVE variables or COMMON blocks is not yet supported, ignoring the ALLOCATE directive for 'myblock'
 
 ! CHECK-LABEL: func.func @_QPsave_allocate_warning
 ! CHECK-NOT: omp.allocate_dir
 ! CHECK-NOT: omp.allocate_free
 
+! CHECK-LABEL: func.func @_QPimplicit_save_allocate_warning
+! CHECK-NOT: omp.allocate_dir
+! CHECK-NOT: omp.allocate_free
+
 ! CHECK-LABEL: func.func @_QPcommon_allocate_warning
 ! CHECK-NOT: omp.allocate_dir
 ! CHECK-NOT: omp.allocate_free
diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index 7c5e941ed5ca5..7bfb80a7d96f4 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -102,10 +102,10 @@ class ModuleTranslation {
     return valueMapping.lookup(value);
   }
 
-  /// Remap old value with new value. When oldValue is a LLVM constant, builder
-  /// is used to materialize any constant-expression users as instructions.
-  void remapAllValuesWith(llvm::Value *oldValue, llvm::Value *newValue,
-                          llvm::IRBuilderBase *builder = nullptr);
+  /// Remap old value with new value in the MLIR-to-LLVM value map so later
+  /// translations use the replacement. Existing LLVM instructions are not
+  /// rewritten.
+  void remapAllValuesWith(llvm::Value *oldValue, llvm::Value *newValue);
 
   /// Looks up remapped a list of remapped values.
   SmallVector<llvm::Value *> lookupValues(ValueRange values);
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index bba4628a5dbd6..132eed3e53374 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6924,13 +6924,12 @@ static Value getBaseValueForTypeLookup(Value value) {
 // 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,
+getAllocatedLlvmTypeForVariable(Value var, Value baseVar,
                                 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());
@@ -6948,17 +6947,17 @@ getAllocatedLlvmTypeForVariable(Value var,
 
 // 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);
+static std::optional<llvm::Value *> getDynamicAllocatedSize(
+    Value var, Value baseVar, LLVM::ModuleTranslation &moduleTranslation,
+    llvm::IRBuilderBase &builder, const llvm::DataLayout &dataLayout) {
   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);
+      if (!numElems->getType()->isIntegerTy(64))
+        numElems = builder.CreateZExt(numElems, builder.getInt64Ty());
       uint64_t elemSize = dataLayout.getTypeStoreSize(elemTy).getFixedValue();
       return builder.CreateMul(numElems, builder.getInt64(elemSize));
     }
@@ -10256,12 +10255,13 @@ convertAllocateDirOp(Operation &opInst, llvm::IRBuilderBase &builder,
   }
 
   for (Value var : vars) {
+    Value baseVar = getBaseValueForTypeLookup(var);
     llvm::Type *typeToInspect =
-        getAllocatedLlvmTypeForVariable(var, moduleTranslation);
+        getAllocatedLlvmTypeForVariable(var, baseVar, moduleTranslation);
 
     llvm::Value *size;
     if (std::optional<llvm::Value *> dynamicSize = getDynamicAllocatedSize(
-            var, moduleTranslation, builder, dataLayout)) {
+            var, baseVar, moduleTranslation, builder, dataLayout)) {
       size = *dynamicSize;
     } else if (auto arrTy = llvm::dyn_cast<llvm::ArrayType>(typeToInspect)) {
       llvm::Value *elementCount = builder.getInt64(1);
@@ -10302,15 +10302,14 @@ convertAllocateDirOp(Operation &opInst, llvm::IRBuilderBase &builder,
     // Record the alloc pointer keyed by the MLIR variable value.
     ompIface.registerAllocatedPtr(var, allocCall);
 
-    Value baseVar = getBaseValueForTypeLookup(var);
     if (llvm::Value *baseLlvm = moduleTranslation.lookupValue(baseVar)) {
       llvm::Value *boundPtr = builder.CreatePointerBitCastOrAddrSpaceCast(
           allocCall, baseLlvm->getType());
-      moduleTranslation.remapAllValuesWith(baseLlvm, boundPtr, &builder);
+      moduleTranslation.remapAllValuesWith(baseLlvm, boundPtr);
     } else if (llvm::Value *varLlvm = moduleTranslation.lookupValue(var)) {
       llvm::Value *boundPtr = builder.CreatePointerBitCastOrAddrSpaceCast(
           allocCall, varLlvm->getType());
-      moduleTranslation.remapAllValuesWith(varLlvm, boundPtr, &builder);
+      moduleTranslation.remapAllValuesWith(varLlvm, boundPtr);
     }
   }
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index ba69c44cf6caa..b19ffadfa26e5 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -43,11 +43,9 @@
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
-#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
-#include "llvm/IR/Operator.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -2488,58 +2486,14 @@ SmallVector<llvm::Value *> ModuleTranslation::lookupValues(ValueRange values) {
   return remapped;
 }
 
-static void remapConstantPointerUses(
-    llvm::Constant *oldPtr, llvm::Value *newPtr, llvm::IRBuilderBase &builder,
-    llvm::DenseMap<llvm::Constant *, llvm::Value *> &replacements) {
-  for (llvm::Use &use : llvm::make_early_inc_range(oldPtr->uses())) {
-    if (auto *constantExpr =
-            llvm::dyn_cast<llvm::ConstantExpr>(use.getUser())) {
-      if (constantExpr->getOpcode() == llvm::Instruction::GetElementPtr) {
-        auto *gep = llvm::cast<llvm::GEPOperator>(constantExpr);
-        llvm::SmallVector<llvm::Value *, 4> indices;
-        for (unsigned i = 1, e = constantExpr->getNumOperands(); i < e; ++i)
-          indices.push_back(constantExpr->getOperand(i));
-        llvm::Value *newGEP =
-            builder.CreateGEP(gep->getSourceElementType(), newPtr, indices);
-        replacements[constantExpr] = newGEP;
-        constantExpr->replaceAllUsesWith(newGEP);
-        continue;
-      }
-      llvm::Instruction *newInst = constantExpr->getAsInstruction();
-      builder.Insert(newInst);
-      replacements[constantExpr] = newInst;
-      constantExpr->replaceAllUsesWith(newInst);
-      continue;
-    }
-    use.set(newPtr);
-  }
-}
-
 void ModuleTranslation::remapAllValuesWith(llvm::Value *oldValue,
-                                           llvm::Value *newValue,
-                                           llvm::IRBuilderBase *builder) {
+                                           llvm::Value *newValue) {
   if (oldValue == newValue)
     return;
 
-  llvm::DenseMap<llvm::Constant *, llvm::Value *> constantReplacements;
-  if (auto *constant = llvm::dyn_cast<llvm::Constant>(oldValue)) {
-    assert(builder &&
-           "IRBuilder required when remapping constant storage pointers");
-    remapConstantPointerUses(constant, newValue, *builder,
-                             constantReplacements);
-  } else {
-    oldValue->replaceAllUsesWith(newValue);
-  }
-
-  for (auto &entry : valueMapping) {
-    if (entry.second == oldValue) {
+  for (auto &entry : valueMapping)
+    if (entry.second == oldValue)
       entry.second = newValue;
-      continue;
-    }
-    if (auto *constant = llvm::dyn_cast<llvm::Constant>(entry.second))
-      if (llvm::Value *replacement = constantReplacements.lookup(constant))
-        entry.second = replacement;
-  }
 }
 
 llvm::OpenMPIRBuilder *ModuleTranslation::getOpenMPBuilder() {
diff --git a/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir b/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
index 9cf2894daf8a0..c9e33432918ca 100644
--- a/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-allocate-directive.mlir
@@ -159,33 +159,60 @@ llvm.func @test_allocate_use(%arg0: !llvm.ptr) {
 
 // -----
 
-// Verifies remapping when a global has multiple GEP users (COMMON block shape).
+// Verifies that a use before omp.allocate_dir keeps the original storage while
+// later uses go through the OMP-allocated pointer.
 //
-// CHECK-LABEL: define void @test_allocate_global_gep_users
+// CHECK-LABEL: define void @test_allocate_use_before
+// CHECK-SAME: (ptr %[[ARG0:.*]])
+// CHECK:   %[[PRE:.*]] = load i32, ptr %[[ARG0]]
 // CHECK:   %[[TID:.*]] = call i32 @__kmpc_global_thread_num(
 // CHECK:   %[[ALLOC:.*]] = call ptr @__kmpc_alloc(i32 %[[TID]], i64 8, ptr null)
-// CHECK:   %[[GEP4:.*]] = getelementptr i8, ptr %[[ALLOC]], i64 4
-// CHECK:   store i32 1, ptr %[[ALLOC]], align 4
-// CHECK:   store i32 2, ptr %[[GEP4]], align 4
+// CHECK:   store i32 42, ptr %[[ALLOC]]
 // CHECK:   %[[TID_FREE:.*]] = call i32 @__kmpc_global_thread_num(
 // CHECK:   call void @__kmpc_free(i32 %[[TID_FREE]], ptr %[[ALLOC]], ptr null)
 // CHECK:   ret void
-llvm.mlir.global internal @common_like() : !llvm.array<2 x i32> {
-  %0 = llvm.mlir.zero : !llvm.array<2 x i32>
-  llvm.return %0 : !llvm.array<2 x i32>
+llvm.func @test_allocate_use_before(%arg0: !llvm.ptr) {
+  %pre = llvm.load %arg0 : !llvm.ptr -> i32
+  omp.allocate_dir (%arg0 : !llvm.ptr)
+  %c42 = llvm.mlir.constant(42 : i32) : i32
+  llvm.store %c42, %arg0 : i32, !llvm.ptr
+  omp.allocate_free (%arg0 : !llvm.ptr)
+  llvm.return
 }
 
-llvm.func @test_allocate_global_gep_users() {
-  %base = llvm.mlir.addressof @common_like : !llvm.ptr
-  %c0 = llvm.mlir.constant(0 : i64) : i64
-  %c1 = llvm.mlir.constant(1 : i64) : i64
-  %m0 = llvm.getelementptr %base[%c0, %c0] : (!llvm.ptr, i64, i64) -> !llvm.ptr, !llvm.array<2 x i32>
-  %m1 = llvm.getelementptr %base[%c0, %c1] : (!llvm.ptr, i64, i64) -> !llvm.ptr, !llvm.array<2 x i32>
-  omp.allocate_dir (%base : !llvm.ptr)
-  %one = llvm.mlir.constant(1 : i32) : i32
-  %two = llvm.mlir.constant(2 : i32) : i32
-  llvm.store %one, %m0 : i32, !llvm.ptr
-  llvm.store %two, %m1 : i32, !llvm.ptr
-  omp.allocate_free (%base : !llvm.ptr)
+// -----
+
+// Verifies dynamic array size with an i32 element count on llvm.alloca.
+//
+// CHECK-LABEL: define void @test_allocate_dynamic_i32_count
+// CHECK:   %[[TID:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   %[[ALLOC:.*]] = call ptr @__kmpc_alloc(i32 %[[TID]], i64 40, ptr null)
+// CHECK:   %[[TID_FREE:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   call void @__kmpc_free(i32 %[[TID_FREE]], ptr %[[ALLOC]], ptr null)
+// CHECK:   ret void
+llvm.func @test_allocate_dynamic_i32_count() {
+  %count = llvm.mlir.constant(10 : i32) : i32
+  %arr = llvm.alloca %count x i32 : (i32) -> !llvm.ptr
+  omp.allocate_dir (%arr : !llvm.ptr)
+  omp.allocate_free (%arr : !llvm.ptr)
+  llvm.return
+}
+
+// -----
+
+// Verifies runtime dynamic array size from a stack alloca: count * 4 bytes.
+//
+// CHECK-LABEL: define void @test_allocate_dynamic_runtime_count
+// CHECK-SAME: (i64 %[[COUNT:.*]])
+// CHECK:   %[[MUL:.*]] = mul i64 %[[COUNT]], 4
+// CHECK:   %[[TID:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   %[[ALLOC:.*]] = call ptr @__kmpc_alloc(i32 %[[TID]], i64 {{.*}}, ptr null)
+// CHECK:   %[[TID_FREE:.*]] = call i32 @__kmpc_global_thread_num(
+// CHECK:   call void @__kmpc_free(i32 %[[TID_FREE]], ptr %[[ALLOC]], ptr null)
+// CHECK:   ret void
+llvm.func @test_allocate_dynamic_runtime_count(%count: i64) {
+  %arr = llvm.alloca %count x i32 : (i64) -> !llvm.ptr
+  omp.allocate_dir (%arr : !llvm.ptr)
+  omp.allocate_free (%arr : !llvm.ptr)
   llvm.return
 }



More information about the flang-commits mailing list