[Mlir-commits] [flang] [mlir] [acc][flang] Add isDeviceData APIs for device data detection (PR #176219)

[Razvan Lupusoru]

================
@@ -1487,4 +1489,131 @@ template bool OpenACCPointerLikeModel<fir::LLVMPointerType>::genStore(
     mlir::Value valueToStore,
     mlir::TypedValue<mlir::acc::PointerLikeType> destPtr) const;
 
+/// Helper function to check if a CUDA attribute represents device data.
+static bool isCUDADeviceAttribute(cuf::DataAttribute attr) {
+  return attr == cuf::DataAttribute::Device ||
+         attr == cuf::DataAttribute::Managed ||
+         attr == cuf::DataAttribute::Constant ||
+         attr == cuf::DataAttribute::Shared ||
+         attr == cuf::DataAttribute::Unified;
+}
+
+/// Helper function to check if an operation has CUDA device data attributes.
+static bool hasCUDADeviceDataAttr(mlir::Operation *op) {
+  if (!op)
+    return false;
+
+  // Check for CUF data attribute on the operation
+  if (auto dataAttr = cuf::getDataAttr(op)) {
+    if (isCUDADeviceAttribute(dataAttr.getValue()))
+      return true;
+  }
+
+  return false;
+}
+
+/// Check CUDA attributes on a function argument.
+static bool hasCUDADeviceAttrOnFuncArg(mlir::BlockArgument blockArg) {
+  auto *owner = blockArg.getOwner();
+  if (!owner)
+    return false;
+
+  auto *parentOp = owner->getParentOp();
+  if (!parentOp)
+    return false;
+
+  if (auto funcLike = mlir::dyn_cast<mlir::FunctionOpInterface>(parentOp)) {
+    unsigned argIndex = blockArg.getArgNumber();
+    if (argIndex < funcLike.getNumArguments()) {
+      if (auto attr = funcLike.getArgAttr(argIndex, cuf::getDataAttrName())) {
+        if (auto cudaAttr = mlir::dyn_cast<cuf::DataAttributeAttr>(attr))
+          return isCUDADeviceAttribute(cudaAttr.getValue());
+      }
+    }
+  }
+  return false;
+}
+
+/// Shared implementation for checking if a value represents device data.
+static bool isDeviceDataImpl(mlir::Value var) {
+  // Strip casts to find the underlying value.
+  mlir::Value currentVal = stripCasts(var, /*stripDeclare=*/false);
+
+  // Handle block arguments (function parameters)
+  if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(currentVal))
+    return hasCUDADeviceAttrOnFuncArg(blockArg);
+
+  mlir::Operation *defOp = currentVal.getDefiningOp();
+  if (!defOp)
+    return false;
+
+  // Check for CUDA attributes on the defining operation.
+  if (hasCUDADeviceDataAttr(defOp))
+    return true;
+
+  // Handle operations that access a partial entity - check if the base entity
+  // is device data.
+  if (auto partialAccess =
+          mlir::dyn_cast<mlir::acc::PartialEntityAccessOpInterface>(defOp)) {
+    if (mlir::Value base = partialAccess.getBaseEntity())
+      return isDeviceDataImpl(base);
+  }
+
+  // Handle fir.rebox - if the underlying box is device data, so is the result.
----------------
razvanlupusoru wrote:

Nice idea - done. I will rerun testing to make sure everything still works OK but seems better because for example this also handles box_addr.

https://github.com/llvm/llvm-project/pull/176219