[flang-commits] [flang] [flang][cuda] Route dynamic autos through malloc_unified/free_unified (PR #212965)

[via flang-commits]

================
@@ -309,3 +314,63 @@ bool fir::replaceAllocas(mlir::RewriterBase &rewriter,
   rewriter.restoreInsertionPoint(insertPoint);
   return replacedAllRequestedAlloca;
 }
+
+/// Device code keeps its stack allocations: the unified/managed entry points
+/// are host-only, and a kernel-side heap allocation would be a large
+/// regression over a device stack array.
+static bool isDeviceCode(mlir::Operation *func, mlir::ModuleOp mod) {
+  if (func->getParentOfType<mlir::gpu::GPUModuleOp>())
+    return true;
+  if (auto procAttr =
+          func->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName()))
+    return procAttr.getValue() != cuf::ProcAttribute::Host;
+  if (mlir::acc::isAccRoutine(func))
+    return true;
+  if (auto offloadMod =
+          llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(mod.getOperation()))
+    return offloadMod.getIsTargetDevice();
+  return false;
+}
+
+bool fir::promoteDynamicAllocasToCudaHeap(mlir::RewriterBase &rewriter,
+                                          mlir::Operation *func) {
+  auto mod = func->getParentOfType<mlir::ModuleOp>();
+  if (!mod)
+    return false;
+  fir::CudaHeapAllocMode mode = fir::getCudaHeapAllocMode(mod);
+  if (mode == fir::CudaHeapAllocMode::None || isDeviceCode(func, mod))
+    return false;
+
+  bool changed = false;
+  // Named locals only: automatic arrays and automatic character. Compiler
+  // temporaries do not need unified memory and would turn a stack save/restore
+  // into a malloc/free pair, possibly per loop iteration.
+  auto mustReplace = [](fir::AllocaOp alloca) {
+    std::optional<llvm::StringRef> uniqName = alloca.getUniqName();
+    return alloca.isDynamic() && uniqName && !uniqName->empty();
----------------
jeanPerier wrote:

Please also exclude the allocas marked as fir.must_be_stack (like function results), otherwise you will leave unused pairs of malloc/free inside functions returning arrays. 

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