[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

8448–8466

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

8451–8454

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

8921–8935

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

8921–8935

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

8921–8935

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

8451–8454

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

private:
 
   /// Map between device pointer declarations and their expression components.
   /// The key value for declarations in 'this' is null.
   llvm::DenseMap<
       const ValueDecl *,
       SmallVector<OMPClauseMappableExprCommon::MappableExprComponentListRef, 4>>
       DevPointersMap;
 
+  /// Map between lambda declarations and their map type.
+  llvm::DenseMap<const ValueDecl *, const OMPMapClause *> LambdasMap;
+
   llvm::Value *getExprTypeSize(const Expr *E) const {
     QualType ExprTy = E->getType().getCanonicalType();
 
     // Calculate the size for array shaping expression.
     if (const auto *OAE = dyn_cast<OMPArrayShapingExpr>(E)) {
       llvm::Value *Size =
           CGF.getTypeSize(OAE->getBase()->getType()->getPointeeType());
       for (const Expr *SE : OAE->getDimensions()) {
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;
     }
+    auto I = LambdasMap.find(Cap.getCapturedVar()->getCanonicalDecl());
+    if (I != LambdasMap.end())
+      // for map(to: lambda): using user specified map type.
+      return getMapTypeBits(
+          I->getSecond()->getMapType(), I->getSecond()->getMapTypeModifiers(),
+          /*MotionModifiers=*/llvm::None, I->getSecond()->isImplicit(),
+          /*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());
   }
 
   static void setCorrectMemberOfFlag(OpenMPOffloadMappingFlags &Flags,
                                      OpenMPOffloadMappingFlags MemberOfFlag) {
     // If the entry is PTR_AND_OBJ but has not been marked with the special
     // placeholder value 0xFFFF in the MEMBER_OF field, then it should not be
     // marked as MEMBER_OF.
     if ((Flags & OMP_MAP_PTR_AND_OBJ) &&
         ((Flags & OMP_MAP_MEMBER_OF) != OMP_MAP_MEMBER_OF))
       return;
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));
+    // Extract map information.
+    for (const auto *C : Dir.getClausesOfKind<OMPMapClause>()) {
+      if (C->getMapType() != OMPC_MAP_to)
+        continue;
+      for (auto L : C->component_lists()) {
+        const ValueDecl *VD = std::get<0>(L);
+        const auto *RD = VD ? VD->getType()
+                                  .getCanonicalType()
+                                  .getNonReferenceType()
+                                  ->getAsCXXRecordDecl()
+                            : nullptr;
+        if (RD && RD->isLambda())
+          LambdasMap.try_emplace(std::get<0>(L), C);
+      }
+    }
   }
 
   /// 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
void generateInfoForCapture(const CapturedStmt::Capture *Cap,
     assert(!Cap->capturesVariableArrayType() &&
            "Not expecting to generate map info for a variable array type!");
 
     // We need to know when we generating information for the first component
     const ValueDecl *VD = Cap->capturesThis()
                               ? nullptr
                               : Cap->getCapturedVar()->getCanonicalDecl();
 
+    // for map(to: lambda): skip here, processing it in
+    // generateDefaultMapInfo
+    if (LambdasMap.count(VD))
+      return;
+
     // If this declaration appears in a is_device_ptr clause we just have to
     // pass the pointer by value. If it is a reference to a declaration, we just
     // pass its value.
     if (DevPointersMap.count(VD)) {
       CombinedInfo.Exprs.push_back(VD);
       CombinedInfo.BasePointers.emplace_back(Arg, VD);
       CombinedInfo.Pointers.push_back(Arg);
       CombinedInfo.Sizes.push_back(CGF.Builder.CreateIntCast(

clang/test/OpenMP/nvptx_lambda_pointer_capturing.cpp

 #ifndef HEADER
 #define HEADER
 
 // CHECK: [[ANON_T:%.+]] = type { i32*, i32* }
 // CHECK-DAG: [[SIZES_TEMPLATE:@.+]] = private {{.+}} constant [5 x i[[PTRSZ:32|64]]] [i{{32|64}} 4, i{{32|64}} 4, i{{32|64}} {{8|16}}, i{{32|64}} 0, i{{32|64}} 0]
 // CHECK-DAG: [[TYPES_TEMPLATE:@.+]] = private {{.+}} constant [5 x i64] [i64 800, i64 800, i64 673, i64 844424930132752, i64 844424930132752]
 // CHECK-DAG: [[SIZES:@.+]] = private {{.+}} constant [3 x i[[PTRSZ:32|64]]] [i{{32|64}} {{8|16}}, i{{32|64}} 0, i{{32|64}} 0]
 // CHECK-DAG: [[TYPES:@.+]] = private {{.+}} constant [3 x i64] [i64 673, i64 281474976711440, i64 281474976711440]
+// CHECK-DAG: [[TYPES3:@.+]] = private {{.+}} constant [3 x i64] [i64 545, i64 281474976711440, i64 800]
+// CHECK-DAG: [[TYPES11:@.+]] = private {{.+}} constant [5 x i64] [i64 800, i64 800, i64 549, i64 844424930132752, i64 844424930132752]
+// CHECK-DAG: [[TYPES13:@.+]] = private {{.+}} constant [2 x i64] [i64 545, i64 281474976711440]
+// CHECK-DAG: [[TYPES15:@.+]] = private {{.+}} constant [2 x i64] [i64 673, i64 281474976711440]
 
 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(always, 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);
+  #pragma omp target
+  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