[flang] Implement !DIR$ VECTOR ALWAYS (#93830)

This patch implements support for the VECTOR ALWAYS directive, which
forces
vectorization to occurr when possible regardless of a decision by the
cost
model. This is done by adding an attribute to the branch into the loop
in LLVM
to indicate that the loop should always be vectorized.

This patch only implements this directive on plan structured do loops 
without labels. Support for unstructured loops and array
expressions is planned for future patches.

Author: DavidTruby

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

@@ -350,6 +350,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

@@ -3337,14 +3337,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

@@ -1935,7 +1935,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();
@@ -1980,8 +1980,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) {
@@ -2046,6 +2058,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;
       }
 
@@ -2579,8 +2600,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

@@ -1277,10 +1277,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

@@ -1834,6 +1834,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_