[flang-commits] [flang] [flang][cuda] Lower device is active as an operation (PR #225583)
Valentin Clement バレンタイン クレメン via flang-commits
flang-commits at lists.llvm.org
Tue Sep 22 21:10:26 PDT 2026
https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/225583
>From 38038fc09aa6e01bd5b3cf43ef666d533dcf78e2 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 22 Sep 2026 18:22:21 -0700
Subject: [PATCH 1/2] [flang][cuda] Lower device is active as an operation
Instead of lowering straight to a functin call, go through an op so it can be
lowered differently between device and host.
---
.../Optimizer/Builder/Runtime/CUDA/Support.h | 5 +--
.../flang/Optimizer/Dialect/CUF/CUFOps.td | 17 ++++++++++
.../flang/Optimizer/Transforms/Passes.td | 4 ++-
.../Builder/Runtime/CUDA/Support.cpp | 10 ++----
.../Transforms/CUDA/CUFOpConversion.cpp | 1 +
.../Transforms/CUDA/CUFOpConversionLate.cpp | 31 +++++++++++++++++++
flang/test/Fir/cuf.mlir | 11 ++++++-
flang/test/Lower/CUDA/cuda-allocatable.cuf | 2 +-
flang/test/Lower/CUDA/cuda-derived.cuf | 2 +-
flang/test/Lower/CUDA/cuda-return01.cuf | 6 ++--
flang/test/Lower/CUDA/cuda-return02.cuf | 6 ++--
11 files changed, 75 insertions(+), 20 deletions(-)
diff --git a/flang/include/flang/Optimizer/Builder/Runtime/CUDA/Support.h b/flang/include/flang/Optimizer/Builder/Runtime/CUDA/Support.h
index df1c90eb53515c..8cac513d6a84ca 100644
--- a/flang/include/flang/Optimizer/Builder/Runtime/CUDA/Support.h
+++ b/flang/include/flang/Optimizer/Builder/Runtime/CUDA/Support.h
@@ -23,8 +23,9 @@ namespace fir::runtime::cuda {
/// Generate runtime call to synchronize the CUDA device.
void genCUDADeviceSynchronize(fir::FirOpBuilder &builder, mlir::Location loc);
-/// Generate a call returning (as i1) whether the device's primary context is
-/// alive, to guard scope-exit frees against a user cudaDeviceReset().
+/// Generate a `cuf.device_is_active` query (i1) used to guard host-side
+/// scope-exit frees against a user cudaDeviceReset(). Converted later to a
+/// runtime call on the host and to false in device code.
mlir::Value genDeviceIsActive(fir::FirOpBuilder &builder, mlir::Location loc);
} // namespace fir::runtime::cuda
diff --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index 1c4a53b00e4289..bae9794d91f91d 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -396,6 +396,23 @@ def cuf_SharedMemoryOp
CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>];
}
+def cuf_DeviceIsActiveOp : cuf_Op<"device_is_active", []> {
+ let summary = "Query whether the CUDA device context is still active";
+
+ let description = [{
+ Returns true if the CUDA device's primary context is alive. Used to
+ guard host-side scope-exit deallocation against a user cudaDeviceReset().
+
+ Converted to a `_FortranACUFDeviceIsActive` runtime call on the host and
+ to `false` in device code, so host-only cleanup is DCE'd from GPU
+ functions.
+ }];
+
+ let results = (outs I1:$res);
+
+ let assemblyFormat = "attr-dict `:` type($res)";
+}
+
def cuf_StreamCastOp : cuf_Op<"stream_cast", [NoMemoryEffect]> {
let summary = "Adapt a stream value to a GPU async token";
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 619f44e53bac5f..af6785eba80aa5 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -590,7 +590,9 @@ def CUFOpConversion : Pass<"cuf-convert", "mlir::ModuleOp"> {
def CUFOpConversionLate : Pass<"cuf-convert-late", "mlir::ModuleOp"> {
let summary = "Convert CUF operations to runtime calls late in the pipeline";
- let dependentDialects = ["fir::FIROpsDialect"];
+ let dependentDialects = [
+ "fir::FIROpsDialect", "mlir::arith::ArithDialect", "mlir::gpu::GPUDialect"
+ ];
}
def CUFDeviceGlobal :
diff --git a/flang/lib/Optimizer/Builder/Runtime/CUDA/Support.cpp b/flang/lib/Optimizer/Builder/Runtime/CUDA/Support.cpp
index eff0068ebb634d..81edb56a2c99a2 100644
--- a/flang/lib/Optimizer/Builder/Runtime/CUDA/Support.cpp
+++ b/flang/lib/Optimizer/Builder/Runtime/CUDA/Support.cpp
@@ -8,14 +8,11 @@
#include "flang/Optimizer/Builder/Runtime/CUDA/Support.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
-#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
-#include "flang/Runtime/CUDA/allocatable.h"
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
static constexpr llvm::StringRef kCudaDeviceSynchronizeName =
"_QPcudadevicesynchronize";
-using namespace Fortran::runtime::cuda;
-
void fir::runtime::cuda::genCUDADeviceSynchronize(fir::FirOpBuilder &builder,
mlir::Location loc) {
mlir::func::FuncOp func =
@@ -37,8 +34,5 @@ void fir::runtime::cuda::genCUDADeviceSynchronize(fir::FirOpBuilder &builder,
mlir::Value fir::runtime::cuda::genDeviceIsActive(fir::FirOpBuilder &builder,
mlir::Location loc) {
- mlir::func::FuncOp func =
- fir::runtime::getRuntimeFunc<mkRTKey(CUFDeviceIsActive)>(loc, builder);
- auto call = fir::CallOp::create(builder, loc, func, mlir::ValueRange{});
- return builder.createConvert(loc, builder.getI1Type(), call.getResult(0));
+ return cuf::DeviceIsActiveOp::create(builder, loc);
}
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
index 05344102878e26..7a3576d37a0301 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
@@ -617,6 +617,7 @@ class CUFOpConversion : public fir::impl::CUFOpConversionBase<CUFOpConversion> {
mlir::gpu::GPUDialect>();
target.addLegalOp<cuf::StreamCastOp>();
target.addLegalOp<cuf::DeviceAddressOp>();
+ target.addLegalOp<cuf::DeviceIsActiveOp>();
cuf::populateCUFToFIRConversionPatterns(typeConverter, *dl, symtab,
patterns);
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
index 3bbba8753addc4..291e55fb1d3bc9 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "flang/Optimizer/Builder/CUFCommon.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/Runtime/CUDA/Descriptor.h"
#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
@@ -15,8 +16,10 @@
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Support/InternalNames.h"
#include "flang/Optimizer/Transforms/Passes.h"
+#include "flang/Runtime/CUDA/allocatable.h"
#include "flang/Runtime/CUDA/descriptor.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
@@ -107,6 +110,33 @@ struct CUFDeviceAddressOpConversion
const mlir::SymbolTable &symTab;
};
+struct CUFDeviceIsActiveOpConversion
+ : public mlir::OpRewritePattern<cuf::DeviceIsActiveOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ mlir::LogicalResult
+ matchAndRewrite(cuf::DeviceIsActiveOp op,
+ mlir::PatternRewriter &rewriter) const override {
+ mlir::Location loc = op.getLoc();
+ // Fold to false in device code so host-only cleanup (and nested calls
+ // such as cudaDeviceSynchronize) is DCE'd from GPU functions.
+ if (op->getParentOfType<mlir::gpu::GPUModuleOp>() ||
+ cuf::isCUDADeviceContext(op.getOperation())) {
+ rewriter.replaceOpWithNewOp<mlir::arith::ConstantOp>(
+ op, rewriter.getBoolAttr(false));
+ return mlir::success();
+ }
+ auto mod = op->getParentOfType<mlir::ModuleOp>();
+ fir::FirOpBuilder builder(rewriter, mod);
+ mlir::func::FuncOp callee =
+ fir::runtime::getRuntimeFunc<mkRTKey(CUFDeviceIsActive)>(loc, builder);
+ auto call = fir::CallOp::create(rewriter, loc, callee, mlir::ValueRange{});
+ rewriter.replaceOp(op, createConvertOp(rewriter, loc, rewriter.getI1Type(),
+ call.getResult(0)));
+ return mlir::success();
+ }
+};
+
class CUFOpConversionLate
: public fir::impl::CUFOpConversionLateBase<CUFOpConversionLate> {
using CUFOpConversionLateBase::CUFOpConversionLateBase;
@@ -125,6 +155,7 @@ class CUFOpConversionLate
mlir::gpu::GPUDialect>();
patterns.insert<CUFDeviceAddressOpConversion>(patterns.getContext(),
symtab);
+ patterns.insert<CUFDeviceIsActiveOpConversion>(patterns.getContext());
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
std::move(patterns)))) {
mlir::emitError(mlir::UnknownLoc::get(ctx),
diff --git a/flang/test/Fir/cuf.mlir b/flang/test/Fir/cuf.mlir
index f80a70eca34a37..75b9a1d20868bb 100644
--- a/flang/test/Fir/cuf.mlir
+++ b/flang/test/Fir/cuf.mlir
@@ -111,4 +111,13 @@ gpu.module @cuda_device_mod {
}
}
-// CHECK-COUNT-2: cuf.shared_memory
+// CHECK-COUNT-2: cuf.shared_memory
+
+// -----
+
+func.func @_QPdevice_is_active() {
+ %0 = cuf.device_is_active : i1
+ return
+}
+
+// CHECK: %{{.*}} = cuf.device_is_active : i1
diff --git a/flang/test/Lower/CUDA/cuda-allocatable.cuf b/flang/test/Lower/CUDA/cuda-allocatable.cuf
index 7e03a56e5d94b5..e0c53cdb575de7 100644
--- a/flang/test/Lower/CUDA/cuda-allocatable.cuf
+++ b/flang/test/Lower/CUDA/cuda-allocatable.cuf
@@ -118,7 +118,7 @@ end subroutine
! CHECK: %[[LOAD_B:.*]] = fir.load %[[BOX_B_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
! CHECK: fir.call @_FortranAAllocatableSetBounds
! CHECK: %{{.*}} = cuf.allocate %[[BOX_A_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> source(%[[LOAD_B]] : !fir.box<!fir.heap<!fir.array<?xf32>>>) {data_attr = #cuf.cuda<device>} -> i32
-! CHECK: fir.call @_FortranACUFDeviceIsActive()
+! CHECK: cuf.device_is_active : i1
! CHECK: fir.if
! CHECK: fir.if %{{.*}} {
! CHECK: %{{.*}} = cuf.deallocate %[[BOX_A_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
diff --git a/flang/test/Lower/CUDA/cuda-derived.cuf b/flang/test/Lower/CUDA/cuda-derived.cuf
index be0827a2204b73..6fb5b997c095d9 100644
--- a/flang/test/Lower/CUDA/cuda-derived.cuf
+++ b/flang/test/Lower/CUDA/cuda-derived.cuf
@@ -15,7 +15,7 @@ contains
! CHECK-LABEL: func.func @_QMm1Psub1()
! CHECK: %[[ALLOC:.*]] = cuf.alloc !fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}> {bindc_name = "a", data_attr = #cuf.cuda<managed>, uniq_name = "_QMm1Fsub1Ea"} -> !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>
! CHECK: %[[DECL:.*]]:2 = hlfir.declare %[[ALLOC]] {data_attr = #cuf.cuda<managed>, uniq_name = "_QMm1Fsub1Ea"} : (!fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>) -> (!fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>, !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>)
-! CHECK: %[[ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK: %[[ACTIVE:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[ACTIVE]] {
! CHECK: fir.call @_FortranADestroy
! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>> {data_attr = #cuf.cuda<managed>}
diff --git a/flang/test/Lower/CUDA/cuda-return01.cuf b/flang/test/Lower/CUDA/cuda-return01.cuf
index 847a22cc18229c..cb037bd19e80de 100644
--- a/flang/test/Lower/CUDA/cuda-return01.cuf
+++ b/flang/test/Lower/CUDA/cuda-return01.cuf
@@ -9,7 +9,7 @@ end
! CHECK: func.func @_QPsub1()
! CHECK: %[[DECL:.*]]:2 = hlfir.declare
-! CHECK: %[[SUB1_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK: %[[SUB1_ACTIVE:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[SUB1_ACTIVE]] {
! CHECK: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
! CHECK-NEXT: }
@@ -22,7 +22,7 @@ subroutine sub2
end
! CHECK-LABEL: func.func @_QPsub2()
-! CHECK: %[[SUB2_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK: %[[SUB2_ACTIVE:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[SUB2_ACTIVE]] {
! CHECK: fir.if
! CHECK: cuf.deallocate
@@ -41,7 +41,7 @@ end
! CHECK-LABEL: func.func @_QQmain() attributes {fir.bindc_name = "MAIN"}
! 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>>>>
! CHECK: %[[PTR:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFEp"}
-! CHECK: %[[ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK: %[[ACTIVE:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[ACTIVE]] {
! CHECK-NEXT: fir.call @_QPcudadevicesynchronize()
! CHECK-NEXT: cuf.free %[[PTR]]#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>{{.*}}
diff --git a/flang/test/Lower/CUDA/cuda-return02.cuf b/flang/test/Lower/CUDA/cuda-return02.cuf
index aa82cd5039cb2e..d34a44f29884fd 100644
--- a/flang/test/Lower/CUDA/cuda-return02.cuf
+++ b/flang/test/Lower/CUDA/cuda-return02.cuf
@@ -17,14 +17,14 @@ end
! CHECK: %[[DECL:.*]]:2 = hlfir.declare
! CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
! CHECK-NEXT: ^bb1:
-! CHECK-NEXT: %[[ACTIVE1:.*]] = fir.call @_FortranACUFDeviceIsActive()
+! CHECK-NEXT: %[[ACTIVE1:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[ACTIVE1]] {
! CHECK-NEXT: fir.call @_QPcudadevicesynchronize() proc_attrs<intrinsic> fastmath<contract> : () -> i32
! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
! CHECK-NEXT: }
! CHECK-NEXT: return
! CHECK-NEXT: ^bb2:
-! CHECK-NEXT: %[[ACTIVE2:.*]] = fir.call @_FortranACUFDeviceIsActive()
+! CHECK-NEXT: %[[ACTIVE2:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[ACTIVE2]] {
! CHECK-NEXT: fir.call @_QPcudadevicesynchronize() proc_attrs<intrinsic> fastmath<contract> : () -> i32
! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
@@ -52,7 +52,7 @@ end
! CHECK: ^bb2:
! CHECK: cf.br ^bb3
! CHECK: ^bb3:
-! CHECK-NEXT: %[[SUB_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive()
+! CHECK-NEXT: %[[SUB_ACTIVE:.*]] = cuf.device_is_active : i1
! CHECK-NEXT: fir.if %[[SUB_ACTIVE]] {
! CHECK-NEXT: fir.call @_QPcudadevicesynchronize() proc_attrs<intrinsic> fastmath<contract> : () -> i32
! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
>From ff12a29b0333dbf2c85bc4d75da59c2ca784cd99 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 22 Sep 2026 21:09:59 -0700
Subject: [PATCH 2/2] Add tests
---
flang/test/Lower/CUDA/cuda-return01.cuf | 29 +++++++++++++++++++++++++
1 file changed, 29 insertions(+)
diff --git a/flang/test/Lower/CUDA/cuda-return01.cuf b/flang/test/Lower/CUDA/cuda-return01.cuf
index cb037bd19e80de..465c001e15f5f6 100644
--- a/flang/test/Lower/CUDA/cuda-return01.cuf
+++ b/flang/test/Lower/CUDA/cuda-return01.cuf
@@ -1,4 +1,5 @@
! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir -fcuda %s -o - | fir-opt --cuf-convert-late | FileCheck %s --check-prefix=LATE
! Check if finalization works with a return statement
@@ -16,6 +17,15 @@ end
! CHECK: return
! CHECK: }
+! LATE: func.func @_QPsub1()
+! LATE: %[[DECL:.*]]:2 = hlfir.declare
+! LATE: %[[SUB1_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() : () -> i1
+! LATE-NEXT: fir.if %[[SUB1_ACTIVE]] {
+! LATE: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
+! LATE-NEXT: }
+! LATE: return
+! LATE: }
+
subroutine sub2
integer, device, allocatable :: a(:)
return
@@ -29,6 +39,14 @@ end
! CHECK: cuf.free{{.*}}
! CHECK-NEXT: }
+! LATE-LABEL: func.func @_QPsub2()
+! LATE: %[[SUB2_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() : () -> i1
+! LATE-NEXT: fir.if %[[SUB2_ACTIVE]] {
+! LATE: fir.if
+! LATE: cuf.deallocate
+! LATE: cuf.free{{.*}}
+! LATE-NEXT: }
+
program main
integer, allocatable, device :: a(:)
integer, device, pointer :: p(:)
@@ -49,3 +67,14 @@ end
! CHECK: cuf.free{{.*}}
! CHECK-NEXT: }
! CHECK-NEXT: return
+
+! LATE-LABEL: func.func @_QQmain() attributes {fir.bindc_name = "MAIN"}
+! LATE: %[[PTR:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFEp"}
+! LATE: %[[ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() : () -> i1
+! LATE-NEXT: fir.if %[[ACTIVE]] {
+! LATE-NEXT: fir.call @_QPcudadevicesynchronize()
+! LATE-NEXT: cuf.free %[[PTR]]#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>{{.*}}
+! LATE: cuf.deallocate
+! LATE: cuf.free{{.*}}
+! LATE-NEXT: }
+! LATE-NOT: cuf.device_is_active
More information about the flang-commits
mailing list