[flang][cuda] Do not generate cuf.alloc/cuf.free in device context

[clementval]

cuf.alloc and cuf.free are converted to fir.alloca or deleted when in device context during the CUFOpConversion pass. Do not generate them in lowering to avoid confusion.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

cuf.alloc and cuf.free are converted to fir.alloca or deleted when in device context during the CUFOpConversion pass. Do not generate them in lowering to avoid confusion.

---

Full diff: https://github.com/llvm/llvm-project/pull/141117.diff

3 Files Affected:

• (modified) flang/lib/Lower/ConvertVariable.cpp (+8-4)
• (modified) flang/test/Lower/CUDA/cuda-allocatable.cuf (+1-1)
• (modified) flang/test/Lower/CUDA/cuda-shared.cuf (-1)

diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp
index 372c71b6d4821..a28596bfd0099 100644
--- a/flang/lib/Lower/ConvertVariable.cpp
+++ b/flang/lib/Lower/ConvertVariable.cpp
@@ -25,6 +25,7 @@
 #include "flang/Lower/StatementContext.h"
 #include "flang/Lower/Support/Utils.h"
 #include "flang/Lower/SymbolMap.h"
+#include "flang/Optimizer/Builder/CUFCommon.h"
 #include "flang/Optimizer/Builder/Character.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"
@@ -735,8 +736,10 @@ static mlir::Value createNewLocal(Fortran::lower::AbstractConverter &converter,
     if (dataAttr.getValue() == cuf::DataAttribute::Shared)
       return builder.create<cuf::SharedMemoryOp>(loc, ty, nm, symNm, lenParams,
                                                  indices);
-    return builder.create<cuf::AllocOp>(loc, ty, nm, symNm, dataAttr, lenParams,
-                                        indices);
+
+    if (!cuf::isCUDADeviceContext(builder.getRegion()))
+      return builder.create<cuf::AllocOp>(loc, ty, nm, symNm, dataAttr,
+                                          lenParams, indices);
   }
 
   // Let the builder do all the heavy lifting.
@@ -1072,8 +1075,9 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
   if (mustBeDefaultInitializedAtRuntime(var))
     Fortran::lower::defaultInitializeAtRuntime(converter, var.getSymbol(),
                                                symMap);
-  if (Fortran::semantics::NeedCUDAAlloc(var.getSymbol())) {
-    auto *builder = &converter.getFirOpBuilder();
+  auto *builder = &converter.getFirOpBuilder();
+  if (Fortran::semantics::NeedCUDAAlloc(var.getSymbol()) &&
+      !cuf::isCUDADeviceContext(builder->getRegion())) {
     cuf::DataAttributeAttr dataAttr =
         Fortran::lower::translateSymbolCUFDataAttribute(builder->getContext(),
                                                         var.getSymbol());
diff --git a/flang/test/Lower/CUDA/cuda-allocatable.cuf b/flang/test/Lower/CUDA/cuda-allocatable.cuf
index cec10dda839e9..36e768bd7d92c 100644
--- a/flang/test/Lower/CUDA/cuda-allocatable.cuf
+++ b/flang/test/Lower/CUDA/cuda-allocatable.cuf
@@ -186,7 +186,7 @@ attributes(global) subroutine sub8()
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPsub8() attributes {cuf.proc_attr = #cuf.cuda_proc<global>}
-! CHECK: %[[DESC:.*]] = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFsub8Ea"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: %[[DESC:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", uniq_name = "_QFsub8Ea"}
 ! CHECK: %[[A:.*]]:2 = hlfir.declare %[[DESC]] {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub8Ea"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
 ! CHECK: %[[HEAP:.*]] = fir.allocmem !fir.array<?xf32>, %{{.*}} {fir.must_be_heap = true, uniq_name = "_QFsub8Ea.alloc"}
 ! CHECK: %[[SHAPE:.*]] = fir.shape %{{.*}} : (index) -> !fir.shape<1>
diff --git a/flang/test/Lower/CUDA/cuda-shared.cuf b/flang/test/Lower/CUDA/cuda-shared.cuf
index 565857f01bdb8..f41011df06ae7 100644
--- a/flang/test/Lower/CUDA/cuda-shared.cuf
+++ b/flang/test/Lower/CUDA/cuda-shared.cuf
@@ -9,5 +9,4 @@ end subroutine
 
 ! CHECK-LABEL: func.func @_QPsharedmem() attributes {cuf.proc_attr = #cuf.cuda_proc<global>}
 ! CHECK: %{{.*}} = cuf.shared_memory !fir.array<32xf32> {bindc_name = "s", uniq_name = "_QFsharedmemEs"} -> !fir.ref<!fir.array<32xf32>>
-! CHECK: cuf.free %{{.*}}#0 : !fir.ref<i32> {data_attr = #cuf.cuda<device>}
 ! CHECK-NOT: cuf.free