[flang-commits] [flang] [flang][cuda] Split realloc for host-accessible CUDA Fortran allocatable assigns (PR #224195)
via flang-commits
flang-commits at lists.llvm.org
Wed Sep 16 22:09:14 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-fir-hlfir
Author: Valentin Clement (バレンタイン クレメン) (clementval)
<details>
<summary>Changes</summary>
Allow SeparateAllocatableAssign to peel reallocation off hlfir.assign for
managed, pinned, and unified allocatables, using cuf.alloc/cuf.free so the
CUDA Fortran allocator is honored. Device (and other non-host-accessible)
storage still keeps the runtime realloc path as well as derived-type.
Drop the extra op-level MemFree on cuf.free so OptimizedBufferization can
see the freed pointer. Without that, a valueless Free effect between the
elemental and the assign forced a temporary for cases such as
`a = int(ran*100)` on a managed allocatable in the example below.
```
subroutine foo(n, a, ran)
integer :: n
real, managed, allocatable :: a(:)
real, allocatable :: ran(:)
allocate(a(n))
allocate(ran(n))
call random_number(ran)
a = int(ran*100)
end subroutine
```
---
Full diff: https://github.com/llvm/llvm-project/pull/224195.diff
6 Files Affected:
- (modified) flang/include/flang/Optimizer/Builder/MutableBox.h (+5-2)
- (modified) flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td (+1-1)
- (modified) flang/lib/Optimizer/Builder/MutableBox.cpp (+31-8)
- (modified) flang/lib/Optimizer/HLFIR/Transforms/SeparateAllocatableAssign.cpp (+26-10)
- (modified) flang/test/HLFIR/opt-bufferization-dealloc-conflict.fir (+51)
- (modified) flang/test/HLFIR/separate-allocatable-assign.fir (+28-7)
``````````diff
diff --git a/flang/include/flang/Optimizer/Builder/MutableBox.h b/flang/include/flang/Optimizer/Builder/MutableBox.h
index 9ba09495f328a..ba78a3525a797 100644
--- a/flang/include/flang/Optimizer/Builder/MutableBox.h
+++ b/flang/include/flang/Optimizer/Builder/MutableBox.h
@@ -14,6 +14,7 @@
#define FORTRAN_OPTIMIZER_BUILDER_MUTABLEBOX_H
#include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h"
#include "flang/Runtime/allocator-registry-consts.h"
#include "llvm/ADT/StringRef.h"
@@ -115,12 +116,14 @@ MutableBoxReallocation
genReallocIfNeeded(fir::FirOpBuilder &builder, mlir::Location loc,
const fir::MutableBoxValue &box, mlir::ValueRange shape,
mlir::ValueRange lenParams,
- ReallocStorageHandlerFunc storageHandler = {});
+ ReallocStorageHandlerFunc storageHandler = {},
+ cuf::DataAttributeAttr dataAttr = {});
void finalizeRealloc(fir::FirOpBuilder &builder, mlir::Location loc,
const fir::MutableBoxValue &box, mlir::ValueRange lbounds,
bool takeLboundsIfRealloc,
- const MutableBoxReallocation &realloc);
+ const MutableBoxReallocation &realloc,
+ cuf::DataAttributeAttr dataAttr = {});
/// Deallocate a mutable box with fir.freemem if it is allocated or associated.
/// This only deallocates the storage and does not call finalization, the
diff --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index 5148d5aa88c3d..1c4a53b00e428 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -68,7 +68,7 @@ def cuf_AllocOp : cuf_Op<"alloc", [AttrSizedOperandSegments]> {
let hasVerifier = 1;
}
-def cuf_FreeOp : cuf_Op<"free", [MemoryEffects<[MemFree]>]> {
+def cuf_FreeOp : cuf_Op<"free", []> {
let summary = "Free a device allocated object";
let description = [{
diff --git a/flang/lib/Optimizer/Builder/MutableBox.cpp b/flang/lib/Optimizer/Builder/MutableBox.cpp
index 170fecf9bffa4..8c022116e9d9c 100644
--- a/flang/lib/Optimizer/Builder/MutableBox.cpp
+++ b/flang/lib/Optimizer/Builder/MutableBox.cpp
@@ -16,6 +16,7 @@
#include "flang/Optimizer/Builder/Runtime/Derived.h"
#include "flang/Optimizer/Builder/Runtime/Stop.h"
#include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
#include "flang/Optimizer/Dialect/FIRAttr.h"
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIROpsSupport.h"
@@ -749,11 +750,21 @@ static mlir::Value allocateAndInitNewStorage(fir::FirOpBuilder &builder,
const fir::MutableBoxValue &box,
mlir::ValueRange extents,
mlir::ValueRange lenParams,
- llvm::StringRef allocName) {
+ llvm::StringRef allocName,
+ cuf::DataAttributeAttr dataAttr) {
auto lengths = getNewLengths(builder, loc, box, lenParams);
+
+ if (dataAttr) {
+ auto newStorage =
+ cuf::AllocOp::create(builder, loc, box.getBaseTy(), allocName,
+ allocName, dataAttr, lengths, extents)
+ .getResult();
+ return fir::ConvertOp::create(
+ builder, loc, fir::HeapType::get(box.getBaseTy()), newStorage);
+ }
+
auto newStorage = fir::AllocMemOp::create(builder, loc, box.getBaseTy(),
allocName, lengths, extents);
-
if (mlir::isa<fir::SequenceType>(box.getBaseTy()))
newStorage.setAlignment(fir::defaultArrayGlobalAlignment);
if (mlir::isa<fir::RecordType>(box.getEleTy())) {
@@ -810,7 +821,8 @@ fir::factory::MutableBoxReallocation fir::factory::genReallocIfNeeded(
fir::FirOpBuilder &builder, mlir::Location loc,
const fir::MutableBoxValue &box, mlir::ValueRange shape,
mlir::ValueRange lengthParams,
- fir::factory::ReallocStorageHandlerFunc storageHandler) {
+ fir::factory::ReallocStorageHandlerFunc storageHandler,
+ cuf::DataAttributeAttr dataAttr) {
// Implement 10.2.1.3 point 3 logic when lhs is an array.
auto reader = MutablePropertyReader(builder, loc, box);
auto addr = reader.readBaseAddress();
@@ -886,7 +898,7 @@ fir::factory::MutableBoxReallocation fir::factory::genReallocIfNeeded(
: shape;
auto heap = allocateAndInitNewStorage(
builder, loc, box, extents, lengthParams,
- ".auto.alloc");
+ ".auto.alloc", dataAttr);
if (storageHandler)
storageHandler(getExtValForStorage(heap));
fir::ResultOp::create(builder, loc, heap);
@@ -913,8 +925,9 @@ fir::factory::MutableBoxReallocation fir::factory::genReallocIfNeeded(
fir::ResultOp::create(builder, loc,
mlir::ValueRange{trueValue, addr});
} else {
- auto heap = allocateAndInitNewStorage(
- builder, loc, box, shape, lengthParams, ".auto.alloc");
+ auto heap = allocateAndInitNewStorage(builder, loc, box, shape,
+ lengthParams, ".auto.alloc",
+ dataAttr);
if (storageHandler)
storageHandler(getExtValForStorage(heap));
fir::ResultOp::create(builder, loc,
@@ -934,7 +947,8 @@ void fir::factory::finalizeRealloc(fir::FirOpBuilder &builder,
const fir::MutableBoxValue &box,
mlir::ValueRange lbounds,
bool takeLboundsIfRealloc,
- const MutableBoxReallocation &realloc) {
+ const MutableBoxReallocation &realloc,
+ cuf::DataAttributeAttr dataAttr) {
builder.genIfThen(loc, realloc.wasReallocated)
.genThen([&]() {
auto reader = MutablePropertyReader(builder, loc, box);
@@ -953,7 +967,16 @@ void fir::factory::finalizeRealloc(fir::FirOpBuilder &builder,
auto heap = fir::getBase(realloc.newValue);
auto extents = fir::factory::getExtents(loc, builder, realloc.newValue);
builder.genIfThen(loc, realloc.oldAddressWasAllocated)
- .genThen([&]() { ::genFreemem(builder, loc, realloc.oldAddress); })
+ .genThen([&]() {
+ if (dataAttr) {
+ mlir::Value devPtr = builder.createConvert(
+ loc, fir::ReferenceType::get(box.getBaseTy()),
+ realloc.oldAddress);
+ cuf::FreeOp::create(builder, loc, devPtr, dataAttr);
+ } else {
+ ::genFreemem(builder, loc, realloc.oldAddress);
+ }
+ })
.end();
MutablePropertyWriter{builder, loc, box}.updateMutableBox(
heap, lbs, extents, lengths);
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/SeparateAllocatableAssign.cpp b/flang/lib/Optimizer/HLFIR/Transforms/SeparateAllocatableAssign.cpp
index 52148f83202b3..124c727693927 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/SeparateAllocatableAssign.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/SeparateAllocatableAssign.cpp
@@ -77,12 +77,30 @@ class SeparateAllocatableAssignConversion
if (!fir::isBoxAddress(lhs.getType()))
return rewriter.notifyMatchFailure(assign, "LHS is not a box address");
- // If the LHS allocatable is non-default, its (re)allocation must go through
- // the flang runtime so the allocator recorded in the descriptor is honored.
+ fir::MutableBoxValue mutableBox(lhs.getFirBase(), /*lenParameters=*/{},
+ /*mutableProperties=*/{});
+
+ // An allocatable with a CUDA Fortran data attribute is backed by a
+ // non-default allocator. The reallocation can still be split out, but it
+ // must use that allocator instead of an inline host allocation, so the
+ // attribute is forwarded to the reallocation helpers below.
+ // The plain hlfir.assign left behind copies the elements on the host, so
+ // only allocators that produce host-accessible storage can be handled here.
+ cuf::DataAttributeAttr dataAttr;
if (mlir::Operation *lhsDef = assign.getLhs().getDefiningOp())
- if (cuf::getDataAttr(lhsDef))
- return rewriter.notifyMatchFailure(
- assign, "LHS uses a non-default allocator; keep runtime realloc");
+ if ((dataAttr = cuf::getDataAttr(lhsDef))) {
+ cuf::DataAttribute attr = dataAttr.getValue();
+ if (attr != cuf::DataAttribute::Managed &&
+ attr != cuf::DataAttribute::Pinned &&
+ attr != cuf::DataAttribute::Unified) {
+ return rewriter.notifyMatchFailure(
+ assign, "LHS storage is not host-accessible; keep runtime "
+ "realloc");
+ }
+ if (mlir::isa<fir::RecordType>(mutableBox.getEleTy()))
+ return rewriter.notifyMatchFailure(
+ assign, "LHS storage is a record; keep runtime realloc");
+ }
mlir::Location loc = assign->getLoc();
fir::FirOpBuilder builder(rewriter, assign.getOperation());
@@ -133,16 +151,14 @@ class SeparateAllocatableAssignConversion
rhsLbounds.push_back(lb);
}
- fir::MutableBoxValue mutableBox(lhs.getFirBase(), /*lenParameters=*/{},
- /*mutableProperties=*/{});
-
auto noopHandler = [](fir::ExtendedValue) {};
llvm::SmallVector<mlir::Value> lenParams;
fir::factory::MutableBoxReallocation realloc =
fir::factory::genReallocIfNeeded(builder, loc, mutableBox, rhsExtents,
- lenParams, noopHandler);
+ lenParams, noopHandler, dataAttr);
fir::factory::finalizeRealloc(builder, loc, mutableBox, rhsLbounds,
- /*takeLboundsIfRealloc=*/true, realloc);
+ /*takeLboundsIfRealloc=*/true, realloc,
+ dataAttr);
mlir::Value lhsBox = fir::LoadOp::create(builder, loc, lhs.getFirBase());
hlfir::AssignOp::create(builder, loc, rhs, lhsBox,
diff --git a/flang/test/HLFIR/opt-bufferization-dealloc-conflict.fir b/flang/test/HLFIR/opt-bufferization-dealloc-conflict.fir
index 2b5eda169a7fc..e7bb9b51ae7e1 100644
--- a/flang/test/HLFIR/opt-bufferization-dealloc-conflict.fir
+++ b/flang/test/HLFIR/opt-bufferization-dealloc-conflict.fir
@@ -30,3 +30,54 @@ func.func @dealloc_conflict(%dst: !fir.box<!fir.array<?xf32>>, %n: index) {
// CHECK: hlfir.elemental
// CHECK: fir.freemem
// CHECK: hlfir.assign %{{.*}} to %{{.*}} : !hlfir.expr<?xf32>, !fir.box<!fir.array<?xf32>>
+
+// Same conflict through cuf.free, which frees the memory of its operand.
+func.func @cuf_dealloc_conflict(%dst: !fir.box<!fir.array<?xf32>>, %n: index) {
+ %cst = arith.constant 1.000000e+00 : f32
+ %shape = fir.shape %n : (index) -> !fir.shape<1>
+ %ptr = cuf.alloc !fir.array<?xf32>, %n : index {bindc_name = ".src", data_attr = #cuf.cuda<managed>, uniq_name = ".src"} -> !fir.ref<!fir.array<?xf32>>
+ %heap = fir.convert %ptr : (!fir.ref<!fir.array<?xf32>>) -> !fir.heap<!fir.array<?xf32>>
+ %src = fir.embox %heap(%shape) : (!fir.heap<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.box<!fir.array<?xf32>>
+ %elem = hlfir.elemental %shape unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
+ ^bb0(%i: index):
+ %d = hlfir.designate %src (%i) : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ %v = fir.load %d : !fir.ref<f32>
+ %s = arith.addf %v, %cst fastmath<contract> : f32
+ hlfir.yield_element %s : f32
+ }
+ cuf.free %ptr : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<managed>}
+ hlfir.assign %elem to %dst : !hlfir.expr<?xf32>, !fir.box<!fir.array<?xf32>>
+ hlfir.destroy %elem : !hlfir.expr<?xf32>
+ return
+}
+
+// CHECK-LABEL: func.func @cuf_dealloc_conflict
+// CHECK: hlfir.elemental
+// CHECK: cuf.free
+// CHECK: hlfir.assign %{{.*}} to %{{.*}} : !hlfir.expr<?xf32>, !fir.box<!fir.array<?xf32>>
+
+// A cuf.free of memory that the elemental does not read must not block the
+// fusion: the free effect is attached to the freed pointer, so the alias
+// analysis can tell the two apart.
+func.func @cuf_dealloc_no_conflict(%dst: !fir.box<!fir.array<?xf32>>, %src: !fir.box<!fir.array<?xf32>>, %n: index) {
+ %cst = arith.constant 1.000000e+00 : f32
+ %shape = fir.shape %n : (index) -> !fir.shape<1>
+ %ptr = cuf.alloc !fir.array<?xf32>, %n : index {bindc_name = ".unrelated", data_attr = #cuf.cuda<managed>, uniq_name = ".unrelated"} -> !fir.ref<!fir.array<?xf32>>
+ %elem = hlfir.elemental %shape unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
+ ^bb0(%i: index):
+ %d = hlfir.designate %src (%i) : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ %v = fir.load %d : !fir.ref<f32>
+ %s = arith.addf %v, %cst fastmath<contract> : f32
+ hlfir.yield_element %s : f32
+ }
+ cuf.free %ptr : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<managed>}
+ hlfir.assign %elem to %dst : !hlfir.expr<?xf32>, !fir.box<!fir.array<?xf32>>
+ hlfir.destroy %elem : !hlfir.expr<?xf32>
+ return
+}
+
+// CHECK-LABEL: func.func @cuf_dealloc_no_conflict
+// The elemental is fused into the assignment loop.
+// CHECK-NOT: hlfir.elemental
+// CHECK: fir.do_loop
+// CHECK: hlfir.assign %{{.*}} to %{{.*}} : f32, !fir.ref<f32>
diff --git a/flang/test/HLFIR/separate-allocatable-assign.fir b/flang/test/HLFIR/separate-allocatable-assign.fir
index 747dc33aa2bd6..2fdeb19957776 100644
--- a/flang/test/HLFIR/separate-allocatable-assign.fir
+++ b/flang/test/HLFIR/separate-allocatable-assign.fir
@@ -180,10 +180,10 @@ func.func @test_lower_bounds(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32
// CHECK: %[[BOX:.*]] = fir.load %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
// CHECK: hlfir.assign %{{.*}} to %[[BOX]] : !fir.box<!fir.array<3xi32>>, !fir.box<!fir.heap<!fir.array<?xi32>>>
-// Test: allocatable backed by a non-default (CUF) allocator must NOT be
-// separated. The reallocation has to stay in hlfir.assign ... realloc so that
-// it is routed through the runtime and honors the allocator recorded in the
-// descriptor (pinned/managed/unified), instead of an inline host fir.allocmem.
+// Test: an allocatable backed by a host-accessible CUF allocator is separated
+// like a plain allocatable, except that the reallocation goes through the CUDA
+// Fortran runtime (cuf.alloc/cuf.free) so the memory space requested by the
+// data attribute is preserved.
func.func @test_pinned_allocator(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, %arg1: !fir.ref<!fir.array<10xf64>>) {
%c10 = arith.constant 10 : index
%shape = fir.shape %c10 : (index) -> !fir.shape<1>
@@ -193,10 +193,14 @@ func.func @test_pinned_allocator(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?
return
}
// CHECK-LABEL: func.func @test_pinned_allocator
-// Non-default allocator: keep the realloc assign, do not inline the allocation.
-// CHECK: hlfir.assign %{{.*}} to %{{.*}} realloc
+// CHECK-NOT: hlfir.assign{{.*}}realloc
// CHECK-NOT: fir.allocmem
-// Test: same for a managed allocatable (also host-accessible, non-default alloc).
+// CHECK: cuf.alloc !fir.array<?xf64>{{.*}}data_attr = #cuf.cuda<pinned>
+// CHECK: cuf.free %{{.*}} {data_attr = #cuf.cuda<pinned>}
+// CHECK: %[[BOX:.*]] = fir.load %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+// CHECK: hlfir.assign %{{.*}} to %[[BOX]] : !fir.ref<!fir.array<10xf64>>, !fir.box<!fir.heap<!fir.array<?xf64>>>
+
+// Test: same for a managed allocatable.
func.func @test_managed_allocator(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, %arg1: !fir.ref<!fir.array<10xf64>>) {
%c10 = arith.constant 10 : index
%shape = fir.shape %c10 : (index) -> !fir.shape<1>
@@ -206,5 +210,22 @@ func.func @test_managed_allocator(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<
return
}
// CHECK-LABEL: func.func @test_managed_allocator
+// CHECK-NOT: hlfir.assign{{.*}}realloc
+// CHECK-NOT: fir.allocmem
+// CHECK: cuf.alloc !fir.array<?xf64>{{.*}}data_attr = #cuf.cuda<managed>
+// CHECK: cuf.free %{{.*}} {data_attr = #cuf.cuda<managed>}
+
+// Test: an allocatable in device memory must NOT be separated: the plain
+// hlfir.assign left behind would copy the elements on the host.
+func.func @test_device_allocator(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, %arg1: !fir.ref<!fir.array<10xf64>>) {
+ %c10 = arith.constant 10 : index
+ %shape = fir.shape %c10 : (index) -> !fir.shape<1>
+ %a:2 = hlfir.declare %arg0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
+ %b:2 = hlfir.declare %arg1(%shape) {uniq_name = "_QFEb"} : (!fir.ref<!fir.array<10xf64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf64>>, !fir.ref<!fir.array<10xf64>>)
+ hlfir.assign %b#0 to %a#0 realloc : !fir.ref<!fir.array<10xf64>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ return
+}
+// CHECK-LABEL: func.func @test_device_allocator
// CHECK: hlfir.assign %{{.*}} to %{{.*}} realloc
+// CHECK-NOT: cuf.alloc
// CHECK-NOT: fir.allocmem
``````````
</details>
https://github.com/llvm/llvm-project/pull/224195
More information about the flang-commits
mailing list