[Clang][OpenMP] Add a util to implicitly map attach-ptr-exprs on target, when applicable.

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -9256,6 +9256,156 @@ class MappableExprsHandler {
     }
   }
 
+  /// On a target construct, if there's an implicit map on a struct, or that of
+  /// this[:], and an explicit map with a member of that struct/class as the
+  /// base-pointer, we need to make sure that base-pointer is implicitly mapped,
+  /// to make sure we don't map the full struct/class. For example:
+  ///
+  /// \code
+  ///   struct S {
+  ///     int dummy[10000];
+  ///     int *p;
+  ///     void f1() {
+  ///       #pragma omp target map(p[0:1])
+  ///       (void)this;
+  ///     }
+  ///   }; S s;
+  ///
+  ///   void f2() {
+  ///     #pragma omp target map(s.p[0:10])
+  ///     (void)s;
+  ///   }
+  /// \endcode
+  ///
+  /// Only `this-p` and `s.p` should be mapped in the two cases above.
+  //
+  // OpenMP 6.0: 7.9.6 map clause, pg 285
+  // If a list item with an implicitly determined data-mapping attribute does
+  // not have any corresponding storage in the device data environment prior to
+  // a task encountering the construct associated with the map clause, and one
+  // or more contiguous parts of the original storage are either list items or
+  // base pointers to list items that are explicitly mapped on the construct,
+  // only those parts of the original storage will have corresponding storage in
+  // the device data environment as a result of the map clauses on the
+  // construct.
+  void addImplicitMapForAttachPtrBaseIfMemberOfCapturedVD(
+      const ValueDecl *CapturedVD, MapDataArrayTy &DeclComponentLists,
+      SmallVectorImpl<
+          SmallVector<OMPClauseMappableExprCommon::MappableComponent, 8>>
+          &ComponentVectorStorage) const {
+    bool IsThisCapture = CapturedVD == nullptr;
+
+    for (const auto &ComponentsAndAttachPtr : AttachPtrExprMap) {
+      OMPClauseMappableExprCommon::MappableExprComponentListRef
+          ComponentsWithAttachPtr = ComponentsAndAttachPtr.first;
+      const Expr *AttachPtrExpr = ComponentsAndAttachPtr.second;
+      if (!AttachPtrExpr)
+        continue;
+
+      const auto *ME = dyn_cast<MemberExpr>(AttachPtrExpr);
+      if (!ME)
+        continue;
+
+      const Expr *Base = ME->getBase()->IgnoreParenImpCasts();
+
+      // If we are handling a "this" capture, then we are looking for
+      // attach-ptrs of form `this->p`, either explicitly or implicitly.
+      if (IsThisCapture && !ME->isImplicitCXXThis() && !isa<CXXThisExpr>(Base))
+        continue;
+
+      if (!IsThisCapture && (!isa<DeclRefExpr>(Base) ||
+                             cast<DeclRefExpr>(Base)->getDecl() != CapturedVD))
+        continue;
+
+      // For non-this captures, we are looking for attach-ptrs of form
+      // `s.p`.
+      // For non-this captures, we are looking for attach-ptrs like `s.p`.
+      if (!IsThisCapture && (ME->isArrow() || !isa<DeclRefExpr>(Base) ||
+                             cast<DeclRefExpr>(Base)->getDecl() != CapturedVD))
+        continue;
+
+      // Check if we have an existing map on either:
+      // this[:], s, this->p, or s.p, in which case, we don't need to add
+      // an implicit one for the attach-ptr s.p/this->p.
+      bool FoundExistingMap = false;
+      for (const MapData &ExistingL : DeclComponentLists) {
+        OMPClauseMappableExprCommon::MappableExprComponentListRef
+            ExistingComponents = std::get<0>(ExistingL);
+
+        if (ExistingComponents.empty())
+          continue;
+
+        // First check if we have a map like map(this->p) or map(s.p).
+        const auto &FirstComponent = ExistingComponents.front();
+        const Expr *FirstExpr = FirstComponent.getAssociatedExpression();
+
+        if (!FirstExpr)
+          continue;
+
+        // First check if we have a map like map(this->p) or map(s.p).
+        if (AttachPtrComparator.areEqual(FirstExpr, AttachPtrExpr)) {
+          FoundExistingMap = true;
+          break;
+        }
+
+        // Check if we have a map like this[0:1]
+        if (IsThisCapture) {
+          if (const auto *OASE = dyn_cast<ArraySectionExpr>(FirstExpr)) {
+            if (isa<CXXThisExpr>(OASE->getBase()->IgnoreParenImpCasts())) {
+              FoundExistingMap = true;
+              break;
+            }
+          }
+          continue;
+        }
+
+        // When the attach-ptr is something like `s.p`, check if
+        // `s` itself is mapped explicitly.
+        if (const auto *DRE = dyn_cast<DeclRefExpr>(FirstExpr)) {
+          if (DRE->getDecl() == CapturedVD) {
+            FoundExistingMap = true;
+            break;
+          }
+        }
+      }
+
+      if (FoundExistingMap)
+        continue;
+
+      // If no base map is found, we need to create an implicit map for the
+      // attach-pointer expr.
+
+      ComponentVectorStorage.emplace_back();
+      auto &AttachPtrComponents = ComponentVectorStorage.back();
+
+      static const OpenMPMapModifierKind Unknown = OMPC_MAP_MODIFIER_unknown;
+      bool SeenAttachPtrComponent = false;
+      // For creating a map on the attach-ptr `s.p/this->p`, we copy all
+      // components from the component-list which has `s.p/this->p`
+      // as the attach-ptr, starting from the component which matches
+      // `s.p/this->p`. This way, we'll have component-lists of
+      // `s.p` -> `s`, and `this->p` -> `this`.
+      for (size_t i = 0; i < ComponentsWithAttachPtr.size(); ++i) {
+        const auto &Component = ComponentsWithAttachPtr[i];
+        const Expr *ComponentExpr = Component.getAssociatedExpression();
+
+        if (!SeenAttachPtrComponent && ComponentExpr != AttachPtrExpr)
+          continue;
+        SeenAttachPtrComponent = true;
+
+        AttachPtrComponents.emplace_back(Component.getAssociatedExpression(),
+                                         Component.getAssociatedDeclaration(),
+                                         Component.isNonContiguous());
+      }
+      assert(!AttachPtrComponents.empty() &&
+             "Could not populate component-lists for mapping attach-ptr");
+
+      DeclComponentLists.emplace_back(
+          AttachPtrComponents, OMPC_MAP_tofrom, Unknown,
+          /*IsImplicit=*/true, /*mapper=*/nullptr, AttachPtrExpr);
+    }
+  }
+
   /// For a capture that has an associated clause, generate the base pointers,
   /// section pointers, sizes, map types, and mappers (all included in
   /// \a CurCaptureVarInfo).