[OPENMP] Codegen for teams directive for NVPTX

[carlo.bertolli]

Summary

This patch implements the teams directive for the NVPTX backend. It is different from the host code generation path as it:

• Does not call kmpc_fork_teams. All necessary teams and threads are started upon touching the target region, when launching a CUDA kernel, and their execution is coordinated through sequential and parallel regions within the target region.

• Does not call kmpc_push_num_teams even if a num_teams of thread_limit clause is present. Setting the number of teams and the thread limit is implemented by the nvptx-related runtime.

Please note that I am now passing a Clang Expr * to emitPushNumTeams instead of the originally chosen llvm::Value * type. The reason for that is that I want to avoid emitting expressions for num_teams and thread_limit if they are not needed in the target region.

[ABataev]

lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

375–376

This will cause a crash for captured global variables. Emit as inlined directive

[carlo.bertolli]

Addressed comment in new version of diff.

[carlo.bertolli]

[OPENMP] This new version of the patch uses the inlining machinery of CGOpenMPRuntime.cpp instead of just dumping the teams body statement using EmitStmt.

[ABataev]

Add tests with captured globals to check that this problem is resolved

[carlo.bertolli]

Support for global variables in a target region requires the global to be placed in a pragma declare target region. Pragma declare target region is not currently available in Clang (no parsing, sema, or codegen). I will add a check for this in the regression test once the support becomes available. Thanks!

[carlo.bertolli]

I can actually add a regression test that checks that the compiler does not break when using a global in a teams region, even without declare target, if that is what is required here.

[ABataev]

LG

lib/CodeGen/CGOpenMPRuntimeNVPTX.h

145–172

Remove 'virtual' and add 'override' at the end of each function

[carlo.bertolli]

[OPENMP] Even though this patch was already accepted in its previous form, comments on depending patch D18286 (http://reviews.llvm.org/D18286) revealed that a new approach for this patch was necessary. Instead of committing something that I have to change anyway later on, I decided to provide a new version of this base patch. Please review it again and let me know about any comments you may have.

[ABataev]

LG

[carlo.bertolli]

Committed revision 265304.

lib/CodeGen/CGOpenMPRuntime.h

public:
   virtual void emitTeamsCall(CodeGenFunction &CGF,
                              const OMPExecutableDirective &D,
                              SourceLocation Loc, llvm::Value *OutlinedFn,
                              ArrayRef<llvm::Value *> CapturedVars);
 
   /// \brief Emits call to void __kmpc_push_num_teams(ident_t *loc, kmp_int32
   /// global_tid, kmp_int32 num_teams, kmp_int32 thread_limit) to generate code
   /// for num_teams clause.
-  /// \param NumTeams An integer value of teams.
-  /// \param ThreadLimit An integer value of threads.
-  virtual void emitNumTeamsClause(CodeGenFunction &CGF, llvm::Value *NumTeams,
-                                  llvm::Value *ThreadLimit, SourceLocation Loc);
-
+  /// \param NumTeams An integer expression of teams.
+  /// \param ThreadLimit An integer expression of threads.
+  virtual void emitNumTeamsClause(CodeGenFunction &CGF, const Expr *NumTeams,
+                                  const Expr *ThreadLimit, SourceLocation Loc);
 };
 
 } // namespace CodeGen
 } // namespace clang
 
 #endif

lib/CodeGen/CGOpenMPRuntime.cpp

void CGOpenMPRuntime::emitTeamsCall(CodeGenFunction &CGF,
   RealArgs.append(std::begin(Args), std::end(Args));
   RealArgs.append(CapturedVars.begin(), CapturedVars.end());
 
   auto RTLFn = createRuntimeFunction(OMPRTL__kmpc_fork_teams);
   CGF.EmitRuntimeCall(RTLFn, RealArgs);
 }
 
 void CGOpenMPRuntime::emitNumTeamsClause(CodeGenFunction &CGF,
-                                         llvm::Value *NumTeams,
-                                         llvm::Value *ThreadLimit,
+                                         const Expr *NumTeams,
+                                         const Expr *ThreadLimit,
                                          SourceLocation Loc) {
   if (!CGF.HaveInsertPoint())
     return;
 
   auto *RTLoc = emitUpdateLocation(CGF, Loc);
 
+  llvm::Value *NumTeamsVal =
+      (NumTeams)
+          ? CGF.Builder.CreateIntCast(CGF.EmitScalarExpr(NumTeams),
+                                      CGF.CGM.Int32Ty, /* isSigned = */ true)
+          : CGF.Builder.getInt32(0);
+
+  llvm::Value *ThreadLimitVal =
+      (ThreadLimit)
+          ? CGF.Builder.CreateIntCast(CGF.EmitScalarExpr(ThreadLimit),
+                                      CGF.CGM.Int32Ty, /* isSigned = */ true)
+          : CGF.Builder.getInt32(0);
+
   // Build call __kmpc_push_num_teamss(&loc, global_tid, num_teams, thread_limit)
-  llvm::Value *PushNumTeamsArgs[] = {
-      RTLoc, getThreadID(CGF, Loc), NumTeams, ThreadLimit};
+  llvm::Value *PushNumTeamsArgs[] = {RTLoc, getThreadID(CGF, Loc), NumTeamsVal,
+                                     ThreadLimitVal};
   CGF.EmitRuntimeCall(createRuntimeFunction(OMPRTL__kmpc_push_num_teams),
                       PushNumTeamsArgs);
 }

lib/CodeGen/CGOpenMPRuntimeNVPTX.h

void emitTargetOutlinedFunction(const OMPExecutableDirective &D,
                                   StringRef ParentName,
                                   llvm::Function *&OutlinedFn,
                                   llvm::Constant *&OutlinedFnID,
                                   bool IsOffloadEntry,
                                   const RegionCodeGenTy &CodeGen) override;
 
 public:
   explicit CGOpenMPRuntimeNVPTX(CodeGenModule &CGM);
+
+  /// \brief This function ought to emit, in the general case, a call to
+  // the openmp runtime kmpc_push_num_teams. In NVPTX backend it is not needed
+  // as these numbers are obtained through the PTX grid and block configuration.
+  /// \param NumTeams An integer expression of teams.
+  /// \param ThreadLimit An integer expression of threads.
+  void emitNumTeamsClause(CodeGenFunction &CGF, const Expr *NumTeams,
+                          const Expr *ThreadLimit, SourceLocation Loc) override;
+
+  /// \brief Emits inlined function for the specified OpenMP parallel
+  //  directive but an inlined function for teams.
+  /// \a D. This outlined function has type void(*)(kmp_int32 *ThreadID,
+  /// kmp_int32 BoundID, struct context_vars*).
+  /// \param D OpenMP directive.
+  /// \param ThreadIDVar Variable for thread id in the current OpenMP region.
+  /// \param InnermostKind Kind of innermost directive (for simple directives it
+  /// is a directive itself, for combined - its innermost directive).
+  /// \param CodeGen Code generation sequence for the \a D directive.
+  llvm::Value *
+  emitParallelOrTeamsOutlinedFunction(const OMPExecutableDirective &D,
+                                      const VarDecl *ThreadIDVar,
+                                      OpenMPDirectiveKind InnermostKind,
+                                      const RegionCodeGenTy &CodeGen) override;
+
+  /// \brief Emits code for teams call of the \a OutlinedFn with
+  /// variables captured in a record which address is stored in \a
+  /// CapturedStruct.
+  /// \param OutlinedFn Outlined function to be run by team masters. Type of
+  /// this function is void(*)(kmp_int32 *, kmp_int32, struct context_vars*).
+  /// \param CapturedVars A pointer to the record with the references to
+  /// variables used in \a OutlinedFn function.
+  ///
+  void emitTeamsCall(CodeGenFunction &CGF, const OMPExecutableDirective &D,
+                     SourceLocation Loc, llvm::Value *OutlinedFn,
+                     ArrayRef<llvm::Value *> CapturedVars) override;
 };
 
 } // CodeGen namespace.
 } // clang namespace.
 
 #endif // LLVM_CLANG_LIB_CODEGEN_CGOPENMPRUNTIMENVPTX_H

lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

 //===---- CGOpenMPRuntimeNVPTX.cpp - Interface to OpenMP NVPTX Runtimes ---===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This provides a class for OpenMP runtime code generation specialized to NVPTX
 // targets.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CGOpenMPRuntimeNVPTX.h"
 #include "clang/AST/DeclOpenMP.h"
+#include "CodeGenFunction.h"
+#include "clang/AST/StmtOpenMP.h"
 
 using namespace clang;
 using namespace CodeGen;
 
 /// \brief Get the GPU warp size.
 llvm::Value *CGOpenMPRuntimeNVPTX::getNVPTXWarpSize(CodeGenFunction &CGF) {
   CGBuilderTy &Bld = CGF.Builder;
   return Bld.CreateCall(

 CGOpenMPRuntimeNVPTX::CGOpenMPRuntimeNVPTX(CodeGenModule &CGM)
     : CGOpenMPRuntime(CGM), ActiveWorkers(nullptr), WorkID(nullptr) {
   if (!CGM.getLangOpts().OpenMPIsDevice)
     llvm_unreachable("OpenMP NVPTX can only handle device code.");
 
   // Called once per module during initialization.
   initializeEnvironment();
 }
+
+void CGOpenMPRuntimeNVPTX::emitNumTeamsClause(CodeGenFunction &CGF,
+                                              const Expr *NumTeams,
+                                              const Expr *ThreadLimit,
+                                              SourceLocation Loc) {}
+
+llvm::Value *CGOpenMPRuntimeNVPTX::emitParallelOrTeamsOutlinedFunction(
+    const OMPExecutableDirective &D, const VarDecl *ThreadIDVar,
+    OpenMPDirectiveKind InnermostKind, const RegionCodeGenTy &CodeGen) {
+
+  llvm::Function *OutlinedFun = nullptr;
+  if (isa<OMPTeamsDirective>(D)) {
+    llvm::Value *OutlinedFunVal =
+        CGOpenMPRuntime::emitParallelOrTeamsOutlinedFunction(
+            D, ThreadIDVar, InnermostKind, CodeGen);
+    OutlinedFun = cast<llvm::Function>(OutlinedFunVal);
+    OutlinedFun->addFnAttr(llvm::Attribute::AlwaysInline);
+  } else
+    llvm_unreachable("parallel directive is not yet supported for nvptx "
+                     "backend.");
+
+  return OutlinedFun;
+}
+
+void CGOpenMPRuntimeNVPTX::emitTeamsCall(CodeGenFunction &CGF,
+                                         const OMPExecutableDirective &D,
+                                         SourceLocation Loc,
+                                         llvm::Value *OutlinedFn,
+                                         ArrayRef<llvm::Value *> CapturedVars) {
+  if (!CGF.HaveInsertPoint())
+    return;
+
+  Address ZeroAddr =
+      CGF.CreateTempAlloca(CGF.Int32Ty, CharUnits::fromQuantity(4),
+                           /*Name*/ ".zero.addr");
+  CGF.InitTempAlloca(ZeroAddr, CGF.Builder.getInt32(/*C*/ 0));
+  llvm::SmallVector<llvm::Value *, 16> OutlinedFnArgs;
+  OutlinedFnArgs.push_back(ZeroAddr.getPointer());
+  OutlinedFnArgs.push_back(ZeroAddr.getPointer());
+  OutlinedFnArgs.append(CapturedVars.begin(), CapturedVars.end());
+  CGF.EmitCallOrInvoke(OutlinedFn, OutlinedFnArgs);
+}

lib/CodeGen/CGStmtOpenMP.cpp

static void emitCommonOMPTeamsDirective(CodeGenFunction &CGF,
   auto OutlinedFn = CGF.CGM.getOpenMPRuntime().
       emitParallelOrTeamsOutlinedFunction(S,
           *CS->getCapturedDecl()->param_begin(), InnermostKind, CodeGen);
 
   const OMPTeamsDirective &TD = *dyn_cast<OMPTeamsDirective>(&S);
   const OMPNumTeamsClause *NT = TD.getSingleClause<OMPNumTeamsClause>();
   const OMPThreadLimitClause *TL = TD.getSingleClause<OMPThreadLimitClause>();
   if (NT || TL) {
-    llvm::Value *NumTeamsVal = (NT) ? CGF.Builder.CreateIntCast(
-        CGF.EmitScalarExpr(NT->getNumTeams()), CGF.CGM.Int32Ty,
-        /* isSigned = */ true) :
-        CGF.Builder.getInt32(0);
-
-    llvm::Value *ThreadLimitVal = (TL) ? CGF.Builder.CreateIntCast(
-        CGF.EmitScalarExpr(TL->getThreadLimit()), CGF.CGM.Int32Ty,
-        /* isSigned = */ true) :
-        CGF.Builder.getInt32(0);
+    Expr *NumTeams = (NT) ? NT->getNumTeams() : nullptr;
+    Expr *ThreadLimit = (TL) ? TL->getThreadLimit() : nullptr;
 
-    CGF.CGM.getOpenMPRuntime().emitNumTeamsClause(CGF, NumTeamsVal,
-        ThreadLimitVal, S.getLocStart());
+    CGF.CGM.getOpenMPRuntime().emitNumTeamsClause(CGF, NumTeams, ThreadLimit,
+                                                  S.getLocStart());
   }
 
   OMPLexicalScope Scope(CGF, S);
   llvm::SmallVector<llvm::Value *, 16> CapturedVars;
   CGF.GenerateOpenMPCapturedVars(*CS, CapturedVars);
   CGF.CGM.getOpenMPRuntime().emitTeamsCall(CGF, S, S.getLocStart(), OutlinedFn,
                                            CapturedVars);
 }

test/OpenMP/nvptx_teams_codegen.cpp

+// Test target codegen - host bc file has to be created first.
+// RUN: %clang_cc1 -DCK1 -verify -fopenmp -x c++ -triple powerpc64le-unknown-unknown -fomptargets=nvptx64-nvidia-cuda -emit-llvm-bc %s -o %t-ppc-host.bc
+// RUN: %clang_cc1 -DCK1 -verify -fopenmp -x c++ -triple nvptx64-unknown-unknown -fomptargets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fomp-host-ir-file-path %t-ppc-host.bc -o - | FileCheck %s --check-prefix CK1 --check-prefix CK1-64
+// RUN: %clang_cc1 -DCK1 -verify -fopenmp -x c++ -triple i386-unknown-unknown -fomptargets=nvptx-nvidia-cuda -emit-llvm-bc %s -o %t-x86-host.bc
+// RUN: %clang_cc1 -DCK1 -verify -fopenmp -x c++ -triple nvptx-unknown-unknown -fomptargets=nvptx-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fomp-host-ir-file-path %t-x86-host.bc -o - | FileCheck %s --check-prefix CK1 --check-prefix CK1-32
+// expected-no-diagnostics
+#ifndef HEADER
+#define HEADER
+
+#ifdef CK1
+
+template <typename T>
+int tmain(T argc) {
+#pragma omp target
+#pragma omp teams
+  argc = 0;
+  return 0;
+}
+
+
+int main (int argc, char **argv) {
+#pragma omp target
+#pragma omp teams
+  {
+  argc = 0;
+  }
+  return tmain(argv);
+}
+
+// only nvptx side: do not outline teams region and do not call fork_teams
+// CK1:  define {{.*}}void @{{[^,]+}}(i{{[0-9]+}} [[ARGC:%.+]])
+// CK1:  {{.+}} = alloca i{{[0-9]+}}*,
+// CK1:  {{.+}} = alloca i{{[0-9]+}}*,
+// CK1:  [[ARGCADDR_PTR:%.+]] = alloca i{{[0-9]+}}*,
+// CK1:  [[ARGCADDR:%.+]] = alloca i{{[0-9]+}},
+// CK1:  store {{.+}} 0, {{.+}},
+// CK1:  store i{{[0-9]+}} [[ARGC]], i{{[0-9]+}}* [[ARGCADDR]],
+// CK1-64:  [[CONV:%.+]] = bitcast i{{[0-9]+}}* [[ARGCADDR]] to i{{[0-9]+}}*
+// CK1-64:  store i{{[0-9]+}}* [[CONV]], i{{[0-9]+}}** [[ARGCADDR_PTR]],
+// CK1-32:  store i{{[0-9]+}}* [[ARGCADDR]], i{{[0-9]+}}** [[ARGCADDR_PTR]],
+// CK1:  [[ARGCADDR_PTR_REF:%.+]] = load i{{[0-9]+}}*, i{{[0-9]+}}** [[ARGCADDR_PTR]],
+// CK1:  store i{{[0-9]+}} 0, i{{[0-9]+}}* [[ARGCADDR_PTR_REF]],
+// CK1-NOT: call {{.*}}void (%ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_teams(
+// CK1:  ret void
+// CK1-NEXT: }
+
+// target region in template
+// CK1: define {{.*}}void @{{[^,]+}}(i{{.+}}***{{.+}} [[ARGC:%.+]])
+// CK1: [[ARGCADDR_PTR:%.+]] = alloca i{{.+}}***,
+// CK1: [[ARGCADDR:%.+]] = alloca i{{.+}}***,
+// CK1: store i{{.+}}*** [[ARGC]], i{{.+}}**** [[ARGCADDR]]
+// CK1: [[ARGCADDR_REF:%.+]] = load i{{.+}}***, i{{.+}}**** [[ARGCADDR]],
+// CK1: store i8*** [[ARGCADDR_REF]], i8**** [[ARGCADDR_PTR]],
+// CK1: [[ARGCADDR_PTR_REF:%.+]] = load i{{.+}}***, i{{.+}}**** [[ARGCADDR_PTR]],
+// CK1: store i{{[0-9]+}}** null, i{{[0-9]+}}*** [[ARGCADDR_PTR_REF]],
+// CK1-NOT: call {{.*}}void (%ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_teams(
+// CK1:  ret void
+// CK1-NEXT: }
+
+
+#endif // CK1
+
+// Test target codegen - host bc file has to be created first.
+// RUN: %clang_cc1 -DCK2 -verify -fopenmp -x c++ -triple powerpc64le-unknown-unknown -fomptargets=nvptx64-nvidia-cuda -emit-llvm-bc %s -o %t-ppc-host.bc
+// RUN: %clang_cc1 -DCK2 -verify -fopenmp -x c++ -triple nvptx64-unknown-unknown -fomptargets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fomp-host-ir-file-path %t-ppc-host.bc -o - | FileCheck %s --check-prefix CK2 --check-prefix CK2-64
+// RUN: %clang_cc1 -DCK2 -verify -fopenmp -x c++ -triple i386-unknown-unknown -fomptargets=nvptx-nvidia-cuda -emit-llvm-bc %s -o %t-x86-host.bc
+// RUN: %clang_cc1 -DCK2 -verify -fopenmp -x c++ -triple nvptx-unknown-unknown -fomptargets=nvptx-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fomp-host-ir-file-path %t-x86-host.bc -o - | FileCheck %s --check-prefix CK2 --check-prefix CK2-32
+// expected-no-diagnostics
+#ifdef CK2
+
+template <typename T>
+int tmain(T argc) {
+  int a = 10;
+  int b = 5;
+#pragma omp target
+#pragma omp teams num_teams(a) thread_limit(b)
+  {
+  argc = 0;
+  }
+  return 0;
+}
+
+int main (int argc, char **argv) {
+  int a = 20;
+  int b = 5;
+#pragma omp target
+#pragma omp teams num_teams(a) thread_limit(b)
+  {
+  argc = 0;
+  }
+  return tmain(argv);
+}
+
+// CK2: define {{.*}}void @{{[^,]+}}(i{{[0-9]+}} [[A_IN:%.+]], i{{[0-9]+}} [[B_IN:%.+]], i{{[0-9]+}} [[ARGC_IN:.+]])
+// CK2: {{.}} = alloca i{{[0-9]+}}*,
+// CK2: {{.}} = alloca i{{[0-9]+}}*,
+// CK2: [[ARGCADDR_PTR:%.+]] = alloca i{{[0-9]+}}*,
+// CK2: [[AADDR:%.+]] = alloca i{{[0-9]+}},
+// CK2: [[BADDR:%.+]] = alloca i{{[0-9]+}},
+// CK2: [[ARGCADDR:%.+]] = alloca i{{[0-9]+}},
+// CK2-NOT:  {{%.+}} = call i32 @__kmpc_global_thread_num(
+// CK2: store i{{[0-9]+}} [[A_IN]], i{{[0-9]+}}* [[AADDR]],
+// CK2: store i{{[0-9]+}} [[B_IN]], i{{[0-9]+}}* [[BADDR]],
+// CK2: store i{{[0-9]+}} [[ARGC_IN]], i{{[0-9]+}}* [[ARGCADDR]],
+// CK2-64: [[ACONV:%.+]] = bitcast i64* [[AADDR]] to i32*
+// CK2-64: [[BCONV:%.+]] = bitcast i64* [[BADDR]] to i32*
+// CK2-64: [[CONV:%.+]] = bitcast i64* [[ARGCADDR]] to i32*
+// CK2-64:  store i{{[0-9]+}}* [[CONV]], i{{[0-9]+}}** [[ARGCADDR_PTR]],
+// CK2-32:  store i{{[0-9]+}}* [[ARGCADDR]], i{{[0-9]+}}** [[ARGCADDR_PTR]],
+// CK2:  [[ARGCADDR_PTR_REF:%.+]] = load i{{[0-9]+}}*, i{{[0-9]+}}** [[ARGCADDR_PTR]],
+// CK2: store i{{[0-9]+}} 0, i{{[0-9]+}}* [[ARGCADDR_PTR_REF]],
+// CK2-NOT:  {{.+}} = call i32 @__kmpc_push_num_teams(
+// CK2-NOT:  call void (%ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_teams(
+// CK2: ret
+
+// CK2: define {{.*}}void @{{[^,]+}}(i{{[0-9]+}}*{{.+}} [[A_IN:%.+]], i{{[0-9]+}}*{{.+}} [[BP:%.+]], i{{[0-9]+}}***{{.+}}  [[ARGC:%.+]])
+// CK2: [[ARGCADDR_PTR:%.+]] = alloca i{{[0-9]+}}***,
+// CK2: [[AADDR:%.+]] = alloca i{{[0-9]+}}*,
+// CK2: [[BADDR:%.+]] = alloca i{{[0-9]+}}*,
+// CK2: [[ARGCADDR:%.+]] = alloca i{{[0-9]+}}***,
+// CK2-NOT: {{%.+}} = call i32 @__kmpc_global_thread_num(
+// CK2: store i{{[0-9]+}}* [[A_IN]], i{{[0-9]+}}** [[AADDR]],
+// CK2: store i{{[0-9]+}}* [[B_IN]], i{{[0-9]+}}** [[BADDR]],
+// CK2: store i{{[0-9]+}}*** [[ARGC]], i{{[0-9]+}}**** [[ARGCADDR]],
+// CK2: [[A_ADDR_VAL:%.+]] = load i32*, i32** [[AADDR]]
+// CK2: [[B_ADDR_VAL:%.+]] = load i32*, i32** [[BADDR]]
+// CK2: [[ARGC_ADDR_VAL:%.+]] = load i{{[0-9]+}}***, i{{[0-9]+}}**** [[ARGCADDR]]
+// CK2: store i{{[0-9]+}}*** [[ARGC_ADDR_VAL]], i{{[0-9]+}}**** [[ARGCADDR_PTR]],
+// CK2: [[ARGCADDR_PTR_REF:%.+]] = load i{{[0-9]+}}***, i{{[0-9]+}}**** [[ARGCADDR_PTR]],
+// CK2: store i{{[0-9]+}}** null, i{{[0-9]+}}*** [[ARGCADDR_PTR_REF]],
+// CK2-NOT: {{.+}} = call i32 @__kmpc_push_num_teams(
+// CK2-NOT: call void (%ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_teams(
+// CK2:  ret void
+
+#endif // CK2
+#endif