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

[jeanPerier]

Summary

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.

[clementval]

This looks nice! Just couple of remark/questions but nothing blocking for me.

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

845

Maybe it's just me but I find it a bit counter intuitive to have the rhs on the left and the lhs on the right.

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

980

Was it too tricky to do it with the inlined assembly format?

[jeanPerier]

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

845

Oups, I did not update the format here... This is actually: (%rhs_elt : !fir.ref<f32>) to (%lhs_elt:!fir.ref<!fir.type<t>>), I moved to custom parser to get the to and make it clear what is the RHS vs LHS. I followed fir.store/hlfir.assign/... logic of "RHS" to "LHS".

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

980

The I was not sure how to get the () to () part of the user assignment block arguments with assembly format. I thought is was nicer to be heavy here to disambiguate the RHS/LHS.

Do you know if one can manipulate region arguments in inlined assembly format?

I also failed to get the terminator to be added on the the user defined region, while the other region are not implicitly terminated, but I did not try hard here.

[clementval]

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

845

Ok this makes sense with the new format then.

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

980

I don't think you can manipulate region argument with the inlined assembly format. That's a good reason to make it custom. Thanks for the explanation.

[tblah]

Great work on this. I like the overall idea. Other than some nits, this looks good to me.

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

838–844

I think maybe something like

hlfir.region_assign {

hlfir.yield %x ...

} = {

hlfir.yield %y ...

}

Would be easier. I know there is already precedent (fir.store) for "RHS to LHS" but when I read "assign" I will be already thinking "LHS = RHS".

This is a matter of opinion so feel free to continue without changing if you prefer "RHS to LHS".

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

1046

I think this and printYieldOpCleanup are unused. Are they for something added in a later patch?

[jeanPerier]

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

838–844

I do not have a strong opination either. If it was not for the fir.store / hlfir.assign / llvm store precedent, I would probably have gone LHS "something" RHS. I am not sure why LLVM store chose to put the value as its first operand.

I tried your suggestion, and unfortunately it seems "=" is a reserved keyword in MLIR parser and cannot be used here. I get the funny error:

`
region_assign.fir:15:5: error: custom op 'hlfir.region_assign' expected '='
  } = {
    ^

I find "equal" too predicate like, maybe "with" could do, also I find "assign y to x" a bit more natural than "assign x with y" (but I am not a native English speaker).

Do you see good short word I could use instead of "="?

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

1046

They are used in code generated by MLIR because the .td inline assembly format for YieldOp uses custom<YieldOpCleanup>($cleanup).

See the custom parsing directives https://mlir.llvm.org/docs/DefiningDialects/Operations/#custom-directives

I try to use this when possible instead of defining a complete cpp print/parse implementation.

[tblah]

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

838–844

How about hlfir.region_assign LHS from RHS?

I don't feel strongly about this so feel free to leave it how it is. If we can't use =, I can't think of any very good options.

[jeanPerier]

Rebase and update example to reflect the correct assembly format
for user assignment region block arguments.

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

 //===-- HLFIROps.h - FIR operations -----------------------------*- 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_OPTIMIZER_HLFIR_HLFIROPS_H
 #define FORTRAN_OPTIMIZER_HLFIR_HLFIROPS_H
 
 #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"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #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", []> {
+  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

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

mlir::LogicalResult hlfir::GetExtentOp::verify() {
   fir::ShapeType shapeTy = getShape().getType().cast<fir::ShapeType>();
   std::uint64_t rank = shapeTy.getRank();
   llvm::APInt dim = getDim();
   if (dim.sge(rank))
     return emitOpError("dimension index out of bounds");
   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"

flang/test/HLFIR/invalid.fir

func.func @bad_shapeof2(%arg0: !hlfir.expr<10xi32>) {
   %0 = hlfir.shape_of %arg0 : (!hlfir.expr<10xi32>) -> !fir.shape<42>
 }
 
 // -----
 func.func @bad_getextent(%arg0: !fir.shape<1>) {
   // expected-error@+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@+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@+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
+}

flang/test/HLFIR/region-assign.fir

+// 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>>) -> ()