[clang][OpenMP] Diagnose badly-formed collapsed imperfect loop nests (#60678)

This patch fixes a couple of cases where Clang aborts with loop nests
that are being collapsed (via the relevant OpenMP clause) into a new,
combined loop.

The problematic cases happen when a variable declared within the
loop nest is used in the (init, condition, iter) statement of a more
deeply-nested loop.  I don't think these cases (generally?) fall under
the non-rectangular loop nest rules as defined in OpenMP 5.0+, but I
could be wrong (and anyway, emitting an error is better than crashing).

In terms of implementation: the crash happens because (to a first
approximation) all the loop bounds calculations are pulled out to the
start of the new, combined loop, but variables declared in the loop nest
"haven't been seen yet".  I believe there is special handling for
iteration variables declared in "for" init statements, but not for
variables declared elsewhere in the "imperfect" parts of a loop nest.

So, this patch tries to diagnose the troublesome cases before they can
cause a crash.  This is slightly awkward because at the point where we
want to do the diagnosis (SemaOpenMP.cpp), we don't have scope information
readily available.  Instead we "manually" scan through the AST of the
loop nest looking for var decls (ForVarDeclFinder), then we ensure we're
not using any of those in loop control subexprs (ForSubExprChecker).
All that is only done when we have a "collapse" clause.

Range-for loops can also cause crashes at present without this patch,
so are handled too.

Author: jtb20

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -11152,6 +11152,8 @@ def err_omp_loop_diff_cxx : Error<
   "upper and lower loop bounds">;
 def err_omp_loop_cannot_use_stmt : Error<
   "'%0' statement cannot be used in OpenMP for loop">;
+def err_omp_loop_bad_collapse_var : Error<
+  "cannot use variable %1 in collapsed imperfectly-nested loop %select{init|condition|increment}0 statement">;
 def err_omp_simd_region_cannot_use_stmt : Error<
   "'%0' statement cannot be used in OpenMP simd region">;
 def warn_omp_loop_64_bit_var : Warning<

clang/lib/AST/StmtOpenMP.cpp

@@ -10,8 +10,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "clang/AST/ASTContext.h"
 #include "clang/AST/StmtOpenMP.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 
 using namespace clang;
 using namespace llvm::omp;

clang/lib/Sema/SemaOpenMP.cpp

@@ -21,6 +21,7 @@
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclOpenMP.h"
 #include "clang/AST/OpenMPClause.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/StmtCXX.h"
 #include "clang/AST/StmtOpenMP.h"
 #include "clang/AST/StmtVisitor.h"
@@ -7668,6 +7669,52 @@ struct LoopIterationSpace final {
   Expr *FinalCondition = nullptr;
 };
 
+/// Scan an AST subtree, checking that no decls in the CollapsedLoopVarDecls
+/// set are referenced.  Used for verifying loop nest structure before
+/// performing a loop collapse operation.
+class ForSubExprChecker final : public RecursiveASTVisitor<ForSubExprChecker> {
+  const llvm::SmallPtrSetImpl<const Decl *> &CollapsedLoopVarDecls;
+  VarDecl *ForbiddenVar = nullptr;
+  SourceRange ErrLoc;
+
+public:
+  explicit ForSubExprChecker(
+      const llvm::SmallPtrSetImpl<const Decl *> &CollapsedLoopVarDecls)
+      : CollapsedLoopVarDecls(CollapsedLoopVarDecls) {}
+
+  // We want to visit implicit code, i.e. synthetic initialisation statements
+  // created during range-for lowering.
+  bool shouldVisitImplicitCode() const { return true; }
+
+  bool VisitDeclRefExpr(DeclRefExpr *E) {
+    ValueDecl *VD = E->getDecl();
+    if (!isa<VarDecl, BindingDecl>(VD))
+      return true;
+    VarDecl *V = VD->getPotentiallyDecomposedVarDecl();
+    if (V->getType()->isReferenceType()) {
+      VarDecl *VD = V->getDefinition();
+      if (VD->hasInit()) {
+        Expr *I = VD->getInit();
+        DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(I);
+        if (!DRE)
+          return true;
+        V = DRE->getDecl()->getPotentiallyDecomposedVarDecl();
+      }
+    }
+    Decl *Canon = V->getCanonicalDecl();
+    if (CollapsedLoopVarDecls.contains(Canon)) {
+      ForbiddenVar = V;
+      ErrLoc = E->getSourceRange();
+      return false;
+    }
+
+    return true;
+  }
+
+  VarDecl *getForbiddenVar() const { return ForbiddenVar; }
+  SourceRange getErrRange() const { return ErrLoc; }
+};
+
 /// Helper class for checking canonical form of the OpenMP loops and
 /// extracting iteration space of each loop in the loop nest, that will be used
 /// for IR generation.
@@ -7682,6 +7729,8 @@ class OpenMPIterationSpaceChecker {
   SourceLocation DefaultLoc;
   /// A location for diagnostics (when increment is not compatible).
   SourceLocation ConditionLoc;
+  /// The set of variables declared within the (to be collapsed) loop nest.
+  const llvm::SmallPtrSetImpl<const Decl *> *CollapsedLoopVarDecls;
   /// A source location for referring to loop init later.
   SourceRange InitSrcRange;
   /// A source location for referring to condition later.
@@ -7725,10 +7774,13 @@ class OpenMPIterationSpaceChecker {
   Expr *Condition = nullptr;
 
 public:
-  OpenMPIterationSpaceChecker(Sema &SemaRef, bool SupportsNonRectangular,
-                              DSAStackTy &Stack, SourceLocation DefaultLoc)
+  OpenMPIterationSpaceChecker(
+      Sema &SemaRef, bool SupportsNonRectangular, DSAStackTy &Stack,
+      SourceLocation DefaultLoc,
+      const llvm::SmallPtrSetImpl<const Decl *> *CollapsedLoopDecls)
       : SemaRef(SemaRef), SupportsNonRectangular(SupportsNonRectangular),
-        Stack(Stack), DefaultLoc(DefaultLoc), ConditionLoc(DefaultLoc) {}
+        Stack(Stack), DefaultLoc(DefaultLoc), ConditionLoc(DefaultLoc),
+        CollapsedLoopVarDecls(CollapsedLoopDecls) {}
   /// Check init-expr for canonical loop form and save loop counter
   /// variable - #Var and its initialization value - #LB.
   bool checkAndSetInit(Stmt *S, bool EmitDiags = true);
@@ -8049,6 +8101,16 @@ bool OpenMPIterationSpaceChecker::checkAndSetInit(Stmt *S, bool EmitDiags) {
     if (!ExprTemp->cleanupsHaveSideEffects())
       S = ExprTemp->getSubExpr();
 
+  if (CollapsedLoopVarDecls) {
+    ForSubExprChecker FSEC{*CollapsedLoopVarDecls};
+    if (!FSEC.TraverseStmt(S)) {
+      SourceRange Range = FSEC.getErrRange();
+      SemaRef.Diag(Range.getBegin(), diag::err_omp_loop_bad_collapse_var)
+          << Range.getEnd() << 0 << FSEC.getForbiddenVar();
+      return true;
+    }
+  }
+
   InitSrcRange = S->getSourceRange();
   if (Expr *E = dyn_cast<Expr>(S))
     S = E->IgnoreParens();
@@ -8152,6 +8214,17 @@ bool OpenMPIterationSpaceChecker::checkAndSetCond(Expr *S) {
   }
   Condition = S;
   S = getExprAsWritten(S);
+
+  if (CollapsedLoopVarDecls) {
+    ForSubExprChecker FSEC{*CollapsedLoopVarDecls};
+    if (!FSEC.TraverseStmt(S)) {
+      SourceRange Range = FSEC.getErrRange();
+      SemaRef.Diag(Range.getBegin(), diag::err_omp_loop_bad_collapse_var)
+          << Range.getEnd() << 1 << FSEC.getForbiddenVar();
+      return true;
+    }
+  }
+
   SourceLocation CondLoc = S->getBeginLoc();
   auto &&CheckAndSetCond =
       [this, IneqCondIsCanonical](BinaryOperatorKind Opcode, const Expr *LHS,
@@ -8250,6 +8323,16 @@ bool OpenMPIterationSpaceChecker::checkAndSetInc(Expr *S) {
     if (!ExprTemp->cleanupsHaveSideEffects())
       S = ExprTemp->getSubExpr();
 
+  if (CollapsedLoopVarDecls) {
+    ForSubExprChecker FSEC{*CollapsedLoopVarDecls};
+    if (!FSEC.TraverseStmt(S)) {
+      SourceRange Range = FSEC.getErrRange();
+      SemaRef.Diag(Range.getBegin(), diag::err_omp_loop_bad_collapse_var)
+          << Range.getEnd() << 2 << FSEC.getForbiddenVar();
+      return true;
+    }
+  }
+
   IncrementSrcRange = S->getSourceRange();
   S = S->IgnoreParens();
   if (auto *UO = dyn_cast<UnaryOperator>(S)) {
@@ -8972,7 +9055,7 @@ void SemaOpenMP::ActOnOpenMPLoopInitialization(SourceLocation ForLoc,
 
   DSAStack->loopStart();
   OpenMPIterationSpaceChecker ISC(SemaRef, /*SupportsNonRectangular=*/true,
-                                  *DSAStack, ForLoc);
+                                  *DSAStack, ForLoc, nullptr);
   if (!ISC.checkAndSetInit(Init, /*EmitDiags=*/false)) {
     if (ValueDecl *D = ISC.getLoopDecl()) {
       auto *VD = dyn_cast<VarDecl>(D);
@@ -9069,7 +9152,8 @@ static bool checkOpenMPIterationSpace(
     Expr *OrderedLoopCountExpr,
     SemaOpenMP::VarsWithInheritedDSAType &VarsWithImplicitDSA,
     llvm::MutableArrayRef<LoopIterationSpace> ResultIterSpaces,
-    llvm::MapVector<const Expr *, DeclRefExpr *> &Captures) {
+    llvm::MapVector<const Expr *, DeclRefExpr *> &Captures,
+    const llvm::SmallPtrSetImpl<const Decl *> &CollapsedLoopVarDecls) {
   bool SupportsNonRectangular = !isOpenMPLoopTransformationDirective(DKind);
   // OpenMP [2.9.1, Canonical Loop Form]
   //   for (init-expr; test-expr; incr-expr) structured-block
@@ -9108,7 +9192,8 @@ static bool checkOpenMPIterationSpace(
     return false;
 
   OpenMPIterationSpaceChecker ISC(SemaRef, SupportsNonRectangular, DSA,
-                                  For ? For->getForLoc() : CXXFor->getForLoc());
+                                  For ? For->getForLoc() : CXXFor->getForLoc(),
+                                  &CollapsedLoopVarDecls);
 
   // Check init.
   Stmt *Init = For ? For->getInit() : CXXFor->getBeginStmt();
@@ -9475,6 +9560,39 @@ static Expr *buildPostUpdate(Sema &S, ArrayRef<Expr *> PostUpdates) {
   return PostUpdate;
 }
 
+/// Look for variables declared in the body parts of a for-loop nest.  Used
+/// for verifying loop nest structure before performing a loop collapse
+/// operation.
+class ForVarDeclFinder final : public RecursiveASTVisitor<ForVarDeclFinder> {
+  int NestingDepth = 0;
+  llvm::SmallPtrSetImpl<const Decl *> &VarDecls;
+
+public:
+  explicit ForVarDeclFinder(llvm::SmallPtrSetImpl<const Decl *> &VD)
+      : VarDecls(VD) {}
+
+  bool VisitForStmt(ForStmt *F) {
+    ++NestingDepth;
+    TraverseStmt(F->getBody());
+    --NestingDepth;
+    return false;
+  }
+
+  bool VisitCXXForRangeStmt(CXXForRangeStmt *RF) {
+    ++NestingDepth;
+    TraverseStmt(RF->getBody());
+    --NestingDepth;
+    return false;
+  }
+
+  bool VisitVarDecl(VarDecl *D) {
+    Decl *C = D->getCanonicalDecl();
+    if (NestingDepth > 0)
+      VarDecls.insert(C);
+    return true;
+  }
+};
+
 /// Called on a for stmt to check itself and nested loops (if any).
 /// \return Returns 0 if one of the collapsed stmts is not canonical for loop,
 /// number of collapsed loops otherwise.
@@ -9487,13 +9605,17 @@ checkOpenMPLoop(OpenMPDirectiveKind DKind, Expr *CollapseLoopCountExpr,
   unsigned NestedLoopCount = 1;
   bool SupportsNonPerfectlyNested = (SemaRef.LangOpts.OpenMP >= 50) &&
                                     !isOpenMPLoopTransformationDirective(DKind);
+  llvm::SmallPtrSet<const Decl *, 4> CollapsedLoopVarDecls{};
 
   if (CollapseLoopCountExpr) {
     // Found 'collapse' clause - calculate collapse number.
     Expr::EvalResult Result;
     if (!CollapseLoopCountExpr->isValueDependent() &&
         CollapseLoopCountExpr->EvaluateAsInt(Result, SemaRef.getASTContext())) {
       NestedLoopCount = Result.Val.getInt().getLimitedValue();
+
+      ForVarDeclFinder FVDF{CollapsedLoopVarDecls};
+      FVDF.TraverseStmt(AStmt);
     } else {
       Built.clear(/*Size=*/1);
       return 1;
@@ -9531,11 +9653,13 @@ checkOpenMPLoop(OpenMPDirectiveKind DKind, Expr *CollapseLoopCountExpr,
           SupportsNonPerfectlyNested, NumLoops,
           [DKind, &SemaRef, &DSA, NumLoops, NestedLoopCount,
            CollapseLoopCountExpr, OrderedLoopCountExpr, &VarsWithImplicitDSA,
-           &IterSpaces, &Captures](unsigned Cnt, Stmt *CurStmt) {
+           &IterSpaces, &Captures,
+           &CollapsedLoopVarDecls](unsigned Cnt, Stmt *CurStmt) {
             if (checkOpenMPIterationSpace(
                     DKind, CurStmt, SemaRef, DSA, Cnt, NestedLoopCount,
                     NumLoops, CollapseLoopCountExpr, OrderedLoopCountExpr,
-                    VarsWithImplicitDSA, IterSpaces, Captures))
+                    VarsWithImplicitDSA, IterSpaces, Captures,
+                    CollapsedLoopVarDecls))
               return true;
             if (Cnt > 0 && Cnt >= NestedLoopCount &&
                 IterSpaces[Cnt].CounterVar) {

clang/test/OpenMP/loop_collapse_1.c

@@ -0,0 +1,40 @@
+// RUN: %clang_cc1 -fopenmp -fopenmp-version=50 -verify %s
+
+void func( double *A, int N, int M, int NB ) {
+#pragma omp parallel
+  {
+    int nblks = (N-1)/NB;
+    int lnb = ((N-1)/NB)*NB;
+
+#pragma omp for collapse(2)
+    for (int jblk = 0 ; jblk < nblks ; jblk++ ) {
+      int jb = (jblk == nblks - 1 ? lnb : NB);
+      for (int jk = 0; jk < N; jk+=jb) {  // expected-error{{cannot use variable 'jb' in collapsed imperfectly-nested loop increment statement}}
+      }
+    }
+
+#pragma omp for collapse(2)
+    for (int a = 0; a < N; a++) {
+      for (int b = 0; b < M; b++) {
+        int cx = a+b < NB ? a : b;
+        for (int c = 0; c < cx; c++) {
+        }
+      }
+    }
+
+#pragma omp for collapse(3)
+    for (int a = 0; a < N; a++) {
+      for (int b = 0; b < M; b++) {
+        int cx = a+b < NB ? a : b;
+        for (int c = 0; c < cx; c++) {  // expected-error{{cannot use variable 'cx' in collapsed imperfectly-nested loop condition statement}}
+        }
+      }
+    }
+  }
+}
+
+int main(void) {
+  double arr[256];
+  func (arr, 16, 16, 16);
+  return 0;
+}

clang/test/OpenMP/loop_collapse_2.cpp

@@ -0,0 +1,80 @@
+// RUN: %clang_cc1 -fopenmp -fopenmp-version=50 -verify %s
+
+// We just want to try out a range for statement... this seems a bit OTT.
+template<typename T>
+class fakevector {
+  T *contents;
+  long size;
+public:
+  fakevector(long sz) : size(sz) {
+    contents = new T[sz];
+  }
+  ~fakevector() {
+    delete[] contents;
+  }
+  T& operator[](long x) { return contents[x]; }
+  typedef T *iterator;
+  fakevector<T>::iterator begin() {
+    return &contents[0];
+  }
+  fakevector<T>::iterator end() {
+    return &contents[size];
+  }
+};
+
+void func( double *A, int N, int M, int NB ) {
+#pragma omp parallel
+  {
+    int nblks = (N-1)/NB;
+    int lnb = ((N-1)/NB)*NB;
+#pragma omp for collapse(2)
+    for (int jblk = 0 ; jblk < nblks ; jblk++ ) {
+      int jb = (jblk == nblks - 1 ? lnb : NB);
+      for (int jk = 0; jk < N; jk+=jb) {  // expected-error{{cannot use variable 'jb' in collapsed imperfectly-nested loop increment statement}}
+      }
+    }
+
+#pragma omp for collapse(2)
+    for (int a = 0; a < N; a++) {
+        for (int b = 0; b < M; b++) {
+          int cx = a+b < NB ? a : b;
+          for (int c = 0; c < cx; c++) {
+          }
+        }
+    }
+
+    fakevector<float> myvec{N};
+#pragma omp for collapse(2)
+    for (auto &a : myvec) {
+      fakevector<float> myvec3{M};
+      for (auto &b : myvec3) {  // expected-error{{cannot use variable 'myvec3' in collapsed imperfectly-nested loop init statement}}
+      }
+    }
+
+    fakevector<float> myvec2{M};
+
+#pragma omp for collapse(3)
+    for (auto &a : myvec) {
+      for (auto &b : myvec2) {
+        int cx = a < b ? N : M;
+        for (int c = 0; c < cx; c++) {  // expected-error {{cannot use variable 'cx' in collapsed imperfectly-nested loop condition statement}}
+        }
+      }
+    }
+
+#pragma omp for collapse(3)
+    for (auto &a : myvec) {
+      int cx = a < 5 ? M : N;
+      for (auto &b : myvec2) {
+        for (int c = 0; c < cx; c++) {  // expected-error{{cannot use variable 'cx' in collapsed imperfectly-nested loop condition statement}}
+        }
+      }
+    }
+  }
+}
+
+int main(void) {
+  double arr[256];
+  func (arr, 16, 16, 16);
+  return 0;
+}