Automerge: [mlir][acc] Add isValidValueUse to OpenACCSupport (#171538)

Add a new API `isValidValueUse ` to OpenACCSupport. This is used in
ACCImplicitData to check value that are already legal in the OpenACC
region and do not require implicit clause to be generated. An example
would be a CUDA Fortran device variable that is already on the GPU.

Author: clementval

mlir/include/mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h

@@ -85,6 +85,9 @@ struct OpenACCSupportTraits {
     /// Check if a symbol use is valid for use in an OpenACC region.
     virtual bool isValidSymbolUse(Operation *user, SymbolRefAttr symbol,
                                   Operation **definingOpPtr) = 0;
+
+    /// Check if a value use is legal in an OpenACC region.
+    virtual bool isValidValueUse(Value v, mlir::Region &region) = 0;
   };
 
   /// SFINAE helpers to detect if implementation has optional methods
@@ -97,6 +100,14 @@ struct OpenACCSupportTraits {
       llvm::is_detected<isValidSymbolUse_t, ImplT, Operation *, SymbolRefAttr,
                         Operation **>;
 
+  template <typename ImplT, typename... Args>
+  using isValidValueUse_t =
+      decltype(std::declval<ImplT>().isValidValueUse(std::declval<Args>()...));
+
+  template <typename ImplT>
+  using has_isValidValueUse =
+      llvm::is_detected<isValidValueUse_t, ImplT, Value, Region &>;
+
   /// This class wraps a concrete OpenACCSupport implementation and forwards
   /// interface calls to it. This provides type erasure, allowing different
   /// implementation types to be used interchangeably without inheritance.
@@ -128,6 +139,13 @@ struct OpenACCSupportTraits {
         return acc::isValidSymbolUse(user, symbol, definingOpPtr);
     }
 
+    bool isValidValueUse(Value v, Region &region) final {
+      if constexpr (has_isValidSymbolUse<ImplT>::value)
+        return impl.isValidValueUse(v, region);
+      else
+        return false;
+    }
+
   private:
     ImplT impl;
   };
@@ -189,6 +207,12 @@ class OpenACCSupport {
   bool isValidSymbolUse(Operation *user, SymbolRefAttr symbol,
                         Operation **definingOpPtr = nullptr);
 
+  /// Check if a value use is legal in an OpenACC region.
+  ///
+  /// \param v The MLIR value to check for legality.
+  /// \param region The MLIR region in which the legality is checked.
+  bool isValidValueUse(Value v, Region &region);
+
   /// Signal that this analysis should always be preserved so that
   /// underlying implementation registration is not lost.
   bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {

mlir/lib/Dialect/OpenACC/Analysis/OpenACCSupport.cpp

@@ -48,5 +48,11 @@ bool OpenACCSupport::isValidSymbolUse(Operation *user, SymbolRefAttr symbol,
   return acc::isValidSymbolUse(user, symbol, definingOpPtr);
 }
 
+bool OpenACCSupport::isValidValueUse(Value v, Region &region) {
+  if (impl)
+    return impl->isValidValueUse(v, region);
+  return false;
+}
+
 } // namespace acc
 } // namespace mlir

mlir/lib/Dialect/OpenACC/Transforms/ACCImplicitData.cpp

@@ -254,9 +254,10 @@ class ACCImplicitData : public acc::impl::ACCImplicitDataBase<ACCImplicitData> {
 
   /// Generates the implicit data ops for a compute construct.
   template <typename OpT>
-  void generateImplicitDataOps(
-      ModuleOp &module, OpT computeConstructOp,
-      std::optional<acc::ClauseDefaultValue> &defaultClause);
+  void
+  generateImplicitDataOps(ModuleOp &module, OpT computeConstructOp,
+                          std::optional<acc::ClauseDefaultValue> &defaultClause,
+                          acc::OpenACCSupport &accSupport);
 
   /// Generates a private recipe for a variable.
   acc::PrivateRecipeOp generatePrivateRecipe(ModuleOp &module, Value var,
@@ -277,12 +278,16 @@ class ACCImplicitData : public acc::impl::ACCImplicitDataBase<ACCImplicitData> {
 
 /// Determines if a variable is a candidate for implicit data mapping.
 /// Returns true if the variable is a candidate, false otherwise.
-static bool isCandidateForImplicitData(Value val, Region &accRegion) {
+static bool isCandidateForImplicitData(Value val, Region &accRegion,
+                                       acc::OpenACCSupport &accSupport) {
   // Ensure the variable is an allowed type for data clause.
   if (!acc::isPointerLikeType(val.getType()) &&
       !acc::isMappableType(val.getType()))
     return false;
 
+  if (accSupport.isValidValueUse(val, accRegion))
+    return false;
+
   // If this is already coming from a data clause, we do not need to generate
   // another.
   if (isa_and_nonnull<ACC_DATA_ENTRY_OPS>(val.getDefiningOp()))
@@ -683,7 +688,8 @@ static void insertInSortedOrder(SmallVector<Value> &sortedDataClauseOperands,
 template <typename OpT>
 void ACCImplicitData::generateImplicitDataOps(
     ModuleOp &module, OpT computeConstructOp,
-    std::optional<acc::ClauseDefaultValue> &defaultClause) {
+    std::optional<acc::ClauseDefaultValue> &defaultClause,
+    acc::OpenACCSupport &accSupport) {
   // Implicit data attributes are only applied if "[t]here is no default(none)
   // clause visible at the compute construct."
   if (defaultClause.has_value() &&
@@ -699,7 +705,7 @@ void ACCImplicitData::generateImplicitDataOps(
 
   // 2) Run the filtering to find relevant pointers that need copied.
   auto isCandidate{[&](Value val) -> bool {
-    return isCandidateForImplicitData(val, accRegion);
+    return isCandidateForImplicitData(val, accRegion, accSupport);
   }};
   auto candidateVars(
       llvm::to_vector(llvm::make_filter_range(liveInValues, isCandidate)));
@@ -763,6 +769,9 @@ void ACCImplicitData::generateImplicitDataOps(
 
 void ACCImplicitData::runOnOperation() {
   ModuleOp module = this->getOperation();
+
+  acc::OpenACCSupport &accSupport = getAnalysis<acc::OpenACCSupport>();
+
   module.walk([&](Operation *op) {
     if (isa<ACC_COMPUTE_CONSTRUCT_OPS, acc::KernelEnvironmentOp>(op)) {
       assert(op->getNumRegions() == 1 && "must have 1 region");
@@ -771,7 +780,7 @@ void ACCImplicitData::runOnOperation() {
       llvm::TypeSwitch<Operation *, void>(op)
           .Case<ACC_COMPUTE_CONSTRUCT_OPS, acc::KernelEnvironmentOp>(
               [&](auto op) {
-                generateImplicitDataOps(module, op, defaultClause);
+                generateImplicitDataOps(module, op, defaultClause, accSupport);
               })
           .Default([&](Operation *) {});
     }