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

This commit is contained in:
Valentin Clement (バレンタイン クレメン) 2025-08-19 14:25:32 -07:00 committed by GitHub
parent 402109e1c4
commit af8a149546
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 0 deletions

View File

@ -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

View File

@ -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