[flang] Implement !DIR$ VECTOR ALWAYS

flang/docs/Directives.md

@@ -36,3 +36,62 @@ A list of non-standard directives supported by Flang
   and is limited to 256.
   [This directive is currently recognised by the parser, but not
   handled by the other parts of the compiler].
+* `!dir$ vector always` forces vectorization on the following loop regardless
+  of cost model decisions. The loop must still be vectorizable.
+  [This directive currently only works on plain do loops without labels].
+
+# Directive Details
+
+## Introduction
+Directives are commonly used in Fortran programs to specify additional actions 
+to be performed by the compiler. The directives are always specified with the 
+`!dir$` or `cdir$` prefix. 
+
+## Loop Directives
+Some directives are associated with the following construct, for example loop
+directives. Directives on loops are used to specify additional transformation to
+be performed by the compiler like enabling vectorisation, unrolling, interchange
+etc.
+
+Currently loop directives are not accepted in the presence of OpenMP or OpenACC
+constructs on the loop. This should be implemented as it is used in some
+applications.
+
+### Array Expressions
+It is to be decided whether loop directives should also be able to be associated
+with array expressions.
+
+## Semantics
+Directives that are associated with constructs must appear in the same section
+as the construct they are associated with, for example loop directives must
+appear in the executable section as the loops appear there. To facilitate this
+the parse tree is corrected to move such directives that appear in the
+specification part into the execution part.
+
+When a directive that must be associated with a construct appears, a search
+forward from that directive to the next non-directive construct is performed to
+check that that construct matches the expected construct for the directive.
+Skipping other intermediate directives allows multiple directives to appear on
+the same construct.
+
+## Lowering 
+Evaluation is extended with a new field called dirs for representing directives
+associated with that Evaluation. When lowering loop directives, the associated
+Do Loop's evaluation is found and the directive is added to it. This information
+is used only during the lowering of the loop.
+
+### Representation in LLVM
+The `llvm.loop` metadata is used in LLVM to provide information to the optimizer
+about the loop. For example, the `llvm.loop.vectorize.enable` metadata informs
+the optimizer that a loop can be vectorized without considering its cost-model.
+This attribute is added to the loop condition branch.
+
+### Representation in MLIR 
+The MLIR LLVM dialect models this by an attribute called LoopAnnotation
+Attribute. The attribute can be added to the latch of the loop in the cf
+dialect and is then carried through lowering to the LLVM dialect.
+
+## Testing
+Since directives must maintain a flow from source to LLVM IR, an integration
+test is provided that tests the `vector always` directive, as well as individual
+lit tests for each of the parsing, semantics and lowering stages.

flang/include/flang/Lower/PFTBuilder.h

@@ -347,6 +347,8 @@ struct Evaluation : EvaluationVariant {
   parser::CharBlock position{};
   std::optional<parser::Label> label{};
   std::unique_ptr<EvaluationList> evaluationList; // nested evaluations
+  // associated compiler directives
+  llvm::SmallVector<const parser::CompilerDirective *, 1> dirs;
   Evaluation *parentConstruct{nullptr};  // set for nodes below the top level
   Evaluation *lexicalSuccessor{nullptr}; // set for leaf nodes, some directives
   Evaluation *controlSuccessor{nullptr}; // set for some leaf nodes

flang/include/flang/Optimizer/Dialect/FIROps.h

@@ -16,6 +16,7 @@
 #include "flang/Optimizer/Dialect/FortranVariableInterface.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 

flang/include/flang/Optimizer/Dialect/FIROps.td

@@ -2160,7 +2160,8 @@ def fir_DoLoopOp : region_Op<"do_loop", [AttrSizedOperandSegments,
     Variadic<AnyType>:$initArgs,
     OptionalAttr<UnitAttr>:$unordered,
     OptionalAttr<UnitAttr>:$finalValue,
-    OptionalAttr<ArrayAttr>:$reduceAttrs
+    OptionalAttr<ArrayAttr>:$reduceAttrs,
+    OptionalAttr<LoopAnnotationAttr>:$loopAnnotation
   );
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);

flang/include/flang/Parser/dump-parse-tree.h

@@ -201,11 +201,12 @@ class ParseTreeDumper {
   NODE(parser, CommonStmt)
   NODE(CommonStmt, Block)
   NODE(parser, CompilerDirective)
+  NODE(CompilerDirective, AssumeAligned)
   NODE(CompilerDirective, IgnoreTKR)
   NODE(CompilerDirective, LoopCount)
-  NODE(CompilerDirective, AssumeAligned)
   NODE(CompilerDirective, NameValue)
   NODE(CompilerDirective, Unrecognized)
+  NODE(CompilerDirective, VectorAlways)
   NODE(parser, ComplexLiteralConstant)
   NODE(parser, ComplexPart)
   NODE(parser, ComponentArraySpec)

flang/include/flang/Parser/parse-tree.h

@@ -3334,14 +3334,15 @@ struct CompilerDirective {
     TUPLE_CLASS_BOILERPLATE(AssumeAligned);
     std::tuple<common::Indirection<Designator>, uint64_t> t;
   };
+  EMPTY_CLASS(VectorAlways);
   struct NameValue {
     TUPLE_CLASS_BOILERPLATE(NameValue);
     std::tuple<Name, std::optional<std::uint64_t>> t;
   };
   EMPTY_CLASS(Unrecognized);
   CharBlock source;
   std::variant<std::list<IgnoreTKR>, LoopCount, std::list<AssumeAligned>,
-      std::list<NameValue>, Unrecognized>
+      VectorAlways, std::list<NameValue>, Unrecognized>
       u;
 };
 

flang/lib/Lower/Bridge.cpp

@@ -1929,7 +1929,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 
     // Increment loop begin code. (Infinite/while code was already generated.)
     if (!infiniteLoop && !whileCondition)
-      genFIRIncrementLoopBegin(incrementLoopNestInfo);
+      genFIRIncrementLoopBegin(incrementLoopNestInfo, doStmtEval.dirs);
 
     // Loop body code.
     auto iter = eval.getNestedEvaluations().begin();
@@ -1974,8 +1974,20 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     return builder->createIntegerConstant(loc, controlType, 1); // step
   }
 
+  void addLoopAnnotationAttr(IncrementLoopInfo &info) {
+    mlir::BoolAttr f = mlir::BoolAttr::get(builder->getContext(), false);
+    mlir::LLVM::LoopVectorizeAttr va = mlir::LLVM::LoopVectorizeAttr::get(
+        builder->getContext(), /*disable=*/f, {}, {}, {}, {}, {}, {});
+    mlir::LLVM::LoopAnnotationAttr la = mlir::LLVM::LoopAnnotationAttr::get(
+        builder->getContext(), {}, /*vectorize=*/va, {}, {}, {}, {}, {}, {}, {},
+        {}, {}, {}, {}, {}, {});
+    info.doLoop.setLoopAnnotationAttr(la);
+  }
+
   /// Generate FIR to begin a structured or unstructured increment loop nest.
-  void genFIRIncrementLoopBegin(IncrementLoopNestInfo &incrementLoopNestInfo) {
+  void genFIRIncrementLoopBegin(
+      IncrementLoopNestInfo &incrementLoopNestInfo,
+      llvm::SmallVectorImpl<const Fortran::parser::CompilerDirective *> &dirs) {
     assert(!incrementLoopNestInfo.empty() && "empty loop nest");
     mlir::Location loc = toLocation();
     for (IncrementLoopInfo &info : incrementLoopNestInfo) {
@@ -2040,6 +2052,15 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         }
         if (info.hasLocalitySpecs())
           handleLocalitySpecs(info);
+
+        for (const auto *dir : dirs) {
+          std::visit(
+              Fortran::common::visitors{
+                  [&](const Fortran::parser::CompilerDirective::VectorAlways
+                          &d) { addLoopAnnotationAttr(info); },
+                  [&](const auto &) {}},
+              dir->u);
+        }
         continue;
       }
 
@@ -2573,8 +2594,28 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     }
   }
 
-  void genFIR(const Fortran::parser::CompilerDirective &) {
-    // TODO
+  void attachDirectiveToLoop(const Fortran::parser::CompilerDirective &dir,
+                             Fortran::lower::pft::Evaluation *e) {
+    while (e->isDirective())
+      e = e->lexicalSuccessor;
+
+    if (e->isA<Fortran::parser::NonLabelDoStmt>())
+      e->dirs.push_back(&dir);
+    else
+      fir::emitFatalError(toLocation(),
+                          "loop directive must appear before a loop");
+  }
+
+  void genFIR(const Fortran::parser::CompilerDirective &dir) {
+    Fortran::lower::pft::Evaluation &eval = getEval();
+
+    std::visit(
+        Fortran::common::visitors{
+            [&](const Fortran::parser::CompilerDirective::VectorAlways &) {
+              attachDirectiveToLoop(dir, &eval);
+            },
+            [&](const auto &) {}},
+        dir.u);
   }
 
   void genFIR(const Fortran::parser::OpenACCConstruct &acc) {

flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp

@@ -132,10 +132,14 @@ class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> {
     auto comparison = rewriter.create<mlir::arith::CmpIOp>(
         loc, arith::CmpIPredicate::sgt, itersLeft, zero);
 
-    rewriter.create<mlir::cf::CondBranchOp>(
+    auto cond = rewriter.create<mlir::cf::CondBranchOp>(
         loc, comparison, firstBlock, llvm::ArrayRef<mlir::Value>(), endBlock,
         llvm::ArrayRef<mlir::Value>());
 
+    // Copy loop annotations from the do loop to the loop entry condition.
+    if (auto ann = loop.getLoopAnnotation())
+      cond->setAttr("loop_annotation", *ann);
+
     // The result of the loop operation is the values of the condition block
     // arguments except the induction variable on the last iteration.
     auto args = loop.getFinalValue()

flang/lib/Parser/Fortran-parsers.cpp

@@ -1276,10 +1276,13 @@ constexpr auto loopCount{
 constexpr auto assumeAligned{"ASSUME_ALIGNED" >>
     optionalList(construct<CompilerDirective::AssumeAligned>(
         indirect(designator), ":"_tok >> digitString64))};
+constexpr auto vectorAlways{
+    "VECTOR ALWAYS" >> construct<CompilerDirective::VectorAlways>()};
 TYPE_PARSER(beginDirective >> "DIR$ "_tok >>
     sourced((construct<CompilerDirective>(ignore_tkr) ||
                 construct<CompilerDirective>(loopCount) ||
                 construct<CompilerDirective>(assumeAligned) ||
+                construct<CompilerDirective>(vectorAlways) ||
                 construct<CompilerDirective>(
                     many(construct<CompilerDirective::NameValue>(
                         name, maybe(("="_tok || ":"_tok) >> digitString64))))) /

flang/lib/Parser/unparse.cpp

@@ -1828,6 +1828,9 @@ class UnparseVisitor {
               Word("!DIR$ ASSUME_ALIGNED ");
               Walk(" ", assumeAligned, ", ");
             },
+            [&](const CompilerDirective::VectorAlways &valways) {
+              Word("!DIR$ VECTOR ALWAYS");
+            },
             [&](const std::list<CompilerDirective::NameValue> &names) {
               Walk("!DIR$ ", names, " ");
             },

flang/lib/Semantics/CMakeLists.txt

@@ -2,6 +2,7 @@ add_flang_library(FortranSemantics
   assignment.cpp
   attr.cpp
   canonicalize-acc.cpp
+  canonicalize-directives.cpp
   canonicalize-do.cpp
   canonicalize-omp.cpp
   check-acc-structure.cpp

flang/lib/Semantics/canonicalize-directives.cpp

@@ -0,0 +1,124 @@
+//===-- lib/Semantics/canonicalize-directives.cpp -------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "canonicalize-directives.h"
+#include "flang/Parser/parse-tree-visitor.h"
+
+namespace Fortran::semantics {
+
+using namespace parser::literals;
+
+// Check that directives are associated with the correct constructs.
+// Directives that need to be associated with other constructs in the execution
+// part are moved to the execution part so they can be checked there.
+class CanonicalizationOfDirectives {
+public:
+  CanonicalizationOfDirectives(parser::Messages &messages)
+      : messages_{messages} {}
+
+  template <typename T> bool Pre(T &) { return true; }
+  template <typename T> void Post(T &) {}
+
+  // Move directives that must appear in the Execution part out of the
+  // Specification part.
+  void Post(parser::SpecificationPart &spec);
+  bool Pre(parser::ExecutionPart &x);
+
+  // Ensure that directives associated with constructs appear accompanying the
+  // construct.
+  void Post(parser::Block &block);
+
+private:
+  // Ensure that loop directives appear immediately before a loop.
+  void CheckLoopDirective(parser::CompilerDirective &dir, parser::Block &block,
+      std::list<parser::ExecutionPartConstruct>::iterator it);
+
+  parser::Messages &messages_;
+
+  // Directives to be moved to the Execution part from the Specification part.
+  std::list<common::Indirection<parser::CompilerDirective>>
+      directivesToConvert_;
+};
+
+bool CanonicalizeDirectives(
+    parser::Messages &messages, parser::Program &program) {
+  CanonicalizationOfDirectives dirs{messages};
+  Walk(program, dirs);
+  return !messages.AnyFatalError();
+}
+
+static bool IsExecutionDirective(const parser::CompilerDirective &dir) {
+  return std::holds_alternative<parser::CompilerDirective::VectorAlways>(dir.u);
+}
+
+void CanonicalizationOfDirectives::Post(parser::SpecificationPart &spec) {
+  auto &list{
+      std::get<std::list<common::Indirection<parser::CompilerDirective>>>(
+          spec.t)};
+  for (auto it{list.begin()}; it != list.end();) {
+    if (IsExecutionDirective(it->value())) {
+      directivesToConvert_.emplace_back(std::move(*it));
+      it = list.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
+bool CanonicalizationOfDirectives::Pre(parser::ExecutionPart &x) {
+  auto origFirst{x.v.begin()};
+  for (auto &dir : directivesToConvert_) {
+    x.v.insert(origFirst,
+        parser::ExecutionPartConstruct{
+            parser::ExecutableConstruct{std::move(dir)}});
+  }
+
+  directivesToConvert_.clear();
+  return true;
+}
+
+template <typename T> T *GetConstructIf(parser::ExecutionPartConstruct &x) {
+  if (auto *y{std::get_if<parser::ExecutableConstruct>(&x.u)}) {
+    if (auto *z{std::get_if<common::Indirection<T>>(&y->u)}) {
+      return &z->value();
+    }
+  }
+  return nullptr;
+}
+
+void CanonicalizationOfDirectives::CheckLoopDirective(
+    parser::CompilerDirective &dir, parser::Block &block,
+    std::list<parser::ExecutionPartConstruct>::iterator it) {
+
+  // Skip over this and other compiler directives
+  while (GetConstructIf<parser::CompilerDirective>(*it)) {
+    ++it;
+  }
+
+  if (it == block.end() || !GetConstructIf<parser::DoConstruct>(*it)) {
+    std::string s{parser::ToUpperCaseLetters(dir.source.ToString())};
+    s.pop_back(); // Remove trailing newline from source string
+    messages_.Say(
+        dir.source, "A DO loop must follow the %s directive"_err_en_US, s);
+  }
+}
+
+void CanonicalizationOfDirectives::Post(parser::Block &block) {
+  for (auto it{block.begin()}; it != block.end(); ++it) {
+    if (auto *dir{GetConstructIf<parser::CompilerDirective>(*it)}) {
+      std::visit(
+          common::visitors{[&](parser::CompilerDirective::VectorAlways &) {
+                             CheckLoopDirective(*dir, block, it);
+                           },
+              [&](auto &) {}},
+          dir->u);
+    }
+  }
+}
+
+} // namespace Fortran::semantics

flang/lib/Semantics/canonicalize-directives.h

@@ -0,0 +1,22 @@
+//===-- lib/Semantics/canonicalize-directives.h -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_SEMANTICS_CANONICALIZE_DIRECTIVES_H_
+#define FORTRAN_SEMANTICS_CANONICALIZE_DIRECTIVES_H_
+
+namespace Fortran::parser {
+struct Program;
+class Messages;
+} // namespace Fortran::parser
+
+namespace Fortran::semantics {
+bool CanonicalizeDirectives(
+    parser::Messages &messages, parser::Program &program);
+}
+
+#endif // FORTRAN_SEMANTICS_CANONICALIZE_DIRECTIVES_H_