diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td index dff1cdb20cbfe..c181c7ed62dff 100644 --- a/flang/include/flang/Optimizer/Dialect/FIROps.td +++ b/flang/include/flang/Optimizer/Dialect/FIROps.td @@ -3190,4 +3190,36 @@ def fir_CUDADataTransferOp : fir_Op<"cuda_data_transfer", []> { }]; } +def fir_CUDAAllocateOp : fir_Op<"cuda_allocate", [AttrSizedOperandSegments, + MemoryEffects<[MemAlloc]>]> { + let summary = "Perform the device allocation of data of an allocatable"; + + let description = [{ + The fir.cuda_allocate operation performs the allocation on the device + of the data of an allocatable. The descriptor passed to the operation + is initialized before with the standard flang runtime calls. + }]; + + let arguments = (ins Arg:$box, + Arg, "", [MemWrite]>:$errmsg, + Optional:$stream, + Arg, "", [MemWrite]>:$pinned, + Arg, "", [MemRead]>:$source, + fir_CUDADataAttributeAttr:$cuda_attr, + UnitAttr:$hasStat); + + let results = (outs AnyIntegerType:$stat); + + let assemblyFormat = [{ + $box `:` qualified(type($box)) + ( `source` `(` $source^ `:` qualified(type($source) )`)` )? + ( `errmsg` `(` $errmsg^ `:` type($errmsg) `)` )? + ( `stream` `(` $stream^ `:` type($stream) `)` )? + ( `pinned` `(` $pinned^ `:` type($pinned) `)` )? + attr-dict `->` type($stat) + }]; + + let hasVerifier = 1; +} + #endif diff --git a/flang/include/flang/Optimizer/Dialect/FIRTypes.td b/flang/include/flang/Optimizer/Dialect/FIRTypes.td index 4c6a8064991ab..3b876e4642da9 100644 --- a/flang/include/flang/Optimizer/Dialect/FIRTypes.td +++ b/flang/include/flang/Optimizer/Dialect/FIRTypes.td @@ -625,6 +625,7 @@ def AnyRefOrBoxLike : TypeConstraint, "any reference or box">; +def AnyRefOrBoxType : Type; def AnyShapeLike : TypeConstraint, "any legal shape type">; diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp index 8ab74103cb6a8..88710880174d2 100644 --- a/flang/lib/Optimizer/Dialect/FIROps.cpp +++ b/flang/lib/Optimizer/Dialect/FIROps.cpp @@ -3993,6 +3993,25 @@ mlir::LogicalResult fir::CUDAKernelOp::verify() { return mlir::success(); } +mlir::LogicalResult fir::CUDAAllocateOp::verify() { + if (getPinned() && getStream()) + return emitOpError("pinned and stream cannot appears at the same time"); + if (!fir::unwrapRefType(getBox().getType()).isa()) + return emitOpError( + "expect box to be a reference to/or a class or box type value"); + if (getSource() && + !fir::unwrapRefType(getSource().getType()).isa()) + return emitOpError( + "expect source to be a reference to/or a class or box type value"); + if (getErrmsg() && + !fir::unwrapRefType(getErrmsg().getType()).isa()) + return emitOpError( + "expect errmsg to be a reference to/or a box type value"); + if (getErrmsg() && !getHasStat()) + return emitOpError("expect stat attribute when errmsg is provided"); + return mlir::success(); +} + //===----------------------------------------------------------------------===// // FIROpsDialect //===----------------------------------------------------------------------===// diff --git a/flang/test/Fir/cuf-invalid.fir b/flang/test/Fir/cuf-invalid.fir new file mode 100644 index 0000000000000..9c5ffe7176a3b --- /dev/null +++ b/flang/test/Fir/cuf-invalid.fir @@ -0,0 +1,50 @@ +// RUN: fir-opt -split-input-file -verify-diagnostics %s + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %1 = fir.alloca i32 + %pinned = fir.alloca i1 + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %s = fir.load %1 : !fir.ref + // expected-error@+1{{'fir.cuda_allocate' op pinned and stream cannot appears at the same time}} + %13 = fir.cuda_allocate %11 : !fir.ref> stream(%s : i32) pinned(%pinned : !fir.ref) {cuda_attr = #fir.cuda} -> i32 + return +} + +// ----- + +func.func @_QPsub1() { + %1 = fir.alloca i32 + // expected-error@+1{{'fir.cuda_allocate' op expect box to be a reference to/or a class or box type value}} + %2 = fir.cuda_allocate %1 : !fir.ref {cuda_attr = #fir.cuda} -> i32 + return +} + +// ----- + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %c100 = arith.constant 100 : index + %7 = fir.alloca !fir.char<1,100> {bindc_name = "msg", uniq_name = "_QFsub1Emsg"} + %8:2 = hlfir.declare %7 typeparams %c100 {uniq_name = "_QFsub1Emsg"} : (!fir.ref>, index) -> (!fir.ref>, !fir.ref>) + %9 = fir.embox %8#1 : (!fir.ref>) -> !fir.box> + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %16 = fir.convert %9 : (!fir.box>) -> !fir.box + // expected-error@+1{{'fir.cuda_allocate' op expect stat attribute when errmsg is provided}} + %13 = fir.cuda_allocate %11 : !fir.ref> errmsg(%16 : !fir.box) {cuda_attr = #fir.cuda} -> i32 + return +} + +// ----- + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %1 = fir.alloca i32 + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + // expected-error@+1{{'fir.cuda_allocate' op expect errmsg to be a reference to/or a box type value}} + %13 = fir.cuda_allocate %11 : !fir.ref> errmsg(%1 : !fir.ref) {cuda_attr = #fir.cuda, hasStat} -> i32 + return +} diff --git a/flang/test/Fir/cuf.mlir b/flang/test/Fir/cuf.mlir new file mode 100644 index 0000000000000..67eff31b35b2b --- /dev/null +++ b/flang/test/Fir/cuf.mlir @@ -0,0 +1,70 @@ +// RUN: fir-opt --split-input-file %s | fir-opt --split-input-file | FileCheck %s + +// Simple round trip test of operations. + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %13 = fir.cuda_allocate %11 : !fir.ref> {cuda_attr = #fir.cuda} -> i32 + return +} + +// CHECK: fir.cuda_allocate %{{.*}} : !fir.ref> {cuda_attr = #fir.cuda} -> i32 + +// ----- + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %1 = fir.alloca i32 + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %s = fir.load %1 : !fir.ref + %13 = fir.cuda_allocate %11 : !fir.ref> stream(%s : i32) {cuda_attr = #fir.cuda} -> i32 + return +} + +// CHECK: fir.cuda_allocate %{{.*}} : !fir.ref> stream(%{{.*}} : i32) {cuda_attr = #fir.cuda} -> i32 + +// ----- + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %1 = fir.alloca !fir.box>> {bindc_name = "b", uniq_name = "_QFsub1Eb"} + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %5:2 = hlfir.declare %1 {fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %12 = fir.convert %5#1 : (!fir.ref>>>) -> !fir.ref> + %13 = fir.cuda_allocate %11 : !fir.ref> source(%12 : !fir.ref>) {cuda_attr = #fir.cuda} -> i32 + return +} + +// CHECK: fir.cuda_allocate %{{.*}} : !fir.ref> source(%{{.*}} : !fir.ref>) {cuda_attr = #fir.cuda} -> i32 + +// ----- + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %pinned = fir.alloca i1 + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %13 = fir.cuda_allocate %11 : !fir.ref> pinned(%pinned : !fir.ref) {cuda_attr = #fir.cuda} -> i32 + return +} + +// CHECK: fir.cuda_allocate %{{.*}} : !fir.ref> pinned(%{{.*}} : !fir.ref) {cuda_attr = #fir.cuda} -> i32 + +// ----- + +func.func @_QPsub1() { + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFsub1Ea"} + %4:2 = hlfir.declare %0 {cuda_attr = #fir.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFsub1Ea"} : (!fir.ref>>>) -> (!fir.ref>>>, !fir.ref>>>) + %c100 = arith.constant 100 : index + %7 = fir.alloca !fir.char<1,100> {bindc_name = "msg", uniq_name = "_QFsub1Emsg"} + %8:2 = hlfir.declare %7 typeparams %c100 {uniq_name = "_QFsub1Emsg"} : (!fir.ref>, index) -> (!fir.ref>, !fir.ref>) + %9 = fir.embox %8#1 : (!fir.ref>) -> !fir.box> + %11 = fir.convert %4#1 : (!fir.ref>>>) -> !fir.ref> + %16 = fir.convert %9 : (!fir.box>) -> !fir.box + %13 = fir.cuda_allocate %11 : !fir.ref> errmsg(%16 : !fir.box) {cuda_attr = #fir.cuda, hasStat} -> i32 + return +}