[flang][hlfir] Add hlfir.forall and its OrderAssignmentTreeOpInterface

[jeanPerier]

Summary

This patch adds the hlfir.forall operation and the
OrderAssignmentTreeOpInterface that allows representing Fortran forall.

It uses regions to keep Fortran expression evaluation independent from
each other in the IR. Forall assignments inside hlfir.forall are
represented with hlfir.region_assign which also keeps the IR generated
for each expressions independently.

The goal of this representation is to provide a representation that is
straightforward to generate from Fortran parse tree without any analysis, while
providing enough structure information so that an optimization pass can decide
how to schedule, and save if needed, the evaluations of the Forall and Where
expression and statements. It allows the data dependency analysis to be done at
the HLFIR level.

The OrderAssignmentTreeOpInterface allows ensuring that the Forall/Where
tree structure is kept in the IR. It will allow visiting this tree in
the IR without hard coding the operation structures in the pass.

[clementval]

LGTM

[tblah]

Looks great to me, thanks!

[vzakhari]

Looks great, Jean! I have just one question.

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

1069

Sounds okay, but would we have more readable IR if we support multiple indices with a single hlfir.forall? Having hlfir.yield return a vector of ints and having multiple blocks arguments in the body block does not sound too complex. Will it make the lowering more complex?

[jeanPerier]

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

1069

Fair point, I thought a bit about having hlfir.yield returning multiple values (for the bounds remapping arguments of pointer assignments, that I have not yet covered). While it is possible and would not make lowering much more complex, it would make the data dependency a bit more complex: right now, I am planning to have a dumb approach that looks at all the operation in a region that yields a value, and if any of them has a read effect on a variable being assigned, to create a temporary for the value yielded.

If we move to yields with several values, the analysis will need to be more clever to identify which yielded value is conflicting (it would need to follow the DAG an that is a more complex), or the pass will need to save all yielded values if there is any conflict in the region, just to be safe.

So I am rather inclined to have a heavier syntax until to have a conflict analysis as simple as possible to start with. Then, if we are confident we can have a "finer grain" analysis to detect which yielded entity conflicts, I would happy to move to a lighter syntax.

[vzakhari]

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

1069

Thank you for the explanation! I am not sure I understand the conflict detection 100%, so I am looking forward to reviewing it.

flang/include/flang/Optimizer/HLFIR/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS HLFIROpBase.td)
 mlir_tablegen(HLFIRTypes.h.inc -gen-typedef-decls -typedefs-dialect=hlfir)
 mlir_tablegen(HLFIRTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=hlfir)
 mlir_tablegen(HLFIRDialect.h.inc -gen-dialect-decls -dialect=hlfir)
 mlir_tablegen(HLFIRDialect.cpp.inc -gen-dialect-defs -dialect=hlfir)
 mlir_tablegen(HLFIRAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=hlfir)
 mlir_tablegen(HLFIRAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=hlfir)
 
 set(LLVM_TARGET_DEFINITIONS HLFIROps.td)
+mlir_tablegen(HLFIROpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(HLFIROpInterfaces.cpp.inc -gen-op-interface-defs)
 mlir_tablegen(HLFIROps.h.inc -gen-op-decls)
 mlir_tablegen(HLFIROps.cpp.inc -gen-op-defs)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name HLFIR)
 
 add_public_tablegen_target(HLFIROpsIncGen)

flang/include/flang/Optimizer/HLFIR/HLFIROps.h

 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/FortranVariableInterface.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "flang/Optimizer/HLFIR/HLFIROpInterfaces.h.inc"
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.h.inc"
 
 #endif // FORTRAN_OPTIMIZER_HLFIR_HLFIROPS_H

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

def hlfir_GetExtentOp : hlfir_Op<"get_extent", [Pure]> {
 
   let assemblyFormat = [{
     $shape attr-dict `:` functional-type(operands, results)
   }];
 
   let builders = [OpBuilder<(ins "mlir::Value":$shape, "unsigned":$dim)>];
 }
 
-def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> {
+def hlfir_OrderedAssignmentTreeOpInterface : OpInterface<"OrderedAssignmentTreeOpInterface"> {
+  let description = [{
+    Interface for the operations representing Forall and Where constructs and
+    statements as an mlir::Region tree.
+
+    These operations all have in common that they have "leaf" regions that contains
+    some code that should be evaluated for "all active combinations of Forall
+    index-name values" before the next OrderedAssignmentTreeOpInterface is
+    evaluated.
+
+    These operations are ordered in a tree fashion: Some operations, like
+    hlfir.forall or hlfir.where, contain a list of OrderedAssignmentTreeOpInterface
+    that should be evaluated after the "Leaf" regions, and before the next
+    OrderedAssignmentTreeOpInterface.
+
+    Nested OrderedAssignmentTreeOpInterface operations are affected by the
+    OrderedAssignmentTreeOpInterface operations that contain them (e.g:
+    hlfir.region_assign may be masked by the value of the mask region of
+    an hlfir.where that contains it).
+
+    OrderedAssignmentTreeOpInterface operations that contain nested operation
+    must return a "sub-tree" region that contains the list of nested
+    OrderedAssignmentTreeOpInterface operations.
+
+    There is no constraints over what IR a leaf region may contain. There is also
+    no restriction regarding how many leaf regions an
+    OrderedAssignmentTreeOpInterface operation may contain.
+
+    A "sub-tree" region, if any, must contain only OrderedAssignmentTreeOpInterface
+    operations and, maybe, a fir.end terminator.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/"Get the OrderedAssignmentTreeOpInterface leaf regions that contain evaluation code",
+      /*retTy=*/"void",
+      /*methodName=*/"getLeafRegions",
+      /*args=*/(ins "llvm::SmallVectorImpl<mlir::Region*>&":$regions),
+      /*methodBody=*/[{}]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Get the region, if any, containing the list of sub-tree OrderedAssignmentTreeOpInterface nodes",
+      /*retTy=*/"mlir::Region*",
+      /*methodName=*/"getSubTreeRegion",
+      /*args=*/(ins),
+      /*methodBody=*/[{}]
+    >,
+  ];
+
+  let extraClassDeclaration = [{
+    /// Interface verifier imlementation.
+    mlir::LogicalResult verifyImpl();
+  }];
+
+  let verify = [{
+    return ::mlir::cast<::hlfir::OrderedAssignmentTreeOpInterface>($_op).verifyImpl();
+  }];
+
+  let cppNamespace = "hlfir";
+}
+
+
+def hlfir_RegionAssignOp : hlfir_Op<"region_assign", [hlfir_OrderedAssignmentTreeOpInterface]> {
   let summary = "represent a Fortran assignment using regions for the LHS and RHS evaluation";
   let description = [{
     This operation can represent Forall and Where assignment when inside an
     hlfir.forall or hlfir.where "ordered assignment tree". It can
     also represent user defined assignments and assignment to vector
     subscripted entities without requiring the materialization of the
     right-hand side temporary copy that may be needed to implement Fortran
     assignment semantic.
def hlfir_RegionAssignOp : hlfir_Op<"region_assign", [hlfir_OrderedAssignmentTreeOpInterface]> {
 
   let extraClassDeclaration = [{
     mlir::Value getUserAssignmentRhs() {
       return getUserDefinedAssignment().getArguments()[0];
     }
     mlir::Value getUserAssignmentLhs() {
       return getUserDefinedAssignment().getArguments()[1];
     }
+    void getLeafRegions(llvm::SmallVectorImpl<mlir::Region*>& regions) {
+      regions.push_back(&getRhsRegion());
+      regions.push_back(&getLhsRegion());
+      if (!getUserDefinedAssignment().empty())
+        regions.push_back(&getUserDefinedAssignment());
+    }
+    mlir::Region* getSubTreeRegion() { return nullptr; }
+
   }];
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
-def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", "ElementalAddrOp"]>,
+def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp",
+    "ElementalAddrOp", "ForallOp"]>,
     SingleBlockImplicitTerminator<"fir::FirEndOp">]> {
 
   let summary = "Yield a value or variable inside a forall, where or region assignment";
 
   let description = [{
     Terminator operation that yields an HLFIR value or variable that was computed in
     a region and hold the yielded entity cleanup, if any, into its own region.
     This allows representing any Fortran expression evaluation in its own region so
let extraClassDeclaration = [{
     mlir::Region::BlockArgListType getIndices() {
       return getBody().getArguments();
     }
   }];
 
   let hasVerifier = 1;
 }
 
+/// Define ODS constraints to verify that a region ends with a yield of a
+/// certain type.
+def YieldIntegerOrEmpty : CPred<"yieldsIntegerOrEmpty($_self)">;
+def YieldIntegerRegion : RegionConstraint<
+  And<[SizedRegion<1>.predicate, YieldIntegerOrEmpty]>,
+  "single block region that yields an integer scalar value">;
+def MaybeYieldIntegerRegion : RegionConstraint<
+  And<[MaxSizedRegion<1>.predicate, YieldIntegerOrEmpty]>,
+  "optional single block region that yields an integer scalar value">;
+
+def hlfir_ForallOp : hlfir_Op<"forall", [hlfir_OrderedAssignmentTreeOpInterface]> {
+  let summary = "represent a Fortran forall";
+  let description = [{
+    This operation allows representing Fortran forall. It computes
+    a set of "index-name" values based on lower bound, upper bound,
+    and step values whose evaluations are represented in their own
+    regions.
+
+    Operations nested in its body region are evaluated in order.
+    As opposed to a regular loop, each nested operation is
+    fully evaluated for all the values in the "active set of
+    index-name" before the next nested operation. In practice, the
+    nested operation evaluation may be fused if it is proven that
+    they do not have data dependency.
+
+    The "index-name" value is represented as the argument of the
+    body region.
+
+    The lower, upper, and step region (if provided), must be terminated
+    by hlfir.yield that yields scalar integers.
+
+    The body region must only contain other OrderedAssignmentTreeOpInterface
+    operations (like hlfir.region_assign, or other hlfir.forall).
+
+    A Fortran forall with several indices is represented as a nest
+    of hlfir.forall.
+
+    Example: FORALL(I=1:10) X(I) = FOO(I)
+    ```
+      hlfir.forall lb {
+        hlfir.yield %c1 : index
+      } ub {
+        hlfir.yield %c10 : index
+      } (%i : index) {
+        hlfir.region_assign {
+          %res = fir.call @foo(%i) : (index) -> f32
+          hlfir.yield %res : f32
+        } to {
+          %xi = hlfir.designate %x(%i) : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+          hlfir.yield %xi : !fir.ref<f32>
+        }
+      }
+    ```
+
+  }];
+
+  let regions = (region  YieldIntegerRegion:$lb_region,
+                         YieldIntegerRegion:$ub_region,
+                         MaybeYieldIntegerRegion:$step_region,
+                         SizedRegion<1>:$body);
+
+  let extraClassDeclaration = [{
+    mlir::Value getForallIndexValue() {
+      return getBody().getArguments()[0];
+    }
+    void getLeafRegions(llvm::SmallVectorImpl<mlir::Region*>& regions) {
+      regions.push_back(&getLbRegion());
+      regions.push_back(&getUbRegion());
+      if (!getStepRegion().empty())
+        regions.push_back(&getStepRegion());
+    }
+    mlir::Region* getSubTreeRegion() { return &getBody(); }
+  }];
+
+  let assemblyFormat = [{
+    attr-dict `lb` $lb_region
+    `ub` $ub_region
+    (`step` $step_region^)?
+    custom<ForallOpBody>($body)
+  }];
+}
 
 #endif // FORTRAN_DIALECT_HLFIR_OPS

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

if (!hlfir::isFortranVariableType(elementAddrType) ||
           .isa<fir::SequenceType>())
     return emitOpError("body must compute the address of a scalar entity");
   unsigned shapeRank = getShape().getType().cast<fir::ShapeType>().getRank();
   if (shapeRank != getIndices().size())
     return emitOpError("body number of indices must match shape rank");
   return mlir::success();
 }
 
+//===----------------------------------------------------------------------===//
+// OrderedAssignmentTreeOpInterface
+//===----------------------------------------------------------------------===//
+
+mlir::LogicalResult hlfir::OrderedAssignmentTreeOpInterface::verifyImpl() {
+  if (mlir::Region *body = getSubTreeRegion())
+    if (!body->empty())
+      for (mlir::Operation &op : body->front())
+        if (!mlir::isa<hlfir::OrderedAssignmentTreeOpInterface, fir::FirEndOp>(
+                op))
+          return emitOpError(
+              "body region must only contain OrderedAssignmentTreeOpInterface "
+              "operations or fir.end");
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
+// ForallOp
+//===----------------------------------------------------------------------===//
+
+static mlir::ParseResult parseForallOpBody(mlir::OpAsmParser &parser,
+                                           mlir::Region &body) {
+  mlir::OpAsmParser::Argument bodyArg;
+  if (parser.parseLParen() || parser.parseArgument(bodyArg) ||
+      parser.parseColon() || parser.parseType(bodyArg.type) ||
+      parser.parseRParen())
+    return mlir::failure();
+  if (parser.parseRegion(body, {bodyArg}))
+    return mlir::failure();
+  ensureTerminator(body, parser.getBuilder(),
+                   parser.getBuilder().getUnknownLoc());
+  return mlir::success();
+}
+
+static void printForallOpBody(mlir::OpAsmPrinter &p, hlfir::ForallOp forall,
+                              mlir::Region &body) {
+  mlir::Value forallIndex = forall.getForallIndexValue();
+  p << " (" << forallIndex << ": " << forallIndex.getType() << ") ";
+  p.printRegion(body, /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/false);
+}
+
+/// Predicate implementation of YieldIntegerOrEmpty.
+static bool yieldsIntegerOrEmpty(mlir::Region &region) {
+  if (region.empty())
+    return true;
+  auto yield = mlir::dyn_cast_or_null<hlfir::YieldOp>(getTerminator(region));
+  return yield && fir::isa_integer(yield.getEntity().getType());
+}
+
+#include "flang/Optimizer/HLFIR/HLFIROpInterfaces.cpp.inc"
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

flang/test/HLFIR/forall.fir

+// Test hlfir.forall operation parse, verify (no errors), and unparse.
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+func.func @forall_test(%x: !fir.box<!fir.array<?xf32>>, %x2: !fir.box<!fir.array<?x?xf32>>) {
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  hlfir.forall lb {
+    hlfir.yield %c1 : index
+  } ub {
+    hlfir.yield %c10 : index
+  } (%i : index) {
+    hlfir.region_assign {
+      %res = fir.call @foo(%i) : (index) -> f32
+      hlfir.yield %res : f32
+    } to {
+      %xi = hlfir.designate %x(%i) : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+      hlfir.yield %xi : !fir.ref<f32>
+    }
+    hlfir.forall lb {
+      hlfir.yield %c1 : index
+    } ub {
+      hlfir.yield %c10 : index
+    } (%j : index) {
+      hlfir.region_assign {
+        %jf = fir.convert %j : (index) -> f32
+        hlfir.yield %jf : f32
+      } to {
+        %xij = hlfir.designate %x2(%i, %j) : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32>
+        hlfir.yield %xij : !fir.ref<f32>
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @forall_test(
+// CHECK-SAME:                           %[[VAL_0:.*]]: !fir.box<!fir.array<?xf32>>,
+// CHECK-SAME:                           %[[VAL_1:.*]]: !fir.box<!fir.array<?x?xf32>>) {
+// CHECK:           %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_3:.*]] = arith.constant 10 : index
+// CHECK:           hlfir.forall lb {
+// CHECK:             hlfir.yield %[[VAL_2]] : index
+// CHECK:           } ub {
+// CHECK:             hlfir.yield %[[VAL_3]] : index
+// CHECK:           }  (%[[VAL_4:.*]]: index) {
+// CHECK:             hlfir.region_assign {
+// CHECK:               %[[VAL_5:.*]] = fir.call @foo(%[[VAL_4]]) : (index) -> f32
+// CHECK:               hlfir.yield %[[VAL_5]] : f32
+// CHECK:             } to {
+// CHECK:               %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_4]])  : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+// CHECK:               hlfir.yield %[[VAL_6]] : !fir.ref<f32>
+// CHECK:             }
+// CHECK:             hlfir.forall lb {
+// CHECK:               hlfir.yield %[[VAL_2]] : index
+// CHECK:             } ub {
+// CHECK:               hlfir.yield %[[VAL_3]] : index
+// CHECK:             }  (%[[VAL_7:.*]]: index) {
+// CHECK:               hlfir.region_assign {
+// CHECK:                 %[[VAL_8:.*]] = fir.convert %[[VAL_7]] : (index) -> f32
+// CHECK:                 hlfir.yield %[[VAL_8]] : f32
+// CHECK:               } to {
+// CHECK:                 %[[VAL_9:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_4]], %[[VAL_7]])  : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32>
+// CHECK:                 hlfir.yield %[[VAL_9]] : !fir.ref<f32>
+// CHECK:               }
+// CHECK:             }
+// CHECK:           }
+
+func.func @forall_test_step(%x : !fir.box<!fir.array<10xf32>>, %y: !fir.box<!fir.array<?xf32>>) {
+  hlfir.forall lb {
+    %c1 = arith.constant 1 : index
+    hlfir.yield %c1 : index
+  } ub {
+    %c10 = arith.constant 10 : index
+    hlfir.yield %c10 : index
+  } step {
+    %c2 = arith.constant 2 : index
+    hlfir.yield %c2 : index
+  } (%i : index) {
+    hlfir.region_assign {
+      %yi = hlfir.designate %y(%i) : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+      %val = fir.load %yi : !fir.ref<f32>
+      hlfir.yield %val : f32
+    } to {
+      %xi = hlfir.designate %x(%i) : (!fir.box<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+      hlfir.yield %xi : !fir.ref<f32>
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @forall_test_step(
+// CHECK-SAME:                                %[[VAL_0:.*]]: !fir.box<!fir.array<10xf32>>,
+// CHECK-SAME:                                %[[VAL_1:.*]]: !fir.box<!fir.array<?xf32>>) {
+// CHECK:           hlfir.forall lb {
+// CHECK:             %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK:             hlfir.yield %[[VAL_2]] : index
+// CHECK:           } ub {
+// CHECK:             %[[VAL_3:.*]] = arith.constant 10 : index
+// CHECK:             hlfir.yield %[[VAL_3]] : index
+// CHECK:           } step {
+// CHECK:             %[[VAL_4:.*]] = arith.constant 2 : index
+// CHECK:             hlfir.yield %[[VAL_4]] : index
+// CHECK:           }  (%[[VAL_5:.*]]: index) {
+// CHECK:             hlfir.region_assign {
+// CHECK:               %[[VAL_6:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_5]])  : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+// CHECK:               %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref<f32>
+// CHECK:               hlfir.yield %[[VAL_7]] : f32
+// CHECK:             } to {
+// CHECK:               %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]])  : (!fir.box<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+// CHECK:               hlfir.yield %[[VAL_8]] : !fir.ref<f32>
+// CHECK:             }
+// CHECK:           }

flang/test/HLFIR/invalid.fir

hlfir.region_assign {
     hlfir.elemental_addr %vector_shape  : !fir.shape<1> {
       ^bb0(%i: index, %j: index):
        %elt = hlfir.designate %y(%i, %j) : (!fir.ref<!fir.array<20x20xf32>>, index, index) -> !fir.ref<f32>
        hlfir.yield %elt : !fir.ref<f32>
     }
   }
   return
 }
+
+// -----
+func.func @bad_forall(%x : !fir.box<!fir.array<10xf32>>, %y: f32, %bad : !fir.ref<!fir.array<10xindex>>) {
+  // expected-error@+1 {{'hlfir.forall' op region #0 ('lb_region') failed to verify constraint: single block region that yields an integer scalar value}}
+  hlfir.forall lb {
+    hlfir.yield %bad : !fir.ref<!fir.array<10xindex>>
+  } ub {
+    %c10 = arith.constant 10 : index
+    hlfir.yield %c10 : index
+  } (%i : index) {
+    hlfir.region_assign {
+      hlfir.yield %y : f32
+    } to {
+      %xi = hlfir.designate %x(%i) : (!fir.box<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+      hlfir.yield %xi : !fir.ref<f32>
+    }
+  }
+  return
+}
+
+// -----
+func.func @bad_forall_2(%x : !fir.box<!fir.array<10xf32>>, %y: f32) {
+  // expected-error@+1 {{'hlfir.forall' op body region must only contain OrderedAssignmentTreeOpInterface operations or fir.end}}
+  hlfir.forall lb {
+    %c1 = arith.constant 1 : index
+    hlfir.yield %c1 : index
+  } ub {
+    %c10 = arith.constant 10 : index
+    hlfir.yield %c10 : index
+  } (%i : index) {
+    %xi = hlfir.designate %x(%i) : (!fir.box<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+    hlfir.assign %y to %xi : f32, !fir.ref<f32>
+  }
+  return
+}