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

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.

Differential Revision: https://reviews.llvm.org/D149734

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

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

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

@@ -20,6 +20,7 @@
 #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"
 

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

@@ -814,7 +814,69 @@ def hlfir_GetExtentOp : hlfir_Op<"get_extent", [Pure]> {
   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
@@ -862,13 +924,22 @@ def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> {
     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";
@@ -961,5 +1032,86 @@ def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"R
   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

@@ -1101,5 +1101,56 @@ mlir::LogicalResult hlfir::ElementalAddrOp::verify() {
   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

@@ -0,0 +1,110 @@
+// 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:           }