[flang-commits] [flang] d6d52a4 - [flang][cuda] Do not generate cuf.alloc/cuf.free in device context (#141117)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2025-05-22T13:30:34-07:00
New Revision: d6d52a4abc71c583bd5391c5c0b36483c65081fe

URL: https://github.com/llvm/llvm-project/commit/d6d52a4abc71c583bd5391c5c0b36483c65081fe
DIFF: https://github.com/llvm/llvm-project/commit/d6d52a4abc71c583bd5391c5c0b36483c65081fe.diff

LOG: [flang][cuda] Do not generate cuf.alloc/cuf.free in device context (#141117)

`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.

Added: 
    

Modified: 
    flang/lib/Lower/ConvertVariable.cpp
    flang/test/Lower/CUDA/cuda-allocatable.cuf
    flang/test/Lower/CUDA/cuda-shared.cuf

Removed: 
    


################################################################################
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