llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-fir-hlfir
@llvm/pr-subscribers-flang-codegen
Author: Valentin Clement (バレンタイン クレメン) (clementval)
<details>
<summary>Changes</summary>
isDeviceAllocation() stopped at fir.declare / fircg.ext_declare, so an embox of managed data could land on the stack. That happens with -g, where declare is rewritten to fircg.ext_declare and kept until codegen.
Walk through those ops so the descriptor for managed data is still allocated in managed memory. Also dyn_cast the adaptor block argument instead of assuming it is always a BlockArgument.
Add a cuda-code-gen.mlir case for embox of a CUFMemAlloc result reached through fircg.ext_declare.
---
Full diff: https://github.com/llvm/llvm-project/pull/220361.diff
2 Files Affected:
- (modified) flang/lib/Optimizer/CodeGen/CodeGen.cpp (+12-1)
- (modified) flang/test/Fir/CUDA/cuda-code-gen.mlir (+37)
``````````diff
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 10ec85340fde6..67be3201b82db 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -2323,8 +2323,19 @@ static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal,
mlir::dyn_cast_or_null<fir::ConvertOp>(val.getDefiningOp()))
return isDeviceAllocation(convertOp.getValue(), {},
cudaDescriptorAllocFunction);
+ // fir.declare, and the fircg.ext_declare it becomes when debug info keeps it
+ // alive until codegen, are pass-through on their memref operand. The adaptor
+ // value is forwarded too so a declared dummy argument stays recognizable.
+ if (auto declareOp =
+ mlir::dyn_cast_or_null<fir::DeclareOp>(val.getDefiningOp()))
+ return isDeviceAllocation(declareOp.getMemref(), adaptorVal,
+ cudaDescriptorAllocFunction);
+ if (auto xDeclareOp =
+ mlir::dyn_cast_or_null<fir::cg::XDeclareOp>(val.getDefiningOp()))
+ return isDeviceAllocation(xDeclareOp.getMemref(), adaptorVal,
+ cudaDescriptorAllocFunction);
if (!val.getDefiningOp() && adaptorVal) {
- if (auto blockArg = llvm::cast<mlir::BlockArgument>(adaptorVal)) {
+ if (auto blockArg = llvm::dyn_cast<mlir::BlockArgument>(adaptorVal)) {
if (blockArg.getOwner() && blockArg.getOwner()->getParentOp() &&
blockArg.getOwner()->isEntryBlock()) {
if (auto func = mlir::dyn_cast_or_null<mlir::FunctionOpInterface>(
diff --git a/flang/test/Fir/CUDA/cuda-code-gen.mlir b/flang/test/Fir/CUDA/cuda-code-gen.mlir
index 30e41efdc9700..fc4a14af47c71 100644
--- a/flang/test/Fir/CUDA/cuda-code-gen.mlir
+++ b/flang/test/Fir/CUDA/cuda-code-gen.mlir
@@ -501,3 +501,40 @@ module attributes {gpu.container_module} {
// CHECK-LABEL: llvm.func @_QQmain()
// CHECK: llvm.alloca {{.*}} x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<2 x array<3 x i64>>)>
// CHECK-NOT: llvm.call @_FortranACUFAllocDescriptor
+
+// -----
+
+// fircg.ext_declare, which survives until codegen when debug info is requested,
+// must not hide the device-resident memory it declares. The descriptor for an
+// embox of a managed allocation goes to managed memory whether or not the
+// allocation is reached through a declare.
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>} {
+ func.func @_QPtesti4() {
+ %c1_i32 = arith.constant 1 : i32
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %c29_i32 = arith.constant 29 : i32
+ %0 = fir.address_of(@_QQclX3C737464696E3E00) : !fir.ref<!fir.char<1,8>>
+ %1 = fir.convert %c24 : (index) -> i64
+ %2 = fir.convert %0 : (!fir.ref<!fir.char<1,8>>) -> !fir.ref<i8>
+ %3 = fir.call @_FortranACUFMemAlloc(%1, %c1_i32, %2, %c29_i32) {cuf.data_attr = #cuf.cuda<managed>} : (i64, i32, !fir.ref<i8>, i32) -> !fir.llvm_ptr<i8>
+ %4 = fir.convert %3 : (!fir.llvm_ptr<i8>) -> !fir.ref<!fir.array<6x!fir.logical<4>>>
+ %5 = fircg.ext_declare %4(%c6) {uniq_name = "_QFtesti4Elma"} : (!fir.ref<!fir.array<6x!fir.logical<4>>>, index) -> !fir.ref<!fir.array<6x!fir.logical<4>>>
+ %6 = fircg.ext_embox %5(%c6) : (!fir.ref<!fir.array<6x!fir.logical<4>>>, index) -> !fir.box<!fir.array<6x!fir.logical<4>>>
+ fir.call @_QPcallee(%6) : (!fir.box<!fir.array<6x!fir.logical<4>>>) -> ()
+ return
+ }
+ func.func private @_QPcallee(!fir.box<!fir.array<6x!fir.logical<4>>>)
+ func.func private @_FortranACUFMemAlloc(i64, i32, !fir.ref<i8>, i32) -> !fir.llvm_ptr<i8> attributes {fir.runtime}
+ fir.global linkonce @_QQclX3C737464696E3E00 constant : !fir.char<1,8> {
+ %0 = fir.string_lit "<stdin>\00"(8) : !fir.char<1,8>
+ fir.has_value %0 : !fir.char<1,8>
+ }
+}
+
+// CHECK-LABEL: llvm.func @_QPtesti4()
+// CHECK: llvm.call @_FortranACUFAllocDescriptor(
+
+// CUSTOM-LABEL: llvm.func @_QPtesti4()
+// CUSTOM: llvm.call @custom_alloc_desc(
``````````
</details>
https://github.com/llvm/llvm-project/pull/220361