Skip to content

Commit b7637a8

Browse files
authored
[flang][cuda] Set PINNED variable to false in ALLOCATE (#121593)
When `PINNED=` is used with variables that don't have the `PINNED` attribute, the logical value must be set to false when host allocation is performed.
1 parent ee1adc5 commit b7637a8

File tree

2 files changed

+53
-7
lines changed

2 files changed

+53
-7
lines changed

flang/lib/Lower/Allocatable.cpp

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,19 @@ class AllocateStmtHelper {
454454
alloc.getSymbol());
455455
}
456456

457+
void setPinnedToFalse() {
458+
if (!pinnedExpr)
459+
return;
460+
Fortran::lower::StatementContext stmtCtx;
461+
mlir::Value pinned =
462+
fir::getBase(converter.genExprAddr(loc, *pinnedExpr, stmtCtx));
463+
mlir::Location loc = pinned.getLoc();
464+
mlir::Value falseValue = builder.createBool(loc, false);
465+
mlir::Value falseConv = builder.createConvert(
466+
loc, fir::unwrapRefType(pinned.getType()), falseValue);
467+
builder.create<fir::StoreOp>(loc, falseConv, pinned);
468+
}
469+
457470
void genSimpleAllocation(const Allocation &alloc,
458471
const fir::MutableBoxValue &box) {
459472
bool isCudaSymbol = Fortran::semantics::HasCUDAAttr(alloc.getSymbol());
@@ -469,6 +482,7 @@ class AllocateStmtHelper {
469482
// can be validated.
470483
genInlinedAllocation(alloc, box);
471484
postAllocationAction(alloc);
485+
setPinnedToFalse();
472486
return;
473487
}
474488

@@ -482,11 +496,13 @@ class AllocateStmtHelper {
482496
genSetDeferredLengthParameters(alloc, box);
483497
genAllocateObjectBounds(alloc, box);
484498
mlir::Value stat;
485-
if (!isCudaSymbol)
499+
if (!isCudaSymbol) {
486500
stat = genRuntimeAllocate(builder, loc, box, errorManager);
487-
else
501+
setPinnedToFalse();
502+
} else {
488503
stat =
489504
genCudaAllocate(builder, loc, box, errorManager, alloc.getSymbol());
505+
}
490506
fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
491507
postAllocationAction(alloc);
492508
errorManager.assignStat(builder, loc, stat);
@@ -616,13 +632,16 @@ class AllocateStmtHelper {
616632
genSetDeferredLengthParameters(alloc, box);
617633
genAllocateObjectBounds(alloc, box);
618634
mlir::Value stat;
619-
if (Fortran::semantics::HasCUDAAttr(alloc.getSymbol()))
635+
if (Fortran::semantics::HasCUDAAttr(alloc.getSymbol())) {
620636
stat =
621637
genCudaAllocate(builder, loc, box, errorManager, alloc.getSymbol());
622-
else if (isSource)
623-
stat = genRuntimeAllocateSource(builder, loc, box, exv, errorManager);
624-
else
625-
stat = genRuntimeAllocate(builder, loc, box, errorManager);
638+
} else {
639+
if (isSource)
640+
stat = genRuntimeAllocateSource(builder, loc, box, exv, errorManager);
641+
else
642+
stat = genRuntimeAllocate(builder, loc, box, errorManager);
643+
setPinnedToFalse();
644+
}
626645
fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
627646
postAllocationAction(alloc);
628647
errorManager.assignStat(builder, loc, stat);

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,3 +196,30 @@ end subroutine
196196
! CHECK: %[[BOX:.*]] = fir.load %[[A]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
197197
! CHECK: %[[BOXADDR:.*]] = fir.box_addr %[[BOX]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
198198
! CHECK: fir.freemem %[[BOXADDR]] : !fir.heap<!fir.array<?xf32>>
199+
200+
subroutine setpinned()
201+
integer, allocatable :: i(:)
202+
logical :: plog
203+
allocate(i(10), pinned=plog)
204+
end
205+
206+
! CHECK-LABEL: func.func @_QPsetpinned()
207+
! CHECK: %[[PLOG:.*]] = fir.alloca !fir.logical<4> {bindc_name = "plog", uniq_name = "_QFsetpinnedEplog"}
208+
! CHECK: %[[PLOG_DECL:.*]]:2 = hlfir.declare %[[PLOG]] {uniq_name = "_QFsetpinnedEplog"} : (!fir.ref<!fir.logical<4>>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>)
209+
! CHECK: %[[FALSE:.*]] = arith.constant false
210+
! CHECK: %[[FLASE_CONV:.*]] = fir.convert %[[FALSE]] : (i1) -> !fir.logical<4>
211+
! CHECK: fir.store %[[FLASE_CONV]] to %[[PLOG_DECL]]#1 : !fir.ref<!fir.logical<4>>
212+
213+
subroutine setpinnedpointer()
214+
integer, pointer :: i(:)
215+
logical :: plog
216+
allocate(i(10), pinned=plog)
217+
end
218+
219+
! CHECK-LABEL: func.func @_QPsetpinnedpointer()
220+
! CHECK: %[[PLOG:.*]] = fir.alloca !fir.logical<4> {bindc_name = "plog", uniq_name = "_QFsetpinnedpointerEplog"}
221+
! CHECK: %[[PLOG_DECL:.*]]:2 = hlfir.declare %[[PLOG]] {uniq_name = "_QFsetpinnedpointerEplog"} : (!fir.ref<!fir.logical<4>>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>)
222+
! CHECK: fir.call @_FortranAPointerAllocate
223+
! CHECK: %[[FALSE:.*]] = arith.constant false
224+
! CHECK: %[[FLASE_CONV:.*]] = fir.convert %[[FALSE]] : (i1) -> !fir.logical<4>
225+
! CHECK: fir.store %[[FLASE_CONV]] to %[[PLOG_DECL]]#1 : !fir.ref<!fir.logical<4>>

0 commit comments

Comments
 (0)