Add an implicit map of attach-ptr on a target construct.

Author: abhinavgaba

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -9462,7 +9462,10 @@ class MappableExprsHandler {
   /// Populate component lists for non-lambda captured variables from map,
   /// is_device_ptr and has_device_addr clause info.
   void populateComponentListsForNonLambdaCaptureFromClauses(
-      const ValueDecl *VD, MapDataArrayTy &DeclComponentLists) const {
+      const ValueDecl *VD, MapDataArrayTy &DeclComponentLists,
+      SmallVectorImpl<
+          SmallVector<OMPClauseMappableExprCommon::MappableComponent, 8>>
+          &StorageForImplicitlyAddedComponentLists) const {
     if (VD && LambdasMap.count(VD))
       return;
 
@@ -9501,6 +9504,14 @@ class MappableExprsHandler {
         ++EI;
       }
     }
+
+    // For the target construct, if there's a map with a base-pointer that's
+    // a member of an implicitly captured struct, of the current class,
+    // we need to emit an implicit map on the pointer.
+    if (isOpenMPTargetExecutionDirective(CurExecDir->getDirectiveKind()))
+      addImplicitMapForAttachPtrBaseIfMemberOfCapturedVD(
+          VD, DeclComponentLists, StorageForImplicitlyAddedComponentLists);
+
     llvm::stable_sort(DeclComponentLists, [](const MapData &LHS,
                                              const MapData &RHS) {
       ArrayRef<OpenMPMapModifierKind> MapModifiers = std::get<2>(LHS);
@@ -9517,6 +9528,147 @@ 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 (IsThisCapture && !ME->isImplicitCXXThis() && !isa<CXXThisExpr>(Base))
+        continue;
+
+      if (!IsThisCapture && (!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 (AttachPtrExprComparator(this).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).
@@ -10325,8 +10477,12 @@ static void genMapInfoForCaptures(
 
       // Populate component lists for the captured variable from clauses.
       MappableExprsHandler::MapDataArrayTy DeclComponentLists;
+      SmallVector<
+          SmallVector<OMPClauseMappableExprCommon::MappableComponent, 8>, 4>
+          StorageForImplicitlyAddedComponentLists;
       MEHandler.populateComponentListsForNonLambdaCaptureFromClauses(
-          CapturedVD, DeclComponentLists);
+          CapturedVD, DeclComponentLists,
+          StorageForImplicitlyAddedComponentLists);
 
       // Map clauses on a target construct must either have a base pointer, or a
       // base-variable. So, if we don't have a base-pointer, that means that it

clang/test/OpenMP/target_map_codegen_20.cpp

@@ -74,8 +74,9 @@
 // CK21-NOUSE: [[MTYPE01:@.+]] = private {{.*}}constant [2 x i64] [i64 3, i64 16384]
 
 // CK21-LABEL: @.__omp_offloading_{{.*}}foo{{.*}}_l{{[0-9]+}}.region_id = weak constant i8 0
-// CK21: [[SIZE02:@.+]] = private {{.*}}constant [2 x i64] [i64 500, i64 {{4|8}}]
-// CK21-USE: [[MTYPE02:@.+]] = private unnamed_addr constant [2 x i64] [i64 34, i64 16384]
+// CK21-USE: [[SIZE02:@.+]] = private {{.*}}constant [3 x i64] [i64 {{4|8}}, i64 500, i64 {{4|8}}]
+// CK21-NOUSE: [[SIZE02:@.+]] = private {{.*}}constant [2 x i64] [i64 500, i64 {{4|8}}]
+// CK21-USE: [[MTYPE02:@.+]] = private {{.*}}constant [3 x i64] [i64 547, i64 2, i64 16384]
 // CK21-NOUSE: [[MTYPE02:@.+]] = private {{.*}}constant [2 x i64] [i64 2, i64 16384]
 
 // CK21-LABEL: @.__omp_offloading_{{.*}}foo{{.*}}_l{{[0-9]+}}.region_id = weak constant i8 0
@@ -168,7 +169,8 @@ struct CC {
 
 // Region 02
 
-//  &B[0], &B[X], 2 * sizeof(T), (FROM) / (FROM | PARAM)
+//  [&this[0], &this[0].B, sizeof(B), PARAM | ALLOC | IMPLICIT]
+//  &B[0], &B[X], 2 * sizeof(T), FROM
 //  &B, &B[X], sizeof(T*), ATTACH
 
 // CK21-DAG: call i32 @__tgt_target_kernel(ptr @{{.+}}, i64 -1, i32 -1, i32 0, ptr @.{{.+}}.region_id, ptr [[ARGS:%.+]])
@@ -181,18 +183,29 @@ struct CC {
 
 // CK21-DAG: [[BP0:%.+]] = getelementptr inbounds {{.+}}[[BP]], i{{.+}} 0, i{{.+}} 0
 // CK21-DAG: [[P0:%.+]] = getelementptr inbounds {{.+}}[[P]], i{{.+}} 0, i{{.+}} 0
-// CK21-DAG: store ptr [[VAR0:%.+]], ptr [[BP0]]
-// CK21-DAG: store ptr [[SEC0:%.+]], ptr [[P0]]
+// CK21-USE-DAG: store ptr [[THIS:%.+]], ptr [[BP0]]
+// CK21-USE-DAG: store ptr [[B:%.+]], ptr [[P0]]
+// CK21-USE-DAG: [[B:%B]] = getelementptr inbounds {{.*}}[[THIS]], i{{.*}} 0, i{{.*}} 2
+
+// CK21-DAG: [[BP1:%.+]] = getelementptr inbounds {{.+}}[[BP]], i{{.+}} 0, i{{.+}} 1
+// CK21-DAG: [[P1:%.+]] = getelementptr inbounds {{.+}}[[P]], i{{.+}} 0, i{{.+}} 1
+// CK21-USE-DAG: store ptr [[VAR0:%.+]], ptr [[BP1]]
+// CK21-USE-DAG: store ptr [[SEC0:%.+]], ptr [[P1]]
+// CK21-NOUSE-DAG: store ptr [[VAR0:%.+]], ptr [[BP0]]
+// CK21-NOUSE-DAG: store ptr [[SEC0:%.+]], ptr [[P0]]
+
 // CK21-DAG: [[VAR0]] = load ptr, ptr [[VAR00:%[^,]+]]
 // CK21-DAG: [[VAR00]] = getelementptr inbounds nuw %struct.CC, ptr [[THIS:%.+]], i{{.*}} 0, i{{.*}} 2
 // CK21-DAG: [[SEC0]] = getelementptr inbounds nuw float, ptr [[SEC00:%.+]], i{{.*}} 123
 // CK21-DAG: [[SEC00]] = load ptr, ptr [[SEC000:%[^,]+]]
 // CK21-DAG: [[SEC000]] = getelementptr inbounds nuw %struct.CC, ptr [[THIS]], i{{.*}} 0, i{{.*}} 2
 
-// CK21-DAG: [[BP1:%.+]] = getelementptr inbounds {{.+}}[[BP]], i{{.+}} 0, i{{.+}} 1
-// CK21-DAG: [[P1:%.+]] = getelementptr inbounds {{.+}}[[P]], i{{.+}} 0, i{{.+}} 1
-// CK21-DAG: store ptr [[VAR00]], ptr [[BP1]]
-// CK21-DAG: store ptr [[SEC0]], ptr [[P1]]
+// CK21-USE-DAG: [[BP2:%.+]] = getelementptr inbounds {{.+}}[[BP]], i{{.+}} 0, i{{.+}} 2
+// CK21-USE-DAG: [[P2:%.+]] = getelementptr inbounds {{.+}}[[P]], i{{.+}} 0, i{{.+}} 2
+// CK21-USE-DAG: store ptr [[VAR00:%.+]], ptr [[BP2]]
+// CK21-USE-DAG: store ptr [[SEC0:%.+]], ptr [[P2]]
+// CK21-NOUSE-DAG: store ptr [[VAR00]], ptr [[BP1]]
+// CK21-NOUSE-DAG: store ptr [[SEC0]], ptr [[P1]]
 
 // CK21-USE: call void [[CALL02:@.+]](ptr {{[^,]+}})
 // CK21-NOUSE: call void [[CALL02:@.+]]()