[OPENMP] Enable correct generation of runtime call when target directive is separated from teams directive by multiple curly brackets

[carlo.bertolli]

Summary

This patch fixes a bug in code generation of the correct OpenMP runtime library call in presence of target and teams, when target is separated by teams with multiple curly brackets. The current implementation will not be able to see the teams directive inside target and issue a call to tgt_target instead of the correct one tgt_target_teams.

[sfantao]

Thanks for the fix!

lib/CodeGen/CGOpenMPRuntime.cpp

4265

I think that this can be simplified to:

while (auto *S = dyn_cast<CompoundStmt>(TargetBody))
  TargetBody = S->body_front();

return dyn_cast<CompoundStmt>(OMPTeamsDirective);

I know that this is currently used only for teams, but I think it would be nice to make this a templated function to look for other possible nests. I suspect this will very useful for other cases like 'target teams parallel distribute'.

[carlo.bertolli]

lib/CodeGen/CGOpenMPRuntime.cpp

4265

Thanks! I agree that having this as a template will help with combined directives. I will produce a new version of the patch with that support.

About your new loop: this would crash if S->body_front() is nullptr. This may happen if you have something like this:

#target
{
}

A solution may be

while (auto *S = dyn_cast_or_null<CompoundStmt>(TargetBody))
  TargetBody = S->body_front();

return dyn_cast_or_null<CompoundStmt>(OMPTeamsDirective);

Thoughts?

[sfantao]

lib/CodeGen/CGOpenMPRuntime.cpp

4265

You're totally right, dyn_cast_or_null is the right cast to use in the return statement.

[carlo.bertolli]

Simplify the search function and generalize by using a template.

[ABataev]

lib/CodeGen/CGOpenMPRuntime.cpp

4265

You don't need SK, because body_front() member function exists only in CompoundStmt.

[carlo.bertolli]

Alexey

Thanks for the review. I addressed your comment in a new version of the revision.
Samuel already reviewed the patch, as you can see in the message exchange. I am very happy to apply any design change you can think of.

• Carlo

[carlo.bertolli]

[OPENMP] Remove template for intervening statement and only allow CompoundStmt.

[ABataev]

lib/CodeGen/CGOpenMPRuntime.cpp

4260–4273

I don't like this function at all. We need to add a template function. Each instantiation will produce a new function with almost the same body. MI prepfer to see another solution, like:

static const Stmt *ignoreCompoundStmts(const Stmt *Body) {
  while (auto *CS = dyn_cast_or_null<CompoundStmt>(Body))
    Body = CS->body_front();
  return Body;
}
...
if (auto *TeamsDir = dyn_cast<OMPTeamsDirective>(ignoreCompoundStmts(CS.getCapturedStmt()))) {
...

[carlo.bertolli]

Great, thanks for the hint. I do much prefer your version of the function. I am updating the diff.

[carlo.bertolli]

[OPENMP] Remove unnecessary templating of function.

[ABataev]

LG, except for the small comment

lib/CodeGen/CGOpenMPRuntime.cpp

4260

\brief tag is not required anymore, just a brief description is required.

[carlo.bertolli]

Committed revision 267972.

lib/CodeGen/CGOpenMPRuntime.cpp

OutlinedFnID = new llvm::GlobalVariable(
         llvm::GlobalValue::PrivateLinkage,
         llvm::Constant::getNullValue(CGM.Int8Ty), ".omp_offload.region_id");
 
   // Register the information for the entry associated with this target region.
   OffloadEntriesInfoManager.registerTargetRegionEntryInfo(
       DeviceID, FileID, ParentName, Line, OutlinedFn, OutlinedFnID);
 }
 
+/// \brief discard all CompoundStmts intervening between two constructs
+static const Stmt *ignoreCompoundStmts(const Stmt *Body) {
+  while (auto *CS = dyn_cast_or_null<CompoundStmt>(Body))
+    Body = CS->body_front();
+
+  return Body;
+}
+
 /// \brief Emit the num_teams clause of an enclosed teams directive at the
 /// target region scope. If there is no teams directive associated with the
 /// target directive, or if there is no num_teams clause associated with the
 /// enclosed teams directive, return nullptr.
 static llvm::Value *
 emitNumTeamsClauseForTargetDirective(CGOpenMPRuntime &OMPRuntime,
                                      CodeGenFunction &CGF,
                                      const OMPExecutableDirective &D) {
 
   assert(!CGF.getLangOpts().OpenMPIsDevice && "Clauses associated with the "
                                               "teams directive expected to be "
                                               "emitted only for the host!");
 
   // FIXME: For the moment we do not support combined directives with target and
   // teams, so we do not expect to get any num_teams clause in the provided
   // directive. Once we support that, this assertion can be replaced by the
   // actual emission of the clause expression.
   assert(D.getSingleClause<OMPNumTeamsClause>() == nullptr &&
          "Not expecting clause in directive.");
 
   // If the current target region has a teams region enclosed, we need to get
   // the number of teams to pass to the runtime function call. This is done
   // by generating the expression in a inlined region. This is required because
   // the expression is captured in the enclosing target environment when the
   // teams directive is not combined with target.
 
   const CapturedStmt &CS = *cast<CapturedStmt>(D.getAssociatedStmt());
 
   // FIXME: Accommodate other combined directives with teams when they become
   // available.
-  if (auto *TeamsDir = dyn_cast<OMPTeamsDirective>(CS.getCapturedStmt())) {
+  if (auto *TeamsDir = dyn_cast_or_null<OMPTeamsDirective>(
+          ignoreCompoundStmts(CS.getCapturedStmt()))) {
     if (auto *NTE = TeamsDir->getSingleClause<OMPNumTeamsClause>()) {
       CGOpenMPInnerExprInfo CGInfo(CGF, CS);
       CodeGenFunction::CGCapturedStmtRAII CapInfoRAII(CGF, &CGInfo);
       llvm::Value *NumTeams = CGF.EmitScalarExpr(NTE->getNumTeams());
       return CGF.Builder.CreateIntCast(NumTeams, CGF.Int32Ty,
                                        /*IsSigned=*/true);
     }
 
emitThreadLimitClauseForTargetDirective(CGOpenMPRuntime &OMPRuntime,
   // by generating the expression in a inlined region. This is required because
   // the expression is captured in the enclosing target environment when the
   // teams directive is not combined with target.
 
   const CapturedStmt &CS = *cast<CapturedStmt>(D.getAssociatedStmt());
 
   // FIXME: Accommodate other combined directives with teams when they become
   // available.
-  if (auto *TeamsDir = dyn_cast<OMPTeamsDirective>(CS.getCapturedStmt())) {
+  if (auto *TeamsDir = dyn_cast_or_null<OMPTeamsDirective>(
+          ignoreCompoundStmts(CS.getCapturedStmt()))) {
     if (auto *TLE = TeamsDir->getSingleClause<OMPThreadLimitClause>()) {
       CGOpenMPInnerExprInfo CGInfo(CGF, CS);
       CodeGenFunction::CGCapturedStmtRAII CapInfoRAII(CGF, &CGInfo);
       llvm::Value *ThreadLimit = CGF.EmitScalarExpr(TLE->getThreadLimit());
       return CGF.Builder.CreateIntCast(ThreadLimit, CGF.Int32Ty,
                                        /*IsSigned=*/true);
     }
 

test/OpenMP/teams_codegen.cpp

int teams_argument_global_local(int a){
   // CK1: call i32 @__tgt_target_teams(i32 -1, i8* @{{[^,]+}}, i32 1, i8** %{{[^,]+}}, i8** %{{[^,]+}}, i{{64|32}}* {{.+}}@{{[^,]+}}, i32 0, i32 0), i32* {{.+}}@{{[^,]+}}, i32 0, i32 0), i32 0, i32 0)
   // CK1: call void @{{.+}}(i{{64|32}} %{{.+}})
   #pragma omp target
   #pragma omp teams
   {
     ++comp;
   }
 
+  // CK1: call i32 @__tgt_target_teams(i32 -1, i8* @{{[^,]+}}, i32 1, i8** %{{[^,]+}}, i8** %{{[^,]+}}, i{{64|32}}* {{.+}}@{{[^,]+}}, i32 0, i32 0), i32* {{.+}}@{{[^,]+}}, i32 0, i32 0), i32 0, i32 0)
+  // CK1: call void @{{.+}}(i{{64|32}} %{{.+}})
+  #pragma omp target
+  {{{
+    #pragma omp teams
+    {
+      ++comp;
+    }
+  }}}
+  
   // CK1-DAG: call i32 @__tgt_target_teams(i32 -1, i8* @{{[^,]+}}, i32 2, i8** %{{[^,]+}}, i8** %{{[^,]+}}, i{{64|32}}* {{.+}}@{{[^,]+}}, i32 0, i32 0), i32* {{.+}}@{{[^,]+}}, i32 0, i32 0), i32 [[NT:%[^,]+]], i32 0)
   // CK1-DAG: [[NT]] = load i32, i32* [[NTA:%[^,]+]],
 
   // CK1: call void @{{.+}}(i{{64|32}} %{{.+}})
   #pragma omp target
   #pragma omp teams num_teams(la)
   {
     ++comp;