diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td index 496193e25cab6..24950322c3ca9 100644 --- a/flang/include/flang/Optimizer/Dialect/FIROps.td +++ b/flang/include/flang/Optimizer/Dialect/FIROps.td @@ -3247,4 +3247,58 @@ def fir_CUDADeallocateOp : fir_Op<"cuda_deallocate", let hasVerifier = 1; } +def fir_CUDAAllocOp : fir_Op<"cuda_alloc", [AttrSizedOperandSegments, + MemoryEffects<[MemAlloc]>]> { + let summary = "Allocate an object on device"; + + let description = [{ + This is a drop in replacement for fir.alloca and fir.allocmem for device + object. Any device, managed or unified object declared in an host + subprogram needs to be allocated in the device memory through runtime calls. + The fir.cuda_alloc is an abstraction to the runtime calls and works together + with fir.cuda_free. + }]; + + let arguments = (ins + TypeAttr:$in_type, + OptionalAttr:$uniq_name, + OptionalAttr:$bindc_name, + Variadic:$typeparams, + Variadic:$shape, + fir_CUDADataAttributeAttr:$cuda_attr + ); + + let results = (outs fir_ReferenceType:$ptr); + + let assemblyFormat = [{ + $in_type (`(` $typeparams^ `:` type($typeparams) `)`)? + (`,` $shape^ `:` type($shape) )? attr-dict `->` qualified(type($ptr)) + }]; + + let builders = [ + OpBuilder<(ins "mlir::Type":$inType, "llvm::StringRef":$uniqName, + "llvm::StringRef":$bindcName, + "fir::CUDADataAttributeAttr":$cudaAttr, + CArg<"mlir::ValueRange", "{}">:$typeparams, + CArg<"mlir::ValueRange", "{}">:$shape, + CArg<"llvm::ArrayRef", "{}">:$attributes)>]; +} + +def fir_CUDAFreeOp : fir_Op<"cuda_free", [MemoryEffects<[MemFree]>]> { + let summary = "Free a device allocated object"; + + let description = [{ + The fir.cuda_free operation frees the memory allocated by fir.cuda_alloc. + This is used for non-allocatable device, managed and unified device + variables declare in host subprogram. + }]; + + let arguments = (ins + Arg:$devptr, + fir_CUDADataAttributeAttr:$cuda_attr + ); + + let assemblyFormat = "$devptr `:` qualified(type($devptr)) attr-dict"; +} + #endif diff --git a/flang/test/Fir/cuf.mlir b/flang/test/Fir/cuf.mlir index 71f0652067fac..8e2346def43ea 100644 --- a/flang/test/Fir/cuf.mlir +++ b/flang/test/Fir/cuf.mlir @@ -74,3 +74,15 @@ func.func @_QPsub1() { // CHECK: fir.cuda_allocate %{{.*}} : !fir.ref> errmsg(%{{.*}} : !fir.box) {cuda_attr = #fir.cuda, hasStat} -> i32 // CHECK: fir.cuda_deallocate %{{.*}} : !fir.ref> errmsg(%{{.*}} : !fir.box) {cuda_attr = #fir.cuda, hasStat} -> i32 + +// ----- + +func.func @_QPsub1() { + %0 = fir.cuda_alloc f32 {bindc_name = "r", cuda_attr = #fir.cuda, uniq_name = "_QFsub1Er"} -> !fir.ref + fir.cuda_free %0 : !fir.ref {cuda_attr = #fir.cuda} + return +} + +// CHECK: fir.cuda_alloc +// CHECK: fir.cuda_free +