[flang-commits] [flang] [flang][cuda][NFC] Add getDataAttr helper (PR #154586)

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

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

None

>From a1dec1e3cb5346400e7939f180214b41a5efef39 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 20 Aug 2025 10:53:17 -0700
Subject: [PATCH] [flang][cuda][NFC] Add getDataAttr helper

---
 .../Dialect/CUF/Attributes/CUFAttr.h          |  3 +++
 .../Dialect/CUF/Attributes/CUFAttr.cpp        | 23 +++++++++++--------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
index c7d89420eab08..4a250d1cc6c54 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
+++ b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
@@ -106,6 +106,9 @@ getProcAttribute(mlir::MLIRContext *mlirContext,
   return {};
 }
 
+/// Returns the data attribute if the operation has one.
+cuf::DataAttributeAttr getDataAttr(mlir::Operation *op);
+
 /// Returns true if the operation has a data attribute with the given value.
 bool hasDataAttr(mlir::Operation *op, cuf::DataAttribute value);
 
diff --git a/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp b/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
index 9314c46a3bdf7..aafca45b87bde 100644
--- a/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp
@@ -30,19 +30,24 @@ void CUFDialect::registerAttributes() {
                 LaunchBoundsAttr, ProcAttributeAttr>();
 }
 
-bool hasDataAttr(mlir::Operation *op, cuf::DataAttribute value) {
+cuf::DataAttributeAttr getDataAttr(mlir::Operation *op) {
   if (!op)
-    return false;
+    return {};
+
+  if (auto dataAttr = op->getAttrOfType<cuf::DataAttributeAttr>(cuf::getDataAttrName()))
+    return dataAttr;
 
-  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;
+  if (auto dataAttr = op->getAttrOfType<cuf::DataAttributeAttr>(cuf::dataAttrName))
+    return dataAttr;
+    
+  return {};
+}
 
-  return dataAttr.getValue() == value;
+bool hasDataAttr(mlir::Operation *op, cuf::DataAttribute value) {
+  if (auto dataAttr = getDataAttr(op))
+    return dataAttr.getValue() == value;
+  return false;
 }
 
 } // namespace cuf