[flang-commits] [flang] [flang][cuda] Add utility function cuf::hasDataAttr (PR #154422)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/154422

None

>From 99a6f5057ce34481a55cd71646f72dbffe5120f7 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 19 Aug 2025 14:02:48 -0700
Subject: [PATCH] [flang][cuda] Add utility function cuf::hasDataAttr

---
 .../Optimizer/Dialect/CUF/Attributes/CUFAttr.h   |  8 ++++++++
 .../Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp | 16 ++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
index 85615a4d69157..1183fc1299810 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
+++ b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
@@ -20,6 +20,10 @@ namespace llvm {
 class StringRef;
 }
 
+namespace mlir {
+class Operation;
+}
+
 #include "flang/Optimizer/Dialect/CUF/Attributes/CUFEnumAttr.h.inc"
 
 #define GET_ATTRDEF_CLASSES
@@ -28,6 +32,7 @@ class StringRef;
 namespace cuf {
 
 /// Attribute to mark Fortran entities with the CUDA attribute.
+static constexpr llvm::StringRef dataAttrName = "data_attr";
 static constexpr llvm::StringRef getDataAttrName() { return "cuf.data_attr"; }
 static constexpr llvm::StringRef getProcAttrName() { return "cuf.proc_attr"; }
 
@@ -101,6 +106,9 @@ getProcAttribute(mlir::MLIRContext *mlirContext,
   return {};
 }
 
+/// Returns true if the operation has a data attribute with the given value.
+bool hasDataAttr(mlir::Operation *op);
+
 } // namespace cuf
 
 #endif // FORTRAN_OPTIMIZER_DIALECT_CUF_CUFATTR_H
diff --git a/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp b/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
index 52c733dcad6f0..9314c46a3bdf7 100644
--- a/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
@@ -16,6 +16,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #include "flang/Optimizer/Dialect/CUF/Attributes/CUFEnumAttr.cpp.inc"
@@ -29,4 +30,19 @@ void CUFDialect::registerAttributes() {
                 LaunchBoundsAttr, ProcAttributeAttr>();
 }
 
+bool hasDataAttr(mlir::Operation *op, cuf::DataAttribute value) {
+  if (!op)
+    return false;
+
+  cuf::DataAttributeAttr dataAttr =
+      op->getAttrOfType<cuf::DataAttributeAttr>(cuf::getDataAttrName());
+  // When the attribute is declared on the operation, it doesn't have a prefix.
+  if (!dataAttr)
+    dataAttr = op->getAttrOfType<cuf::DataAttributeAttr>(cuf::dataAttrName);
+  if (!dataAttr)
+    return false;
+
+  return dataAttr.getValue() == value;
+}
+
 } // namespace cuf