[flang-commits] [flang] a9c73f6 - [flang][cuda] Add fir.cuda_alloc/fir.cuda_free operations (#90525)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-04-30T08:29:51-07:00
New Revision: a9c73f66ce96421ba1e8952950d6df74e9043589

URL: https://github.com/llvm/llvm-project/commit/a9c73f66ce96421ba1e8952950d6df74e9043589
DIFF: https://github.com/llvm/llvm-project/commit/a9c73f66ce96421ba1e8952950d6df74e9043589.diff

LOG: [flang][cuda] Add fir.cuda_alloc/fir.cuda_free operations (#90525)

This patch introduces fir.cuda_alloc/fir.cuda_free. These operations
will be used instead of fir.alloca for local CUDA device, managed and
unified variables.

Added: 
    

Modified: 
    flang/include/flang/Optimizer/Dialect/FIROps.td
    flang/test/Fir/cuf.mlir

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 496193e25cab61..24950322c3ca9e 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<StrAttr>:$uniq_name,
+    OptionalAttr<StrAttr>:$bindc_name,
+    Variadic<AnyIntegerType>:$typeparams,
+    Variadic<AnyIntegerType>:$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<mlir::NamedAttribute>", "{}">:$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<AnyReferenceLike, "", [MemFree]>:$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 71f0652067facf..8e2346def43ea8 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<!fir.box<none>> errmsg(%{{.*}} : !fir.box<none>) {cuda_attr = #fir.cuda<device>, hasStat} -> i32
 // CHECK: fir.cuda_deallocate %{{.*}} : !fir.ref<!fir.box<none>> errmsg(%{{.*}} : !fir.box<none>) {cuda_attr = #fir.cuda<device>, hasStat} -> i32
+
+// -----
+
+func.func @_QPsub1() {
+  %0 = fir.cuda_alloc f32 {bindc_name = "r", cuda_attr = #fir.cuda<device>, uniq_name = "_QFsub1Er"} -> !fir.ref<f32>
+  fir.cuda_free %0 : !fir.ref<f32> {cuda_attr = #fir.cuda<device>}
+  return
+}
+
+// CHECK: fir.cuda_alloc
+// CHECK: fir.cuda_free
+