Add support for '#pragma unroll'.

llvm-svn: 213574

3 files changed, 223 insertions, 127 deletions

diff --git a/clang/lib/Parse/ParsePragma.cpp b/clang/lib/Parse/ParsePragma.cpp
index 1b343b4291b..a7b5f6ebddb 100644
--- a/clang/lib/Parse/ParsePragma.cpp
+++ b/clang/lib/Parse/ParsePragma.cpp
@@ -148,6 +148,12 @@ struct PragmaLoopHintHandler : public PragmaHandler {
                     Token &FirstToken) override;
 };
 
+struct PragmaUnrollHintHandler : public PragmaHandler {
+  PragmaUnrollHintHandler(const char *name) : PragmaHandler(name) {}
+  void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
+                    Token &FirstToken) override;
+};
+
 }  // end namespace
 
 void Parser::initializePragmaHandlers() {
@@ -218,6 +224,9 @@ void Parser::initializePragmaHandlers() {
 
   LoopHintHandler.reset(new PragmaLoopHintHandler());
   PP.AddPragmaHandler("clang", LoopHintHandler.get());
+
+  UnrollHintHandler.reset(new PragmaUnrollHintHandler("unroll"));
+  PP.AddPragmaHandler(UnrollHintHandler.get());
 }
 
 void Parser::resetPragmaHandlers() {
@@ -278,6 +287,9 @@ void Parser::resetPragmaHandlers() {
 
   PP.RemovePragmaHandler("clang", LoopHintHandler.get());
   LoopHintHandler.reset();
+
+  PP.RemovePragmaHandler(UnrollHintHandler.get());
+  UnrollHintHandler.reset();
 }
 
 /// \brief Handle the annotation token produced for #pragma unused(...)
@@ -649,9 +661,10 @@ bool Parser::HandlePragmaMSInitSeg(StringRef PragmaName,
 }
 
 struct PragmaLoopHintInfo {
-  Token Loop;
-  Token Value;
+  Token PragmaName;
   Token Option;
+  Token Value;
+  bool HasValue;
 };
 
 LoopHint Parser::HandlePragmaLoopHint() {
@@ -660,24 +673,30 @@ LoopHint Parser::HandlePragmaLoopHint() {
       static_cast<PragmaLoopHintInfo *>(Tok.getAnnotationValue());
 
   LoopHint Hint;
-  Hint.LoopLoc =
-      IdentifierLoc::create(Actions.Context, Info->Loop.getLocation(),
-                            Info->Loop.getIdentifierInfo());
+  Hint.PragmaNameLoc =
+      IdentifierLoc::create(Actions.Context, Info->PragmaName.getLocation(),
+                            Info->PragmaName.getIdentifierInfo());
   Hint.OptionLoc =
       IdentifierLoc::create(Actions.Context, Info->Option.getLocation(),
                             Info->Option.getIdentifierInfo());
-  Hint.ValueLoc =
-      IdentifierLoc::create(Actions.Context, Info->Value.getLocation(),
-                            Info->Value.getIdentifierInfo());
-  Hint.Range =
-      SourceRange(Info->Option.getLocation(), Info->Value.getLocation());
-
-  // FIXME: We should allow non-type template parameters for the loop hint
-  // value. See bug report #19610
-  if (Info->Value.is(tok::numeric_constant))
-    Hint.ValueExpr = Actions.ActOnNumericConstant(Info->Value).get();
-  else
+  if (Info->HasValue) {
+    Hint.Range =
+        SourceRange(Info->Option.getLocation(), Info->Value.getLocation());
+    Hint.ValueLoc =
+        IdentifierLoc::create(Actions.Context, Info->Value.getLocation(),
+                              Info->Value.getIdentifierInfo());
+
+    // FIXME: We should allow non-type template parameters for the loop hint
+    // value. See bug report #19610
+    if (Info->Value.is(tok::numeric_constant))
+      Hint.ValueExpr = Actions.ActOnNumericConstant(Info->Value).get();
+    else
+      Hint.ValueExpr = nullptr;
+  } else {
+    Hint.Range = SourceRange(Info->PragmaName.getLocation());
+    Hint.ValueLoc = nullptr;
     Hint.ValueExpr = nullptr;
+  }
 
   return Hint;
 }
@@ -1650,7 +1669,9 @@ void PragmaOptimizeHandler::HandlePragma(Preprocessor &PP,
   Token Tok;
   PP.Lex(Tok);
   if (Tok.is(tok::eod)) {
-    PP.Diag(Tok.getLocation(), diag::err_pragma_optimize_missing_argument);
+    PP.Diag(Tok.getLocation(), diag::err_pragma_missing_argument)
+        << "clang optimize"
+        << "'on' or 'off'";
     return;
   }
   if (Tok.isNot(tok::identifier)) {
@@ -1679,6 +1700,48 @@ void PragmaOptimizeHandler::HandlePragma(Preprocessor &PP,
   Actions.ActOnPragmaOptimize(IsOn, FirstToken.getLocation());
 }
 
+/// \brief Parses loop or unroll pragma hint value and fills in Info.
+static bool ParseLoopHintValue(Preprocessor &PP, Token Tok, Token &PragmaName,
+                               Token &Option, bool &ValueInParens,
+                               PragmaLoopHintInfo &Info) {
+  ValueInParens = Tok.is(tok::l_paren);
+  if (ValueInParens) {
+    PP.Lex(Tok);
+    if (Tok.is(tok::r_paren)) {
+      // Nothing between the parentheses.
+      std::string PragmaString;
+      if (PragmaName.getIdentifierInfo()->getName() == "loop") {
+        PragmaString = "clang loop ";
+        PragmaString += Option.getIdentifierInfo()->getName();
+      } else {
+        assert(PragmaName.getIdentifierInfo()->getName() == "unroll" &&
+               "Unexpected pragma name");
+        PragmaString = "unroll";
+      }
+      PP.Diag(Tok.getLocation(), diag::err_pragma_missing_argument)
+          << PragmaString << "a positive integer value";
+      return true;
+    }
+  }
+
+  // FIXME: Value should be stored and parsed as a constant expression.
+  Token Value = Tok;
+
+  if (ValueInParens) {
+    PP.Lex(Tok);
+    if (Tok.isNot(tok::r_paren)) {
+      PP.Diag(Tok.getLocation(), diag::err_expected) << tok::r_paren;
+      return true;
+    }
+  }
+
+  Info.PragmaName = PragmaName;
+  Info.Option = Option;
+  Info.Value = Value;
+  Info.HasValue = true;
+  return false;
+}
+
 /// \brief Handle the \#pragma clang loop directive.
 ///  #pragma clang 'loop' loop-hints
 ///
@@ -1720,7 +1783,8 @@ void PragmaOptimizeHandler::HandlePragma(Preprocessor &PP,
 void PragmaLoopHintHandler::HandlePragma(Preprocessor &PP,
                                          PragmaIntroducerKind Introducer,
                                          Token &Tok) {
-  Token Loop = Tok;
+  // Incoming token is "loop" from "#pragma clang loop".
+  Token PragmaName = Tok;
   SmallVector<Token, 1> TokenList;
 
   // Lex the optimization option and verify it is an identifier.
@@ -1736,59 +1800,40 @@ void PragmaLoopHintHandler::HandlePragma(Preprocessor &PP,
     IdentifierInfo *OptionInfo = Tok.getIdentifierInfo();
 
     bool OptionValid = llvm::StringSwitch<bool>(OptionInfo->getName())
-        .Case("vectorize", true)
-        .Case("interleave", true)
-        .Case("unroll", true)
-        .Case("vectorize_width", true)
-        .Case("interleave_count", true)
-        .Case("unroll_count", true)
-        .Default(false);
+                           .Case("vectorize", true)
+                           .Case("interleave", true)
+                           .Case("unroll", true)
+                           .Case("vectorize_width", true)
+                           .Case("interleave_count", true)
+                           .Case("unroll_count", true)
+                           .Default(false);
     if (!OptionValid) {
       PP.Diag(Tok.getLocation(), diag::err_pragma_loop_invalid_option)
           << /*MissingOption=*/false << OptionInfo;
       return;
     }
 
-    // Read '('
-    PP.Lex(Tok);
-    if (Tok.isNot(tok::l_paren)) {
-      PP.Diag(Tok.getLocation(), diag::err_expected) << tok::l_paren;
-      return;
-    }
-
-    // FIXME: All tokens between '(' and ')' should be stored and parsed as a
-    // constant expression.
+    auto *Info = new (PP.getPreprocessorAllocator()) PragmaLoopHintInfo;
     PP.Lex(Tok);
-    if (Tok.is(tok::r_paren)) {
-      // Nothing between the parentheses.
-      PP.Diag(Tok.getLocation(), diag::err_pragma_loop_missing_argument)
-          << OptionInfo;
+    bool ValueInParens;
+    if (ParseLoopHintValue(PP, Tok, PragmaName, Option, ValueInParens, *Info))
       return;
-    }
-    Token Value = Tok;
 
-    // Read ')'
-    PP.Lex(Tok);
-    if (Tok.isNot(tok::r_paren)) {
-      PP.Diag(Tok.getLocation(), diag::err_expected) << tok::r_paren;
+    if (!ValueInParens) {
+      PP.Diag(Info->Value.getLocation(), diag::err_expected) << tok::l_paren;
       return;
     }
 
-    // Get next optimization option.
-    PP.Lex(Tok);
-
-    auto *Info = new (PP.getPreprocessorAllocator()) PragmaLoopHintInfo;
-    Info->Loop = Loop;
-    Info->Option = Option;
-    Info->Value = Value;
-
-    // Generate the vectorization hint token.
+    // Generate the loop hint token.
     Token LoopHintTok;
     LoopHintTok.startToken();
     LoopHintTok.setKind(tok::annot_pragma_loop_hint);
-    LoopHintTok.setLocation(Loop.getLocation());
+    LoopHintTok.setLocation(PragmaName.getLocation());
     LoopHintTok.setAnnotationValue(static_cast<void *>(Info));
     TokenList.push_back(LoopHintTok);
+
+    // Get next optimization option.
+    PP.Lex(Tok);
   }
 
   if (Tok.isNot(tok::eod)) {
@@ -1804,3 +1849,62 @@ void PragmaLoopHintHandler::HandlePragma(Preprocessor &PP,
                       /*DisableMacroExpansion=*/false,
                       /*OwnsTokens=*/true);
 }
+
+/// \brief Handle the loop unroll optimization pragmas.
+///  #pragma unroll
+///  #pragma unroll unroll-hint-value
+///  #pragma unroll '(' unroll-hint-value ')'
+///
+///  unroll-hint-value:
+///    constant-expression
+///
+/// Loop unrolling hints are specified with '#pragma unroll'. '#pragma unroll'
+/// can take a numeric argument optionally contained in parentheses. With no
+/// argument the directive instructs llvm to try to unroll the loop
+/// completely. A positive integer argument can be specified to indicate the
+/// number of times the loop should be unrolled.  To maximize compatibility with
+/// other compilers the unroll count argument can be specified with or without
+/// parentheses.
+void PragmaUnrollHintHandler::HandlePragma(Preprocessor &PP,
+                                           PragmaIntroducerKind Introducer,
+                                           Token &Tok) {
+  // Incoming token is "unroll" of "#pragma unroll".
+  Token PragmaName = Tok;
+  PP.Lex(Tok);
+  auto *Info = new (PP.getPreprocessorAllocator()) PragmaLoopHintInfo;
+  if (Tok.is(tok::eod)) {
+    // Unroll pragma without an argument.
+    Info->PragmaName = PragmaName;
+    Info->Option = PragmaName;
+    Info->HasValue = false;
+  } else {
+    // Unroll pragma with an argument: "#pragma unroll N" or
+    // "#pragma unroll(N)".
+    bool ValueInParens;
+    if (ParseLoopHintValue(PP, Tok, PragmaName, PragmaName, ValueInParens,
+                           *Info))
+      return;
+
+    // In CUDA, the argument to '#pragma unroll' should not be contained in
+    // parentheses.
+    if (PP.getLangOpts().CUDA && ValueInParens)
+      PP.Diag(Info->Value.getLocation(),
+              diag::warn_pragma_unroll_cuda_value_in_parens);
+
+    PP.Lex(Tok);
+    if (Tok.isNot(tok::eod)) {
+      PP.Diag(Tok.getLocation(), diag::warn_pragma_extra_tokens_at_eol)
+          << "unroll";
+      return;
+    }
+  }
+
+  // Generate the hint token.
+  Token *TokenArray = new Token[1];
+  TokenArray[0].startToken();
+  TokenArray[0].setKind(tok::annot_pragma_loop_hint);
+  TokenArray[0].setLocation(PragmaName.getLocation());
+  TokenArray[0].setAnnotationValue(static_cast<void *>(Info));
+  PP.EnterTokenStream(TokenArray, 1, /*DisableMacroExpansion=*/false,
+                      /*OwnsTokens=*/true);
+}
diff --git a/clang/lib/Parse/ParseStmt.cpp b/clang/lib/Parse/ParseStmt.cpp
index 1e8c4a69d86..ec0ca6b1fdd 100644
--- a/clang/lib/Parse/ParseStmt.cpp
+++ b/clang/lib/Parse/ParseStmt.cpp
@@ -1833,18 +1833,16 @@ StmtResult Parser::ParsePragmaLoopHint(StmtVector &Stmts, bool OnlyStatement,
   // Create temporary attribute list.
   ParsedAttributesWithRange TempAttrs(AttrFactory);
 
-  // Get vectorize hints and consume annotated token.
+  // Get loop hints and consume annotated token.
   while (Tok.is(tok::annot_pragma_loop_hint)) {
     LoopHint Hint = HandlePragmaLoopHint();
     ConsumeToken();
 
-    if (!Hint.LoopLoc || !Hint.OptionLoc || !Hint.ValueLoc)
-      continue;
-
-    ArgsUnion ArgHints[] = {Hint.OptionLoc, Hint.ValueLoc,
+    ArgsUnion ArgHints[] = {Hint.PragmaNameLoc, Hint.OptionLoc, Hint.ValueLoc,
                             ArgsUnion(Hint.ValueExpr)};
-    TempAttrs.addNew(Hint.LoopLoc->Ident, Hint.Range, nullptr,
-                     Hint.LoopLoc->Loc, ArgHints, 3, AttributeList::AS_Pragma);
+    TempAttrs.addNew(Hint.PragmaNameLoc->Ident, Hint.Range, nullptr,
+                     Hint.PragmaNameLoc->Loc, ArgHints, 4,
+                     AttributeList::AS_Pragma);
   }
 
   // Get the next statement.
diff --git a/clang/lib/Sema/SemaStmtAttr.cpp b/clang/lib/Sema/SemaStmtAttr.cpp
index 44169c2fdce..a32e0fbcb62 100644
--- a/clang/lib/Sema/SemaStmtAttr.cpp
+++ b/clang/lib/Sema/SemaStmtAttr.cpp
@@ -45,35 +45,50 @@ static Attr *handleFallThroughAttr(Sema &S, Stmt *St, const AttributeList &A,
 
 static Attr *handleLoopHintAttr(Sema &S, Stmt *St, const AttributeList &A,
                                 SourceRange) {
+  IdentifierLoc *PragmaNameLoc = A.getArgAsIdent(0);
+  IdentifierLoc *OptionLoc = A.getArgAsIdent(1);
+  IdentifierInfo *OptionInfo = OptionLoc->Ident;
+  IdentifierLoc *ValueLoc = A.getArgAsIdent(2);
+  IdentifierInfo *ValueInfo = ValueLoc ? ValueLoc->Ident : nullptr;
+  Expr *ValueExpr = A.getArgAsExpr(3);
+
+  assert(OptionInfo && "Attribute must have valid option info.");
+
   if (St->getStmtClass() != Stmt::DoStmtClass &&
       St->getStmtClass() != Stmt::ForStmtClass &&
       St->getStmtClass() != Stmt::CXXForRangeStmtClass &&
       St->getStmtClass() != Stmt::WhileStmtClass) {
-    S.Diag(St->getLocStart(), diag::err_pragma_loop_precedes_nonloop);
+    const char *Pragma = PragmaNameLoc->Ident->getName() == "unroll"
+                             ? "#pragma unroll"
+                             : "#pragma clang loop";
+    S.Diag(St->getLocStart(), diag::err_pragma_loop_precedes_nonloop) << Pragma;
     return nullptr;
   }
 
-  IdentifierLoc *OptionLoc = A.getArgAsIdent(0);
-  IdentifierInfo *OptionInfo = OptionLoc->Ident;
-  IdentifierLoc *ValueLoc = A.getArgAsIdent(1);
-  IdentifierInfo *ValueInfo = ValueLoc->Ident;
-  Expr *ValueExpr = A.getArgAsExpr(2);
-
-  assert(OptionInfo && "Attribute must have valid option info.");
-
-  LoopHintAttr::OptionType Option =
-      llvm::StringSwitch<LoopHintAttr::OptionType>(OptionInfo->getName())
-          .Case("vectorize", LoopHintAttr::Vectorize)
-          .Case("vectorize_width", LoopHintAttr::VectorizeWidth)
-          .Case("interleave", LoopHintAttr::Interleave)
-          .Case("interleave_count", LoopHintAttr::InterleaveCount)
-          .Case("unroll", LoopHintAttr::Unroll)
-          .Case("unroll_count", LoopHintAttr::UnrollCount)
-          .Default(LoopHintAttr::Vectorize);
+  LoopHintAttr::OptionType Option;
+  LoopHintAttr::Spelling Spelling;
+  if (PragmaNameLoc->Ident->getName() == "unroll") {
+    Option = ValueLoc ? LoopHintAttr::UnrollCount : LoopHintAttr::Unroll;
+    Spelling = LoopHintAttr::Pragma_unroll;
+  } else {
+    Option = llvm::StringSwitch<LoopHintAttr::OptionType>(OptionInfo->getName())
+                 .Case("vectorize", LoopHintAttr::Vectorize)
+                 .Case("vectorize_width", LoopHintAttr::VectorizeWidth)
+                 .Case("interleave", LoopHintAttr::Interleave)
+                 .Case("interleave_count", LoopHintAttr::InterleaveCount)
+                 .Case("unroll", LoopHintAttr::Unroll)
+                 .Case("unroll_count", LoopHintAttr::UnrollCount)
+                 .Default(LoopHintAttr::Vectorize);
+    Spelling = LoopHintAttr::Pragma_clang_loop;
+  }
 
   int ValueInt;
-  if (Option == LoopHintAttr::Vectorize || Option == LoopHintAttr::Interleave ||
-      Option == LoopHintAttr::Unroll) {
+  if (Option == LoopHintAttr::Unroll &&
+      Spelling == LoopHintAttr::Pragma_unroll) {
+    ValueInt = 1;
+  } else if (Option == LoopHintAttr::Vectorize ||
+             Option == LoopHintAttr::Interleave ||
+             Option == LoopHintAttr::Unroll) {
     if (!ValueInfo) {
       S.Diag(ValueLoc->Loc, diag::err_pragma_loop_invalid_keyword);
       return nullptr;
@@ -100,12 +115,12 @@ static Attr *handleLoopHintAttr(Sema &S, Stmt *St, const AttributeList &A,
   } else
     llvm_unreachable("Unknown loop hint option");
 
-  return LoopHintAttr::CreateImplicit(S.Context, Option, ValueInt,
+  return LoopHintAttr::CreateImplicit(S.Context, Spelling, Option, ValueInt,
                                       A.getRange());
 }
 
-static void
-CheckForIncompatibleAttributes(Sema &S, SmallVectorImpl<const Attr *> &Attrs) {
+static void CheckForIncompatibleAttributes(
+    Sema &S, const SmallVectorImpl<const Attr *> &Attrs) {
   // There are 3 categories of loop hints: vectorize, interleave, and
   // unroll. Each comes in two variants: an enable/disable form and a
   // form which takes a numeric argument. For example:
@@ -113,18 +128,9 @@ CheckForIncompatibleAttributes(Sema &S, SmallVectorImpl<const Attr *> &Attrs) {
   // accumulate the hints encountered while iterating through the
   // attributes to check for compatibility.
   struct {
-    int EnableOptionId;
-    int NumericOptionId;
-    bool EnabledIsSet;
-    bool ValueIsSet;
-    bool Enabled;
-    int Value;
-  } Options[] = {{LoopHintAttr::Vectorize, LoopHintAttr::VectorizeWidth, false,
-                  false, false, 0},
-                 {LoopHintAttr::Interleave, LoopHintAttr::InterleaveCount,
-                  false, false, false, 0},
-                 {LoopHintAttr::Unroll, LoopHintAttr::UnrollCount, false, false,
-                  false, 0}};
+    const LoopHintAttr *EnableAttr;
+    const LoopHintAttr *NumericAttr;
+  } HintAttrs[] = {{nullptr, nullptr}, {nullptr, nullptr}, {nullptr, nullptr}};
 
   for (const auto *I : Attrs) {
     const LoopHintAttr *LH = dyn_cast<LoopHintAttr>(I);
@@ -134,8 +140,6 @@ CheckForIncompatibleAttributes(Sema &S, SmallVectorImpl<const Attr *> &Attrs) {
       continue;
 
     int Option = LH->getOption();
-    int ValueInt = LH->getValue();
-
     int Category;
     switch (Option) {
     case LoopHintAttr::Vectorize:
@@ -152,44 +156,34 @@ CheckForIncompatibleAttributes(Sema &S, SmallVectorImpl<const Attr *> &Attrs) {
       break;
     };
 
-    auto &CategoryState = Options[Category];
-    SourceLocation ValueLoc = LH->getRange().getEnd();
+    auto &CategoryState = HintAttrs[Category];
+    SourceLocation OptionLoc = LH->getRange().getBegin();
+    const LoopHintAttr *PrevAttr;
     if (Option == LoopHintAttr::Vectorize ||
         Option == LoopHintAttr::Interleave || Option == LoopHintAttr::Unroll) {
       // Enable|disable hint.  For example, vectorize(enable).
-      if (CategoryState.EnabledIsSet) {
-        // Cannot specify enable/disable state twice.
-        S.Diag(ValueLoc, diag::err_pragma_loop_compatibility)
-            << /*Duplicate=*/true << LoopHintAttr::getOptionName(Option)
-            << LoopHintAttr::getValueName(CategoryState.Enabled)
-            << LoopHintAttr::getOptionName(Option)
-            << LoopHintAttr::getValueName(ValueInt);
-      }
-      CategoryState.EnabledIsSet = true;
-      CategoryState.Enabled = ValueInt;
+      PrevAttr = CategoryState.EnableAttr;
+      CategoryState.EnableAttr = LH;
     } else {
-      // Numeric hint.  For example, unroll_count(8).
-      if (CategoryState.ValueIsSet) {
-        // Cannot specify numeric hint twice.
-        S.Diag(ValueLoc, diag::err_pragma_loop_compatibility)
-            << /*Duplicate=*/true << LoopHintAttr::getOptionName(Option)
-            << CategoryState.Value << LoopHintAttr::getOptionName(Option)
-            << ValueInt;
-      }
-      CategoryState.ValueIsSet = true;
-      CategoryState.Value = ValueInt;
+      // Numeric hint.  For example, vectorize_width(8).
+      PrevAttr = CategoryState.NumericAttr;
+      CategoryState.NumericAttr = LH;
     }
 
-    if (CategoryState.EnabledIsSet && !CategoryState.Enabled &&
-        CategoryState.ValueIsSet) {
+    if (PrevAttr)
+      // Cannot specify same type of attribute twice.
+      S.Diag(OptionLoc, diag::err_pragma_loop_compatibility)
+          << /*Duplicate=*/true << PrevAttr->getDiagnosticName()
+          << LH->getDiagnosticName();
+
+    if (CategoryState.EnableAttr && !CategoryState.EnableAttr->getValue() &&
+        CategoryState.NumericAttr) {
       // Disable hints are not compatible with numeric hints of the
       // same category.
-      S.Diag(ValueLoc, diag::err_pragma_loop_compatibility)
+      S.Diag(OptionLoc, diag::err_pragma_loop_compatibility)
           << /*Duplicate=*/false
-          << LoopHintAttr::getOptionName(CategoryState.EnableOptionId)
-          << LoopHintAttr::getValueName(CategoryState.Enabled)
-          << LoopHintAttr::getOptionName(CategoryState.NumericOptionId)
-          << CategoryState.Value;
+          << CategoryState.EnableAttr->getDiagnosticName()
+          << CategoryState.NumericAttr->getDiagnosticName();
     }
   }
 }