[OPENMP] Fix assert of "Unable to find base lambda address" from adjustMemberOfForLambdaCaptures.

[jyu2]

Summary

The problem is happening when user passes lambda function with reference
type in the map clause.

The natural of the problem is when processing generateInfoForCapture, the
BasePointer is generated with new load for a lambda variable with reference
type. It is not expected in adjustMemberOfForLambdaCaptures.

One way to fix this is to skipping call to generateInfoForCapture for
map(to:lambda). The map info will be generated later in the call to
generateDefaultMapInfo, samiler as firsprivate clase.

This to fix https://bugs.llvm.org/show_bug.cgi?id=52071

[ABataev]

clang/lib/CodeGen/CGOpenMPRuntime.cpp

8445–8463

Maybe add a new data container for mapped lambdas, similar to FirstPrivateDecls to avoid all those extra checks here?

[jyu2]

Thanks Alexey's review. This patch is address Alexe's comments by adding new field LambdasMap.

[ABataev]

clang/lib/CodeGen/CGOpenMPRuntime.cpp

8448–8451

No need to create empty ArrayRefs, just pass llvm::None instead.

8928–8942

What if we have to(lambda) in data motion directive? What shall we do if there mapping modifiers? Not sure we shall ignore them.

[jyu2]

clang/lib/CodeGen/CGOpenMPRuntime.cpp

8928–8942

to(lambda) is not isOpenMPTargetExecutionDirective. I don't see generateInfoForLambdaCaptures and adjustMemberOfForLambdaCaptures functions get called in that part. Those two functions only called in emitTargetCall.
So I don't think we need that.
What do you think?
Thanks.
Jennifer

[ABataev]

clang/lib/CodeGen/CGOpenMPRuntime.cpp

8928–8942

Ok, let's skip it for now but still need to keep in mind to fix/implement it in the future.

[jyu2]

Address Alexey's comments.

[ABataev]

clang/lib/CodeGen/CGOpenMPRuntime.cpp

8448–8451

What if there are modifiers in the clause?

[jyu2]

Thanks Alexey! Fix missing Modifiers.

[ABataev]

Tests with modifiers?

[jyu2]

Adding check for map type in test. Thanks Alexey.

[ABataev]

LG

clang/lib/CodeGen/CGOpenMPRuntime.cpp

getMapModifiersForPrivateClauses(const CapturedStmt::Capture &Cap) const {
     // declaration is known as first-private in this handler.
     if (FirstPrivateDecls.count(Cap.getCapturedVar())) {
       if (Cap.getCapturedVar()->getType()->isAnyPointerType())
         return MappableExprsHandler::OMP_MAP_TO |
                MappableExprsHandler::OMP_MAP_PTR_AND_OBJ;
       return MappableExprsHandler::OMP_MAP_PRIVATE |
              MappableExprsHandler::OMP_MAP_TO;
     }
+    const ValueDecl *VD = Cap.getCapturedVar()->getCanonicalDecl();
+    const auto *RD = VD->getType()
+                         .getCanonicalType()
+                         .getNonReferenceType()
+                         ->getAsCXXRecordDecl();
+    const auto *CurExecDir = CurDir.get<const OMPExecutableDirective *>();
+    for (const auto *C : CurExecDir->getClausesOfKind<OMPMapClause>())
+      if (!C->decl_component_lists(VD).empty())
+        if (RD && RD->isLambda())
+          if (C->getMapType() == OMPC_MAP_to) {
+            // for map(to: lambda): using user specified map type.
+            ArrayRef<OpenMPMapModifierKind> MapModifiers;
+            ArrayRef<OpenMPMotionModifierKind> MotionModifiers;
+            return getMapTypeBits(C->getMapType(), MapModifiers,
+                                  MotionModifiers, false,
+                                  /*AddPtrFlag=*/false,
+                                  /*AddIsTargetParamFlag=*/false,
+                                  /*isNonContiguous=*/false);
+          }
     return MappableExprsHandler::OMP_MAP_TO |
            MappableExprsHandler::OMP_MAP_FROM;
   }
 
   static OpenMPOffloadMappingFlags getMemberOfFlag(unsigned Position) {
     // Rotate by getFlagMemberOffset() bits.
     return static_cast<OpenMPOffloadMappingFlags>(((uint64_t)Position + 1)
                                                   << getFlagMemberOffset());
for (const auto *C : Dir.getClausesOfKind<OMPUsesAllocatorsClause>()) {
                          ->getDecl()))
           FirstPrivateDecls.try_emplace(VD, /*Implicit=*/true);
       }
     }
     // Extract device pointer clause information.
     for (const auto *C : Dir.getClausesOfKind<OMPIsDevicePtrClause>())
       for (auto L : C->component_lists())
         DevPointersMap[std::get<0>(L)].push_back(std::get<1>(L));
   }
 
   /// Constructor for the declare mapper directive.
   MappableExprsHandler(const OMPDeclareMapperDecl &Dir, CodeGenFunction &CGF)
       : CurDir(&Dir), CGF(CGF) {}
 
   /// Generate code for the combined entry if we have a partially mapped struct
   /// and take care of the mapping flags of the arguments corresponding to
   /// individual struct members.
   void emitCombinedEntry(MapCombinedInfoTy &CombinedInfo,
                          MapFlagsArrayTy &CurTypes,
                          const StructRangeInfoTy &PartialStruct,
                          const ValueDecl *VD = nullptr,
                          bool NotTargetParams = true) const {
     if (CurTypes.size() == 1 &&
         ((CurTypes.back() & OMP_MAP_MEMBER_OF) != OMP_MAP_MEMBER_OF) &&
         !PartialStruct.IsArraySection)
       return;
     Address LBAddr = PartialStruct.LowestElem.second;
     Address HBAddr = PartialStruct.HighestElem.second;
     if (PartialStruct.HasCompleteRecord) {
       LBAddr = PartialStruct.LB;
       HBAddr = PartialStruct.LB;
void generateInfoForCapture(const CapturedStmt::Capture *Cap,
     using MapData =
         std::tuple<OMPClauseMappableExprCommon::MappableExprComponentListRef,
                    OpenMPMapClauseKind, ArrayRef<OpenMPMapModifierKind>, bool,
                    const ValueDecl *, const Expr *>;
     SmallVector<MapData, 4> DeclComponentLists;
     assert(CurDir.is<const OMPExecutableDirective *>() &&
            "Expect a executable directive");
     const auto *CurExecDir = CurDir.get<const OMPExecutableDirective *>();
+    const auto *RD = VD ? VD->getType()
+                              .getCanonicalType()
+                              .getNonReferenceType()
+                              ->getAsCXXRecordDecl()
+                        : nullptr;
     for (const auto *C : CurExecDir->getClausesOfKind<OMPMapClause>()) {
       const auto *EI = C->getVarRefs().begin();
       for (const auto L : C->decl_component_lists(VD)) {
+        if (RD && RD->isLambda())
+          // for map(to: lambda): skip here, processing it in
+          // generateDefaultMapInfo
+          if (C->getMapType() == OMPC_MAP_to)
+            return;
         const ValueDecl *VDecl, *Mapper;
         // The Expression is not correct if the mapping is implicit
         const Expr *E = (C->getMapLoc().isValid()) ? *EI : nullptr;
         OMPClauseMappableExprCommon::MappableExprComponentListRef Components;
         std::tie(VDecl, Components, Mapper) = L;
         assert(VDecl == VD && "We got information for the wrong declaration??");
         assert(!Components.empty() &&
                "Not expecting declaration with no component lists.");

clang/test/OpenMP/nvptx_lambda_pointer_capturing.cpp

 template <typename F>
 void omp_loop(int start, int end, F body) {
 #pragma omp target teams distribute parallel for
   for (int i = start; i < end; ++i) {
     body(i);
   }
 }
 
+template <typename F>
+void omp_loop_ref(int start, int end, F& body) {
+#pragma omp target teams distribute parallel for map(to: body)
+  for (int i = start; i < end; ++i) {
+    body(i);
+  }
+  int *p;
+  const auto &body_ref = [=](int i) {p[i]=0;};
+  #pragma omp target map(to: body_ref)
+  body_ref(10);
+}
+template <class FTy>
+struct C {
+  static void xoo(const FTy& f) {
+    int x = 10;
+    #pragma omp target map(to:f)
+      f(x);
+  }
+};
+
+template <class FTy>
+void zoo(const FTy &functor) {
+  C<FTy>::xoo(functor);
+}
+
+
 // CHECK: define {{.*}}[[MAIN:@.+]](
 int main()
 {
   int* p = new int[100];
   int* q = new int[100];
   auto body = [=](int i){
     p[i] = q[i];
   };
+  zoo([=](int i){p[i] = 0;});
 
 #pragma omp target teams distribute parallel for
   for (int i = 0; i < 100; ++i) {
     body(i);
   }
 
 // CHECK: [[BASE_PTRS:%.+]] = alloca [3 x i8*]{{.+}}
 // CHECK: [[PTRS:%.+]] = alloca [3 x i8*]{{.+}}

 
 // actual target invocation
 // CHECK: [[BASES_GEP:%.+]] = getelementptr {{.+}} [3 x {{.+}}*], [3 x {{.+}}*]* [[BASE_PTRS]], {{.+}} 0, {{.+}} 0
 // CHECK: [[PTRS_GEP:%.+]] = getelementptr {{.+}} [3 x {{.+}}*], [3 x {{.+}}*]* [[PTRS]], {{.+}} 0, {{.+}} 0
 // CHECK: {{%.+}} = call{{.+}} @__tgt_target_teams_mapper(%struct.ident_t* @{{.+}}, {{.+}}, {{.+}}, {{.+}}, i8** [[BASES_GEP]], i8** [[PTRS_GEP]], i[[PTRSZ]]* getelementptr inbounds ([3 x i{{.+}}], [3 x i{{.+}}]* [[SIZES]], i{{.+}} 0, i{{.+}} 0), i64* getelementptr inbounds ([3 x i64], [3 x i64]* [[TYPES]], i{{.+}} 0, i{{.+}} 0), i8** null, i8** null, {{.+}}, {{.+}})
 
 
   omp_loop(0,100,body);
+  omp_loop_ref(0,100,body);
 }
 
 // CHECK: [[BASE_PTRS:%.+]] = alloca [5 x i8*]{{.+}}
 // CHECK: [[PTRS:%.+]] = alloca [5 x i8*]{{.+}}
 
 // access of pointers inside lambdas
 // CHECK: [[BASE_PTR1:%.+]] = getelementptr inbounds [[ANON_T]], [[ANON_T]]* %{{.+}}, i{{.+}} 0, i{{.+}} 0
 // CHECK: [[PTR1:%.+]] = load i32*, i32** [[BASE_PTR1]]

 // CHECK: store i32* [[PTR2]], i32** [[CAST_PTR2]]{{.+}}
 
 
 // actual target invocation
 // CHECK: [[BASES_GEP:%.+]] = getelementptr {{.+}} [5 x {{.+}}*], [5 x {{.+}}*]* [[BASE_PTRS]], {{.+}} 0, {{.+}} 0
 // CHECK: [[PTRS_GEP:%.+]] = getelementptr {{.+}} [5 x {{.+}}*], [5 x {{.+}}*]* [[PTRS]], {{.+}} 0, {{.+}} 0
 // CHECK: {{%.+}} = call{{.+}} @__tgt_target_teams_mapper(%struct.ident_t* @{{.+}}, {{.+}}, {{.+}}, {{.+}}, i8** [[BASES_GEP]], i8** [[PTRS_GEP]], i[[PTRSZ]]* getelementptr inbounds ([5 x i{{.+}}], [5 x i{{.+}}]* [[SIZES_TEMPLATE]], i{{.+}} 0, i{{.+}} 0), i64* getelementptr inbounds ([5 x i64], [5 x i64]* [[TYPES_TEMPLATE]], i{{.+}} 0, i{{.+}} 0), i8** null, i8** null, {{.+}}, {{.+}})
 
+// CHECK: define internal void @{{.+}}omp_loop_ref{{.+}}(
+// CHECK: [[BODY:%body.addr]] = alloca %class.anon*
+// CHECK: [[TMP:%tmp]] = alloca %class.anon*
+// CHECK: [[BODY_REF:%body_ref]] = alloca %class.anon.1*
+// CHECK: [[REF_TMP:%ref.tmp]] = alloca %class.anon.1
+// CHECK: [[TMP8:%tmp.+]] = alloca %class.anon.1*
+// CHECK: [[L0:%.+]] = load %class.anon*, %class.anon** [[BODY]]
+// CHECK: store %class.anon* [[L0]], %class.anon** [[TMP]]
+// CHECK: [[L5:%.+]] = load %class.anon*, %class.anon** [[TMP]]
+// CHECK-NOT [[L6:%.+]] = load %class.anon*, %class.anon** [[TMP]]
+// CHECK-NOT [[L7:%.+]] = load %class.anon*, %class.anon** [[TMP]]
+// CHECK: store %class.anon.1* [[REF_TMP]], %class.anon.1** [[BODY_REF]]
+// CHECK:[[L47:%.+]] =  load %class.anon.1*, %class.anon.1** [[BODY_REF]]
+// CHECK: store %class.anon.1* [[L47]], %class.anon.1** [[TMP8]]
+// CHECK: [[L48:%.+]] = load %class.anon.1*, %class.anon.1** [[TMP8]]
+// CHECK-NOT: [[L49:%.+]] = load %class.anon.1*, %class.anon.1** [[TMP8]]
+// CHECK-NOT: [[L50:%.+]] = load %class.anon.1*, %class.anon.1** [[TMP8]]
+// CHECK: ret void
+
+// CHECK: define internal void @{{.+}}xoo{{.+}}(
+// CHECK: [[FADDR:%f.addr]] = alloca %class.anon.0*
+// CHECK: [[L0:%.+]] = load %class.anon.0*, %class.anon.0** [[FADDR]]
+// CHECK: store %class.anon.0* [[L0]], %class.anon.0** [[TMP:%tmp]]
+// CHECK: [[L1:%.+]] = load %class.anon.0*, %class.anon.0** [[TMP]]
+// CHECK-NOT: %4 = load %class.anon.0*, %class.anon.0** [[TMP]]
+// CHECK-NOT: %5 = load %class.anon.0*, %class.anon.0** [[TMP]]
+// CHECK: [[L4:%.+]] = getelementptr inbounds %class.anon.0, %class.anon.0* [[L1]], i32 0, i32 0
+// CHECK: [[L5:%.+]] = load i{{.*}}*, i{{.*}}** [[L4]]
+// CHECK: ret void
+
 #endif