[OpenMP][NVPTX] Enable default scheduling for parallel for in non-SPMD cases.

Summary: This patch enables the choosing of the default schedule for parallel for loops even in non-SPMD cases.

Reviewers: ABataev, caomhin

Reviewed By: ABataev

Subscribers: jholewinski, guansong, cfe-commits

Differential Revision: https://reviews.llvm.org/D53443

llvm-svn: 345507

2 files changed, 30 insertions, 7 deletions

diff --git a/clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp b/clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp
index 2aa8e83c2c0..c8dd1a6f7b6 100644
--- a/clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp
+++ b/clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp
@@ -4238,16 +4238,17 @@ void CGOpenMPRuntimeNVPTX::getDefaultDistScheduleAndChunk(
     Chunk = CGF.EmitScalarConversion(getNVPTXNumThreads(CGF),
         CGF.getContext().getIntTypeForBitwidth(32, /*Signed=*/0),
         S.getIterationVariable()->getType(), S.getBeginLoc());
+    return;
   }
+  CGOpenMPRuntime::getDefaultDistScheduleAndChunk(
+      CGF, S, ScheduleKind, Chunk);
 }
 
 void CGOpenMPRuntimeNVPTX::getDefaultScheduleAndChunk(
     CodeGenFunction &CGF, const OMPLoopDirective &S,
     OpenMPScheduleClauseKind &ScheduleKind,
     llvm::Value *&Chunk) const {
-  if (getExecutionMode() == CGOpenMPRuntimeNVPTX::EM_SPMD) {
-    ScheduleKind = OMPC_SCHEDULE_static;
-    Chunk = CGF.Builder.getIntN(CGF.getContext().getTypeSize(
-        S.getIterationVariable()->getType()), 1);
-  }
+  ScheduleKind = OMPC_SCHEDULE_static;
+  Chunk = CGF.Builder.getIntN(CGF.getContext().getTypeSize(
+      S.getIterationVariable()->getType()), 1);
 }
diff --git a/clang/test/OpenMP/nvptx_parallel_for_codegen.cpp b/clang/test/OpenMP/nvptx_parallel_for_codegen.cpp
index 286abb71aed..c61ba91bbfd 100644
--- a/clang/test/OpenMP/nvptx_parallel_for_codegen.cpp
+++ b/clang/test/OpenMP/nvptx_parallel_for_codegen.cpp
@@ -57,7 +57,10 @@ int bar(int n){
 // CHECK: store i32 0, {{.*}} [[OMP_LB:%.+]],
 // CHECK: store i32 9, {{.*}} [[OMP_UB:%.+]],
 // CHECK: store i32 1, {{.*}} [[OMP_ST:%.+]],
-// CHECK: call void @__kmpc_for_static_init_4({{.*}} i32 34, {{.*}} [[OMP_LB]], {{.*}} [[OMP_UB]], {{.*}} [[OMP_ST]], i32 1, i32 1)
+// CHECK: call void @__kmpc_for_static_init_4({{.*}} i32 33, {{.*}} [[OMP_LB]], {{.*}} [[OMP_UB]], {{.*}} [[OMP_ST]], i32 1, i32 1)
+// CHECK: br label %[[OMP_DISPATCH_COND:.+]]
+
+// CHECK: [[OMP_DISPATCH_COND]]
 // CHECK: [[OMP_UB_1:%.+]] = load {{.*}} [[OMP_UB]]
 // CHECK: [[COMP_1:%.+]] = icmp sgt {{.*}} [[OMP_UB_1]]
 // CHECK: br i1 [[COMP_1]], label %[[COND_TRUE:.+]], label %[[COND_FALSE:.+]]
@@ -74,6 +77,12 @@ int bar(int n){
 // CHECK: store i32 [[COND_RES]], i32* [[OMP_UB]]
 // CHECK: [[OMP_LB_1:%.+]] = load i32, i32* [[OMP_LB]]
 // CHECK: store i32 [[OMP_LB_1]], i32* [[OMP_IV]]
+// CHECK: [[OMP_IV_1:%.+]] = load i32, i32* [[OMP_IV]]
+// CHECK: [[OMP_UB_3:%.+]] = load i32, i32* [[OMP_UB]]
+// CHECK: [[COMP_2:%.+]] = icmp sle i32 [[OMP_IV_1]], [[OMP_UB_3]]
+// CHECK: br i1 [[COMP_2]], label %[[DISPATCH_BODY:.+]], label %[[DISPATCH_END:.+]]
+
+// CHECK: [[DISPATCH_BODY]]
 // CHECK: br label %[[OMP_INNER_FOR_COND:.+]]
 
 // CHECK: [[OMP_INNER_FOR_COND]]
@@ -94,7 +103,20 @@ int bar(int n){
 // CHECK: store i32 [[ADD_1]], i32* [[OMP_IV]]
 // CHECK: br label %[[OMP_INNER_FOR_COND]]
 
-// CHECK: [[OMP_INNER_FOR_END]]
+// CHECK: [[OMP_INNER_FOR_COND]]
+// CHECK: br label %[[OMP_DISPATCH_INC:.+]]
+
+// CHECK: [[OMP_DISPATCH_INC]]
+// CHECK: [[OMP_LB_2:%.+]] = load i32, i32* [[OMP_LB]]
+// CHECK: [[OMP_ST_1:%.+]] = load i32, i32* [[OMP_ST]]
+// CHECK: [[ADD_2:%.+]] = add nsw i32 [[OMP_LB_2]], [[OMP_ST_1]]
+// CHECK: store i32 [[ADD_2]], i32* [[OMP_LB]]
+// CHECK: [[OMP_UB_5:%.+]] = load i32, i32* [[OMP_UB]]
+// CHECK: [[OMP_ST_2:%.+]] = load i32, i32* [[OMP_ST]]
+// CHECK: [[ADD_3:%.+]] = add nsw i32 [[OMP_UB_5]], [[OMP_ST_2]]
+// CHECK: store i32 [[ADD_3]], i32* [[OMP_UB]]
+
+// CHECK: [[DISPATCH_END]]
 // CHECK: call void @__kmpc_for_static_fini(
 // CHECK: ret void