[flang-commits] [flang] 64b591a - [flang][hlfir] Add hlfir.region_assign and its hlfir.yield terminator

Jean Perier via flang-commits flang-commits at lists.llvm.org
Wed May 3 00:12:51 PDT 2023


Author: Jean Perier
Date: 2023-05-03T09:08:48+02:00
New Revision: 64b591a89ceab9636d5dbc27740a3620d5d64a12

URL: https://github.com/llvm/llvm-project/commit/64b591a89ceab9636d5dbc27740a3620d5d64a12
DIFF: https://github.com/llvm/llvm-project/commit/64b591a89ceab9636d5dbc27740a3620d5d64a12.diff

LOG: [flang][hlfir] Add hlfir.region_assign and its hlfir.yield terminator

hlfir.region_assign is a Region based version of hlfir.assign: the
right-hand side and left-hand-side are evaluated in their own region,
and an optional region can be added to implement user defined
assignment.

This will be used for:
 - assignments inside where and forall
 - user defined assignments
 - assignments to vector subscripted entities.

Rational:

Forall and Where lowering requires solving an expression/assignment
evaluation scheduling problem based on data dependencies between the
variables being assigned and the one used in the expressions.
Keeping left-hand side and right-hand side in their own region will
make it really easy to analyse the dependency and move around the
expression evaluation as a whole. Operation DAGs are hard to scissor out
when the LHS and RHS evaluation are lowered in the same block. The pass
dealing with further forall/where lowering in HLFIR will need to
succeed. It is not acceptable for them to fail splitting the RHS/LHS
evaluation code. Keeping them in independent block is an approach that
cannot fail.

For user defined assignments, having a region allows implementing all
the call details in lowering, and even to allow inlining of the user
assignment, before it is decided if a temporary for the LHS or RHS is
required or not.

The operation description mention "hlfir.elemental_addr" (operation that
will be used for vector subscripted LHS) and "ordered assignment trees"
(concept/inetrface that will be used to represent forall/where structure
in HLFIR). These will be pushed in follow-up patch, but I do not want t
scissor out the descriptions.

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

Added: 
    flang/test/HLFIR/region-assign.fir

Modified: 
    flang/include/flang/Optimizer/HLFIR/HLFIROps.h
    flang/include/flang/Optimizer/HLFIR/HLFIROps.td
    flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
    flang/test/HLFIR/invalid.fir

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.h b/flang/include/flang/Optimizer/HLFIR/HLFIROps.h
index 88f50fd110db5..fe82b691bc526 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.h
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.h
@@ -11,6 +11,7 @@
 
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/FortranVariableInterface.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"

diff  --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index 79bba3dc39e2b..495d72863db56 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
@@ -814,4 +814,88 @@ def hlfir_GetExtentOp : hlfir_Op<"get_extent", [Pure]> {
   let builders = [OpBuilder<(ins "mlir::Value":$shape, "unsigned":$dim)>];
 }
 
+def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> {
+  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.
+
+    The right-hand side and left-hand side evaluations are held in their
+    own regions terminated with hlfir.yield operations (or hlfir.elemental_addr
+    for a left-hand side with vector subscript).
+
+    An optional region may be added to implement user defined assignment.
+    This region provides two block arguments with the same type as the
+    yielded rhs and lhs entities (in that order), or the element type if this
+    is an elemental user defined assignment.
+
+    If this optional region is not provided, intrinsic assignment is performed.
+
+    Example: "X = Y",  where "=" is a user defined elemental assignment "foo"
+    taking Y by value.
+    ```
+    hlfir.region_assign {
+      hlfir.yield %y : !fir.box<!fir.array<?x!f32>>
+    } to {
+      hlfir.yield %x : !fir.box<!fir.array<?x!fir.type<t>>>
+    } user_defined_assignment (%rhs_elt: !fir.ref<f32>) to (%lhs_elt: !fir.ref<!fir.type<t>>) {
+      %0 = fir.load %rhs_elt : !fir.ref<f32>
+      fir.call @foo(%lhs_elt, %0) : (!fir.ref<!fir.type<t>>, f32) -> ()
+    }
+    ```
+
+    TODO: add optional "realloc" semantics like for hlfir.assign.
+  }];
+
+  let regions = (region  SizedRegion<1>:$rhs_region,
+                         SizedRegion<1>:$lhs_region,
+                         MaxSizedRegion<1>:$user_defined_assignment);
+
+  let extraClassDeclaration = [{
+    mlir::Value getUserAssignmentRhs() {
+      return getUserDefinedAssignment().getArguments()[0];
+    }
+    mlir::Value getUserAssignmentLhs() {
+      return getUserDefinedAssignment().getArguments()[1];
+    }
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
+def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp"]>,
+    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
+    that the evaluation can easily be scheduled/moved around in a pass.
+
+    Example: "foo(x)" where foo returns an allocatable array.
+    ```
+    {
+      // In some region.
+      %0 = fir.call @foo(x) (!fir.ref<f32>) -> !fir.box<fir.heap<!fir.array<?xf32>>>
+      hlfir.yield %0 : !fir.box<!fir.heap<!fir.array<?xf32>>> cleanup {
+        %1 = fir.box_addr %0 : !fir.box<!fir.heap<!fir.array<?xf32>>> -> !fir.heap<!fir.array<?xf32>>
+        %fir.freemem %1 : !fir.heap<!fir.array<?xf32>>
+      }
+    }
+    ```
+  }];
+
+  let arguments = (ins AnyFortranEntity:$entity);
+  let regions = (region  MaxSizedRegion<1>:$cleanup);
+
+  let assemblyFormat = "$entity attr-dict `:` type($entity) custom<YieldOpCleanup>($cleanup)";
+}
+
 #endif // FORTRAN_DIALECT_HLFIR_OPS

diff  --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index a479d64f39db8..51d18def03478 100644
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -960,5 +960,111 @@ mlir::LogicalResult hlfir::GetExtentOp::verify() {
   return mlir::success();
 }
 
+//===----------------------------------------------------------------------===//
+// RegionAssignOp
+//===----------------------------------------------------------------------===//
+
+/// Add a fir.end terminator to a parsed region if it does not already has a
+/// terminator.
+static void ensureTerminator(mlir::Region &region, mlir::Builder &builder,
+                             mlir::Location loc) {
+  // Borrow YielOp::ensureTerminator MLIR generated implementation to add a
+  // fir.end if there is no terminator. This has nothing to do with YielOp,
+  // other than the fact that yieldOp has the
+  // SingleBlocklicitTerminator<"fir::FirEndOp"> interface that
+  // cannot be added on other HLFIR operations with several regions which are
+  // not all terminated the same way.
+  hlfir::YieldOp::ensureTerminator(region, builder, loc);
+}
+
+mlir::ParseResult hlfir::RegionAssignOp::parse(mlir::OpAsmParser &parser,
+                                               mlir::OperationState &result) {
+  mlir::Region &rhsRegion = *result.addRegion();
+  if (parser.parseRegion(rhsRegion))
+    return mlir::failure();
+  mlir::Region &lhsRegion = *result.addRegion();
+  if (parser.parseKeyword("to") || parser.parseRegion(lhsRegion))
+    return mlir::failure();
+  mlir::Region &userDefinedAssignmentRegion = *result.addRegion();
+  if (succeeded(parser.parseOptionalKeyword("user_defined_assign"))) {
+    mlir::OpAsmParser::Argument rhsArg, lhsArg;
+    if (parser.parseLParen() || parser.parseArgument(rhsArg) ||
+        parser.parseColon() || parser.parseType(rhsArg.type) ||
+        parser.parseRParen() || parser.parseKeyword("to") ||
+        parser.parseLParen() || parser.parseArgument(lhsArg) ||
+        parser.parseColon() || parser.parseType(lhsArg.type) ||
+        parser.parseRParen())
+      return mlir::failure();
+    if (parser.parseRegion(userDefinedAssignmentRegion, {rhsArg, lhsArg}))
+      return mlir::failure();
+    ensureTerminator(userDefinedAssignmentRegion, parser.getBuilder(),
+                     result.location);
+  }
+  return mlir::success();
+}
+
+void hlfir::RegionAssignOp::print(mlir::OpAsmPrinter &p) {
+  p << " ";
+  p.printRegion(getRhsRegion(), /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/true);
+  p << " to ";
+  p.printRegion(getLhsRegion(), /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/true);
+  if (!getUserDefinedAssignment().empty()) {
+    p << " user_defined_assign ";
+    mlir::Value userAssignmentRhs = getUserAssignmentRhs();
+    mlir::Value userAssignmentLhs = getUserAssignmentLhs();
+    p << " (" << userAssignmentRhs << ": " << userAssignmentRhs.getType()
+      << ") to (";
+    p << userAssignmentLhs << ": " << userAssignmentLhs.getType() << ") ";
+    p.printRegion(getUserDefinedAssignment(), /*printEntryBlockArgs=*/false,
+                  /*printBlockTerminators=*/false);
+  }
+}
+
+static mlir::Operation *getTerminator(mlir::Region &region) {
+  if (region.empty() || region.back().empty())
+    return nullptr;
+  return &region.back().back();
+}
+
+mlir::LogicalResult hlfir::RegionAssignOp::verify() {
+  if (!mlir::isa_and_nonnull<hlfir::YieldOp>(getTerminator(getRhsRegion())))
+    return emitOpError(
+        "right-hand side region must be terminated by an hlfir.yield");
+  // TODO: allow hlfir.elemental_addr.
+  if (!mlir::isa_and_nonnull<hlfir::YieldOp>(getTerminator(getLhsRegion())))
+    return emitOpError("left-hand side region must be terminated by an "
+                       "hlfir.yield or hlfir.elemental_addr");
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+static mlir::ParseResult parseYieldOpCleanup(mlir::OpAsmParser &parser,
+                                             mlir::Region &cleanup) {
+  if (succeeded(parser.parseOptionalKeyword("cleanup"))) {
+    if (parser.parseRegion(cleanup, /*arguments=*/{},
+                           /*argTypes=*/{}))
+      return mlir::failure();
+    hlfir::YieldOp::ensureTerminator(cleanup, parser.getBuilder(),
+                                     parser.getBuilder().getUnknownLoc());
+  }
+
+  return mlir::success();
+}
+
+template <typename YieldOp>
+static void printYieldOpCleanup(mlir::OpAsmPrinter &p, YieldOp yieldOp,
+                                mlir::Region &cleanup) {
+  if (!cleanup.empty()) {
+    p << "cleanup ";
+    p.printRegion(cleanup, /*printEntryBlockArgs=*/false,
+                  /*printBlockTerminators=*/false);
+  }
+}
+
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

diff  --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir
index 9b80cbdc25c4f..a1e532563b4da 100644
--- a/flang/test/HLFIR/invalid.fir
+++ b/flang/test/HLFIR/invalid.fir
@@ -518,3 +518,26 @@ func.func @bad_getextent(%arg0: !fir.shape<1>) {
   // expected-error at +1 {{'hlfir.get_extent' op dimension index out of bounds}}
   %0 = hlfir.get_extent %arg0 {dim = 1 : index} : (!fir.shape<1>) -> index
 }
+
+// -----
+func.func @bad_region_assign_1(%x: !fir.box<!fir.array<?xf32>>) {
+// expected-error at +1 {{'hlfir.region_assign' op right-hand side region must be terminated by an hlfir.yield}}
+  hlfir.region_assign {
+    %c100 = arith.constant 100 : index
+  } to {
+    hlfir.yield %x : !fir.box<!fir.array<?xf32>>
+  }
+  return
+}
+
+// -----
+func.func @bad_region_assign_2(%x: !fir.box<!fir.array<?xf32>>) {
+// expected-error at +1 {{'hlfir.region_assign' op left-hand side region must be terminated by an hlfir.yield or hlfir.elemental_addr}}
+  hlfir.region_assign {
+    hlfir.yield %x : !fir.box<!fir.array<?xf32>>
+  } to {
+    %c100 = arith.constant 100 : index
+  } user_defined_assign  (%rhs: !fir.ref<i64>) to (%lhs: !fir.ref<f32>) {
+  }
+  return
+}

diff  --git a/flang/test/HLFIR/region-assign.fir b/flang/test/HLFIR/region-assign.fir
new file mode 100644
index 0000000000000..39724d818cb02
--- /dev/null
+++ b/flang/test/HLFIR/region-assign.fir
@@ -0,0 +1,87 @@
+// Test hlfir.region_assign and hlfir.yield operation parse, verify (no errors),
+// and unparse.
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+func.func @region_assign_test(%y : !fir.box<!fir.array<?xi64>>, %x: !fir.box<!fir.array<?xf32>>) {
+  %c100 = arith.constant 100 : index
+  %shape = fir.shape %c100 : (index) -> !fir.shape<1>
+  hlfir.region_assign {
+    %expr = hlfir.elemental %shape : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
+    ^bb0(%i : index):
+      %yelt = hlfir.designate %y(%i) : (!fir.box<!fir.array<?xi64>>, index) -> !fir.ref<i64>
+      %elt = fir.call @some_elemental(%yelt) : (!fir.ref<i64>) -> f32
+      hlfir.yield_element %elt : f32
+    }
+    hlfir.yield %expr : !hlfir.expr<?xf32>
+  } to {
+    hlfir.yield %x : !fir.box<!fir.array<?xf32>>
+  }
+  return
+}
+// CHECK-LABEL:   func.func @region_assign_test(
+// CHECK-SAME:                                  %[[VAL_0:.*]]: !fir.box<!fir.array<?xi64>>,
+// CHECK-SAME:                                  %[[VAL_1:.*]]: !fir.box<!fir.array<?xf32>>) {
+// CHECK:           %[[VAL_2:.*]] = arith.constant 100 : index
+// CHECK:           %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK:           hlfir.region_assign {
+// CHECK:             %[[VAL_4:.*]] = hlfir.elemental %[[VAL_3]] : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
+// CHECK:             ^bb0(%[[VAL_5:.*]]: index):
+// CHECK:               %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]])  : (!fir.box<!fir.array<?xi64>>, index) -> !fir.ref<i64>
+// CHECK:               %[[VAL_7:.*]] = fir.call @some_elemental(%[[VAL_6]]) : (!fir.ref<i64>) -> f32
+// CHECK:               hlfir.yield_element %[[VAL_7]] : f32
+// CHECK:             }
+// CHECK:             hlfir.yield %[[VAL_8:.*]] : !hlfir.expr<?xf32>
+// CHECK:           } to {
+// CHECK:             hlfir.yield %[[VAL_1]] : !fir.box<!fir.array<?xf32>>
+// CHECK:           }
+
+func.func @region_user_assign_test(%y : !fir.box<!fir.array<?xi64>>, %x: !fir.box<!fir.array<?xf32>>) {
+  hlfir.region_assign {
+    hlfir.yield %y : !fir.box<!fir.array<?xi64>>
+  } to {
+    hlfir.yield %x : !fir.box<!fir.array<?xf32>>
+  } user_defined_assign (%rhs : !fir.ref<i64>) to (%lhs : !fir.ref<f32>) {
+    %0 = fir.load %rhs : !fir.ref<i64>
+    fir.call @user_assign(%lhs, %0) : (!fir.ref<f32>, i64) -> ()
+  }
+  return
+}
+// CHECK-LABEL:   func.func @region_user_assign_test(
+// CHECK-SAME:                                       %[[VAL_0:.*]]: !fir.box<!fir.array<?xi64>>,
+// CHECK-SAME:                                       %[[VAL_1:.*]]: !fir.box<!fir.array<?xf32>>) {
+// CHECK:           hlfir.region_assign {
+// CHECK:             hlfir.yield %[[VAL_0]] : !fir.box<!fir.array<?xi64>>
+// CHECK:           } to {
+// CHECK:             hlfir.yield %[[VAL_1]] : !fir.box<!fir.array<?xf32>>
+// CHECK:           } user_defined_assign  (%[[VAL_2:.*]]: !fir.ref<i64>) to (%[[VAL_3:.*]]: !fir.ref<f32>) {
+// CHECK:             %[[VAL_4:.*]] = fir.load %[[VAL_2]] : !fir.ref<i64>
+// CHECK:             fir.call @user_assign(%[[VAL_3]], %[[VAL_4]]) : (!fir.ref<f32>, i64) -> ()
+// CHECK:           }
+
+func.func @yield_cleanup(%x: !fir.box<!fir.array<?xf32>>) {
+  hlfir.region_assign {
+    %0 = fir.allocmem !fir.array<100xf32>
+    fir.call @fillin_some_values(%0) : (!fir.heap<!fir.array<100xf32>>) -> ()
+    hlfir.yield %0 : !fir.heap<!fir.array<100xf32>> cleanup {
+      fir.freemem %0 : !fir.heap<!fir.array<100xf32>>
+    }
+  } to {
+    hlfir.yield %x : !fir.box<!fir.array<?xf32>>
+  }
+  return
+}
+// CHECK-LABEL:   func.func @yield_cleanup(
+// CHECK-SAME:                             %[[VAL_0:.*]]: !fir.box<!fir.array<?xf32>>) {
+// CHECK:           hlfir.region_assign {
+// CHECK:             %[[VAL_1:.*]] = fir.allocmem !fir.array<100xf32>
+// CHECK:             fir.call @fillin_some_values(%[[VAL_1]]) : (!fir.heap<!fir.array<100xf32>>) -> ()
+// CHECK:             hlfir.yield %[[VAL_1]] : !fir.heap<!fir.array<100xf32>> cleanup {
+// CHECK:               fir.freemem %[[VAL_1]] : !fir.heap<!fir.array<100xf32>>
+// CHECK:             }
+// CHECK:           } to {
+// CHECK:             hlfir.yield %[[VAL_0]] : !fir.box<!fir.array<?xf32>>
+// CHECK:           }
+
+func.func private @user_assign(!fir.ref<f32>, i64) -> ()
+func.func private @some_elemental(!fir.ref<i64>) -> f32
+func.func private @fillin_some_values(!fir.heap<!fir.array<100xf32>>) -> ()


        


More information about the flang-commits mailing list