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

Author: clementval

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);
 

flang/lib/Optimizer/Dialect/CUF/Attributes/CUFAttr.cpp

@@ -30,19 +30,26 @@ 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 dataAttr.getValue() == value;
+  return {};
+}
+
+bool hasDataAttr(mlir::Operation *op, cuf::DataAttribute value) {
+  if (auto dataAttr = getDataAttr(op))
+    return dataAttr.getValue() == value;
+  return false;
 }
 
 } // namespace cuf