From f8a9973f8c1ef60281ace6f3cfeb24d9dcd5b3c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Valentin=20Clement=20=28=E3=83=90=E3=83=AC=E3=83=B3?=
 =?UTF-8?q?=E3=82=BF=E3=82=A4=E3=83=B3=20=E3=82=AF=E3=83=AC=E3=83=A1?=
 =?UTF-8?q?=E3=83=B3=29?= <clementval@gmail.com>
Date: Fri, 3 May 2024 11:25:34 -0700
Subject: [PATCH] [flang][cuda] Add verifier for cuda_alloc/cuda_free (#90983)

Adding a verifier to check the associated cuda attribute.
---
 .../include/flang/Optimizer/Dialect/FIROps.td  |  4 ++++
 flang/lib/Optimizer/Dialect/FIROps.cpp         | 13 +++++++++++++
 flang/test/Fir/cuf-invalid.fir                 | 18 ++++++++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index dc38e56d93c664..64c5e360b28f79 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -3364,6 +3364,8 @@ def fir_CUDAAllocOp : fir_Op<"cuda_alloc", [AttrSizedOperandSegments,
       CArg<"mlir::ValueRange", "{}">:$typeparams,
       CArg<"mlir::ValueRange", "{}">:$shape,
       CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>];
+
+  let hasVerifier = 1;
 }
 
 def fir_CUDAFreeOp : fir_Op<"cuda_free", [MemoryEffects<[MemFree]>]> {
@@ -3381,6 +3383,8 @@ def fir_CUDAFreeOp : fir_Op<"cuda_free", [MemoryEffects<[MemFree]>]> {
   );
 
   let assemblyFormat = "$devptr `:` qualified(type($devptr)) attr-dict";
+
+  let hasVerifier = 1;
 }
 
 #endif
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 5e6c18af2dd0f9..edf7f7f4b1a966 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4048,6 +4048,19 @@ void fir::CUDAAllocOp::build(
   result.addAttributes(attributes);
 }
 
+template <typename Op>
+static mlir::LogicalResult checkCudaAttr(Op op) {
+  if (op.getCudaAttr() == fir::CUDADataAttribute::Device ||
+      op.getCudaAttr() == fir::CUDADataAttribute::Managed ||
+      op.getCudaAttr() == fir::CUDADataAttribute::Unified)
+    return mlir::success();
+  return op.emitOpError("expect device, managed or unified cuda attribute");
+}
+
+mlir::LogicalResult fir::CUDAAllocOp::verify() { return checkCudaAttr(*this); }
+
+mlir::LogicalResult fir::CUDAFreeOp::verify() { return checkCudaAttr(*this); }
+
 //===----------------------------------------------------------------------===//
 // FIROpsDialect
 //===----------------------------------------------------------------------===//
diff --git a/flang/test/Fir/cuf-invalid.fir b/flang/test/Fir/cuf-invalid.fir
index 6c533a32ccf9ba..5a12e3c1a4bf20 100644
--- a/flang/test/Fir/cuf-invalid.fir
+++ b/flang/test/Fir/cuf-invalid.fir
@@ -85,3 +85,21 @@ func.func @_QPsub1() {
   %13 = fir.cuda_deallocate %11 : !fir.ref<!fir.box<none>> errmsg(%16 : !fir.box<none>) {cuda_attr = #fir.cuda<device>} -> i32
   return
 }
+
+// -----
+
+func.func @_QPsub1() {
+  // expected-error@+1{{'fir.cuda_alloc' op expect device, managed or unified cuda attribute}}
+  %0 = fir.cuda_alloc f32 {bindc_name = "r", cuda_attr = #fir.cuda<pinned>, uniq_name = "_QFsub1Er"} -> !fir.ref<f32>
+  fir.cuda_free %0 : !fir.ref<f32> {cuda_attr = #fir.cuda<constant>}
+  return
+}
+
+// -----
+
+func.func @_QPsub1() {
+  %0 = fir.cuda_alloc f32 {bindc_name = "r", cuda_attr = #fir.cuda<device>, uniq_name = "_QFsub1Er"} -> !fir.ref<f32>
+  // expected-error@+1{{'fir.cuda_free' op expect device, managed or unified cuda attribute}}
+  fir.cuda_free %0 : !fir.ref<f32> {cuda_attr = #fir.cuda<constant>}
+  return
+}