[OpenMP] Fix iterations calculation for dependent counters

[mikerice]

Summary

The number of iterations calculation was failing in some cases with more
than two collpased loops. Now the LoopIterationSpace selected matches
InitDependOnLC and CondDependOnLC.

[ABataev]

LG

clang/lib/Sema/SemaOpenMP.cpp

Expr *OpenMPIterationSpaceChecker::buildNumIterations(
   if (!VarType->isIntegerType() && !VarType->isPointerType() &&
       !SemaRef.getLangOpts().CPlusPlus)
     return nullptr;
   Expr *LBVal = LB;
   Expr *UBVal = UB;
   // LB = TestIsLessOp.getValue() ? min(LB(MinVal), LB(MaxVal)) :
   // max(LB(MinVal), LB(MaxVal))
   if (InitDependOnLC) {
-    const LoopIterationSpace &IS =
-        ResultIterSpaces[ResultIterSpaces.size() - 1 -
-                         InitDependOnLC.getValueOr(
-                             CondDependOnLC.getValueOr(0))];
+    const LoopIterationSpace &IS = ResultIterSpaces[*InitDependOnLC - 1];
     if (!IS.MinValue || !IS.MaxValue)
       return nullptr;
     // OuterVar = Min
     ExprResult MinValue =
         SemaRef.ActOnParenExpr(DefaultLoc, DefaultLoc, IS.MinValue);
     if (!MinValue.isUsable())
       return nullptr;
 
if (TestIsLessOp.getValue()) {
       if (!MaxLB.isUsable())
         return nullptr;
       LBVal = MaxLB.get();
     }
   }
   // UB = TestIsLessOp.getValue() ? max(UB(MinVal), UB(MaxVal)) :
   // min(UB(MinVal), UB(MaxVal))
   if (CondDependOnLC) {
-    const LoopIterationSpace &IS =
-        ResultIterSpaces[ResultIterSpaces.size() - 1 -
-                         InitDependOnLC.getValueOr(
-                             CondDependOnLC.getValueOr(0))];
+    const LoopIterationSpace &IS = ResultIterSpaces[*CondDependOnLC - 1];
     if (!IS.MinValue || !IS.MaxValue)
       return nullptr;
     // OuterVar = Min
     ExprResult MinValue =
         SemaRef.ActOnParenExpr(DefaultLoc, DefaultLoc, IS.MinValue);
     if (!MinValue.isUsable())
       return nullptr;
 

clang/test/OpenMP/for_codegen.cpp

void loop_with_counter_collapse() {
   // LIFETIME: call void @llvm.lifetime.end
   // LIFETIME: call void @llvm.lifetime.end
   #pragma omp for collapse(2)
   for (int i = 0; i < 4; i++) {
     for (int j = i; j < 4 + i; j++) {
     }
   }
 }
+
+// CHECK-LABEL: loop_with_counter_collapse4
+void loop_with_counter_collapse4() {
+
+  // Check bounds calculation when collapse > 2
+  // CHECK: store i32 0, i32* [[I_TMP:%.+]],
+  // CHECK: [[VAL:%.+]] = load i32, i32* [[I_TMP]],
+  // CHECK: store i32 [[VAL]], i32* [[K_LB_MIN:%.+]],
+  // CHECK: store i32 6, i32* [[I_TMP]],
+  // CHECK: [[VAL:%.+]] = load i32, i32* [[I_TMP]],
+  // CHECK: store i32 [[VAL]], i32* [[K_LB_MAX:%.+]],
+  #pragma omp for collapse(4)
+  for (int i = 0; i < 7; i++) {
+    for (int j = 0; j < 11; j++) {
+      for (int k = i; k < 7; k++) {
+        for (int l = 0; l < 11; l++) {
+        }
+      }
+    }
+  }
+}
+
 // CHECK-LABEL: define {{.*void}} @{{.*}}without_schedule_clause{{.*}}(float* {{.+}}, float* {{.+}}, float* {{.+}}, float* {{.+}})
 void without_schedule_clause(float *a, float *b, float *c, float *d) {
 // CHECK: [[GTID:%.+]] = call i32 @__kmpc_global_thread_num([[IDENT_T_TY]]* [[DEFAULT_LOC:[@%].+]])
   #pragma omp for nowait
 // CHECK: call void @__kmpc_for_static_init_4([[IDENT_T_TY]]* [[LOOP_LOC]], i32 [[GTID]], i32 34, i32* [[IS_LAST:%[^,]+]], i32* [[OMP_LB:%[^,]+]], i32* [[OMP_UB:%[^,]+]], i32* [[OMP_ST:%[^,]+]], i32 1, i32 1)
 // UB = min(UB, GlobalUB)
 // CHECK-NEXT: [[UB:%.+]] = load i32, i32* [[OMP_UB]]
 // CHECK-NEXT: [[UBCMP:%.+]] = icmp sgt i32 [[UB]], 4571423

clang/test/OpenMP/for_loop_messages.cpp

 // expected-error@+4 2 {{expected loop invariant expression or '<invariant1> * ii + <invariant2>' kind of expression}}
 // expected-error@+3 {{expected loop invariant expression or '<invariant1> * TC::ii + <invariant2>' kind of expression}}
 #pragma omp for collapse(2)
     for (ii = 10 + 25; ii < 1000; ii += 1)
       for (iii = ii * 10 + 25; iii < ii / ii - 23; iii += 1)
         ;
 
 #pragma omp parallel
-// expected-error@+5 2 {{expected loop invariant expression or '<invariant1> * ii + <invariant2>' kind of expression}}
-// expected-error@+4 {{expected loop invariant expression or '<invariant1> * TC::ii + <invariant2>' kind of expression}}
-// expected-error@+4 {{expected loop invariant expression or '<invariant1> * TC::ii + <invariant2>' kind of expression}}
+// expected-error@+6 {{expected loop invariant expression or '<invariant1> * TC::ii + <invariant2>' kind of expression}}
+// expected-error@+6 {{expected loop invariant expression or '<invariant1> * TC::ii + <invariant2>' kind of expression}}
+// expected-error@+4 2 {{expected loop invariant expression or '<invariant1> * ii + <invariant2>' kind of expression}}
+// expected-error@+4 2 {{expected loop invariant expression or '<invariant1> * ii + <invariant2>' kind of expression}}
 #pragma omp for collapse(3)
     for (ii = 10 + 25; ii < 1000; ii += 1)
       for (iii = ii * 10 + 25; iii < ii / ii - 23; iii += 1)
         for (kk = ii * 10 + 25; kk < iii - 23; kk += 1)
           ;
 
 #pragma omp parallel
 // expected-note@+3 {{loop step is expected to be positive due to this condition}}