[flang][cuda][NFC] Add getDataAttr helper

[clementval]

No description provided.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/154586.diff

2 Files Affected:

• (modified) flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h (+3)
• (modified) flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp (+14-9)

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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.