[flang-commits] [flang] 9b6504e - [flang][cuda] Make sure to issue freemem for the allocated temp (#98078)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-07-11T17:15:54-07:00
New Revision: 9b6504e98359f5d14fdaa353b2789e7e95239f96

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

LOG: [flang][cuda] Make sure to issue freemem for the allocated temp (#98078)

When implicit data transfer is created, make sure we generate the
`freemem` op on the `allocmem` result value and not the declare op
value.

Added: 
    

Modified: 
    flang/lib/Lower/Bridge.cpp
    flang/test/Lower/CUDA/cuda-data-transfer.cuf

Removed: 
    


################################################################################
diff  --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 3d071f6bb8d5a..267a3557ab8c4 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -4296,8 +4296,13 @@ class FirConverter : public Fortran::lower::AbstractConverter {
           auto [temp, cleanup] =
               hlfir::createTempFromMold(loc, builder, entity);
           auto needCleanup = fir::getIntIfConstant(cleanup);
-          if (needCleanup && *needCleanup)
-            temps.push_back(temp);
+          if (needCleanup && *needCleanup) {
+            if (auto declareOp =
+                    mlir::dyn_cast<hlfir::DeclareOp>(temp.getDefiningOp()))
+              temps.push_back(declareOp.getMemref());
+            else
+              temps.push_back(temp);
+          }
           addSymbol(sym,
                     hlfir::translateToExtendedValue(loc, builder, temp).first,
                     /*forced=*/true);

diff  --git a/flang/test/Lower/CUDA/cuda-data-transfer.cuf b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
index 8b04a2c202dc0..1383b73ea44d6 100644
--- a/flang/test/Lower/CUDA/cuda-data-transfer.cuf
+++ b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
@@ -109,7 +109,7 @@ end
 ! CHECK: cuf.data_transfer %[[ADEV]]#1 to %[[DECL_TEMP]]#0 {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.ref<!fir.array<10xi32>>, !fir.heap<!fir.array<10xi32>>
 ! CHECK: %[[ELEMENTAL:.*]] = hlfir.elemental %{{.*}} unordered : (!fir.shape<1>) -> !hlfir.expr<10xi32>
 ! CHECK: hlfir.assign %[[ELEMENTAL]] to %[[BHOST]]#0 : !hlfir.expr<10xi32>, !fir.ref<!fir.array<10xi32>>
-! CHECK: fir.freemem %[[DECL_TEMP]]#0 : !fir.heap<!fir.array<10xi32>>
+! CHECK: fir.freemem %[[TEMP]] : !fir.heap<!fir.array<10xi32>>
 
 subroutine sub3()
   use mod1
@@ -251,3 +251,17 @@ end subroutine
 ! CHECK: %{{.*}} = fir.call @_QMmod1Pdev1
 ! CHECK: hlfir.assign
 ! CHECK-NOT: cuf.data_transfer
+
+subroutine sub13(a, b, n)
+  integer :: n
+  integer :: a(n)
+  integer, allocatable, device :: b(:)
+  integer :: res(10)
+
+  res = a + b
+end subroutine
+
+! CHECK-LABEL: func.func @_QPsub13
+! CHECK: %[[TEMP:.*]] = fir.allocmem !fir.array<?xi32>, %14#1 {bindc_name = ".tmp", uniq_name = ""}
+! CHECK: cuf.data_transfer
+! CHECK: fir.freemem %[[TEMP]] : !fir.heap<!fir.array<?xi32>>