Skip to content

Commit f3000d7

Browse files
authored
[flang][cuda] Do not trigger automatic deallocation in main (#128789)
Similar to host flow, do not trigger automatic deallocation at then end of the main program since anything could happen like a cudaDevcieReset().
1 parent c8b4086 commit f3000d7

File tree

3 files changed

+33
-13
lines changed

3 files changed

+33
-13
lines changed

flang/lib/Lower/ConvertVariable.cpp

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -977,12 +977,15 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
977977
fir::ExtendedValue exv =
978978
converter.getSymbolExtendedValue(var.getSymbol(), &symMap);
979979
auto *sym = &var.getSymbol();
980-
converter.getFctCtx().attachCleanup([builder, loc, exv, sym]() {
981-
cuf::DataAttributeAttr dataAttr =
982-
Fortran::lower::translateSymbolCUFDataAttribute(builder->getContext(),
983-
*sym);
984-
builder->create<cuf::FreeOp>(loc, fir::getBase(exv), dataAttr);
985-
});
980+
const Fortran::semantics::Scope &owner = sym->owner();
981+
if (owner.kind() != Fortran::semantics::Scope::Kind::MainProgram) {
982+
converter.getFctCtx().attachCleanup([builder, loc, exv, sym]() {
983+
cuf::DataAttributeAttr dataAttr =
984+
Fortran::lower::translateSymbolCUFDataAttribute(
985+
builder->getContext(), *sym);
986+
builder->create<cuf::FreeOp>(loc, fir::getBase(exv), dataAttr);
987+
});
988+
}
986989
}
987990
if (std::optional<VariableCleanUp> cleanup =
988991
needDeallocationOrFinalization(var)) {

flang/test/Lower/CUDA/cuda-return01.cuf

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,32 @@
22

33
! Check if finalization works with a return statement
44

5-
program main
5+
subroutine sub1
66
integer, device :: a(10)
77
return
88
end
99

10-
! CHECK: func.func @_QQmain() attributes {fir.bindc_name = "main"} {
10+
! CHECK: func.func @_QPsub1()
1111
! CHECK: %[[DECL:.*]]:2 = hlfir.declare
12-
! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
13-
! CHECK-NEXT: return
14-
! CHECK-NEXT: }
12+
! CHECK: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
13+
! CHECK: return
14+
! CHECK: }
15+
16+
subroutine sub2
17+
integer, device, allocatable :: a(:)
18+
return
19+
end
20+
21+
! CHECK-LABEL: func.func @_QPsub2()
22+
! CHECK: fir.if
23+
! CHECK: cuf.deallocate
24+
! CHECK: cuf.free
25+
26+
program main
27+
integer, allocatable, device :: a(:)
28+
return
29+
end
30+
31+
! CHECK-LABEL: func.func @_QQmain() attributes {fir.bindc_name = "main"}
32+
! CHECK: cuf.alloc !fir.box<!fir.heap<!fir.array<?xi32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
33+
! CHECK-NOT: cuf.free

flang/test/Lower/CUDA/cuda-return02.cuf

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,8 @@ end
1717
! CHECK: %[[DECL:.*]]:2 = hlfir.declare
1818
! CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
1919
! CHECK-NEXT: ^bb1:
20-
! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
2120
! CHECK-NEXT: return
2221
! CHECK-NEXT: ^bb2:
23-
! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
2422
! CHECK-NEXT: return
2523
! CHECK-NEXT: }
2624

0 commit comments

Comments
 (0)