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

Author: clementval

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

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