[flang][hlfir] Add hlfir.elemental_addr for vector subscripted assignment

[jeanPerier]

Summary

See the operation description in HLFIROps.td.

Depends on D149442

[clementval]

LGTM

[vzakhari]

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

1034–1035

TODO is probably redundant now.

[tblah]

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

903

Please could you mention the cleanup region in the description

[jeanPerier]

• Rebase
• Remove obsolete TODO comment
• Document cleanup region

[jeanPerier]

Thanks for the reviews!

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

903

Good point, done.

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

1034–1035

Indeed, thanks.

[tblah]

Looks good to me, thanks for this!

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

mlir::Value getUserAssignmentLhs() {
       return getUserDefinedAssignment().getArguments()[1];
     }
   }];
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
-def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp"]>,
+def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", "ElementalAddrOp"]>,
     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
def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", "ElementalAddrOp"]>,
   }];
 
   let arguments = (ins AnyFortranEntity:$entity);
   let regions = (region  MaxSizedRegion<1>:$cleanup);
 
   let assemblyFormat = "$entity attr-dict `:` type($entity) custom<YieldOpCleanup>($cleanup)";
 }
 
+def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"RegionAssignOp">, RecursiveMemoryEffects, RecursivelySpeculatable]> {
+  let summary = "Yield the address of a vector subscripted variable inside an hlfir.region_assign";
+  let description = [{
+    Special terminator node for the left-hand side region of an hlfir.region_assign
+    to a vector subscripted entity.
+
+    It represents how the address of an element of such entity is computed given
+    one based indices.
+
+    It is very similar to hlfir.elemental, except that it does not produce an SSA
+    value because there is no hlfir type to describe a vector subscripted entity
+    (the codegen of such type would be problematic). Hence, it is tightly linked
+    to an hlfir.region_assign by its terminator property.
+
+    An optional cleanup region may be provided if any of the subscript expressions
+    of the designator require a cleanup.
+    This allows documenting cleanups that cannot be generated after the vector
+    subscripted designator usage (that has not been materizaled yet). The cleanups
+    will be evaluated after the assignment once the related
+    hlfir.region_assign is lowered.
+
+    Example: "X(VECTOR) = Y"
+
+    ```
+    hlfir.region_assign {
+      hlfir.yield %y : !fir.ref<!fir.array<20xf32>>
+    } to {
+      hlfir.elemental_addr %vector_shape  : !fir.shape<1> {
+        ^bb0(%i: index):
+        %0 = hlfir.designate %vector (%i)  : (!fir.ref<!fir.array<20xi32>>, index) -> !fir.ref<i32>
+        %1 = fir.load %0 : !fir.ref<i32>
+        %x_element_addr = hlfir.designate %x (%1)  : (!fir.ref<!fir.array<100xf32>>, i32) -> !fir.ref<f32>
+        hlfir.yield %x_element_addr : !fir.ref<f32>
+      }
+    }
+    ```
+  }];
+
+  let arguments = (ins
+    fir_ShapeType:$shape,
+    Variadic<AnyIntegerType>:$typeparams);
+
+  let regions = (region  SizedRegion<1>:$body,
+                         MaxSizedRegion<1>:$cleanup);
+
+  let builders = [
+    OpBuilder<(ins "mlir::Value":$shape)>
+  ];
+
+  let assemblyFormat = [{
+    $shape (`typeparams` $typeparams^)?
+    attr-dict `:` type(operands) $body
+    custom<YieldOpCleanup>($cleanup)}];
+
+  let extraClassDeclaration = [{
+    mlir::Region::BlockArgListType getIndices() {
+      return getBody().getArguments();
+    }
+  }];
+
+  let hasVerifier = 1;
+}
+
+
 #endif // FORTRAN_DIALECT_HLFIR_OPS

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

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())))
+  if (!mlir::isa_and_nonnull<hlfir::YieldOp, hlfir::ElementalAddrOp>(
+          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);
   }
 }
 
+//===----------------------------------------------------------------------===//
+// ElementalAddrOp
+//===----------------------------------------------------------------------===//
+
+void hlfir::ElementalAddrOp::build(mlir::OpBuilder &builder,
+                                   mlir::OperationState &odsState,
+                                   mlir::Value shape) {
+  odsState.addOperands(shape);
+  mlir::Region *bodyRegion = odsState.addRegion();
+  bodyRegion->push_back(new mlir::Block{});
+  if (auto shapeType = shape.getType().dyn_cast<fir::ShapeType>()) {
+    unsigned dim = shapeType.getRank();
+    mlir::Type indexType = builder.getIndexType();
+    for (unsigned d = 0; d < dim; ++d)
+      bodyRegion->front().addArgument(indexType, odsState.location);
+  }
+  // Push cleanUp region.
+  odsState.addRegion();
+}
+
+mlir::LogicalResult hlfir::ElementalAddrOp::verify() {
+  hlfir::YieldOp yieldOp =
+      mlir::dyn_cast_or_null<hlfir::YieldOp>(getTerminator(getBody()));
+  if (!yieldOp)
+    return emitOpError("body region must be terminated by an hlfir.yield");
+  mlir::Type elementAddrType = yieldOp.getEntity().getType();
+  if (!hlfir::isFortranVariableType(elementAddrType) ||
+      hlfir::getFortranElementOrSequenceType(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();
+}
+
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

flang/test/HLFIR/element-addr.fir

+// Test hlfir.element_addr operation parse, verify (no errors), and unparse.
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+func.func @test_element_addr(%x: !fir.ref<!fir.array<100xf32>>, %vector: !fir.ref<!fir.array<20xi32>>, %y: !fir.ref<!fir.array<20xf32>>) {
+  %c20 = arith.constant 20 : index
+  %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
+  hlfir.region_assign {
+    hlfir.yield %y : !fir.ref<!fir.array<20xf32>>
+  } to {
+    hlfir.elemental_addr %vector_shape  : !fir.shape<1> {
+      ^bb0(%i: index):
+      %0 = hlfir.designate %vector (%i)  : (!fir.ref<!fir.array<20xi32>>, index) -> !fir.ref<i32>
+      %1 = fir.load %0 : !fir.ref<i32>
+      %x_element_addr = hlfir.designate %x (%1)  : (!fir.ref<!fir.array<100xf32>>, i32) -> !fir.ref<f32>
+      hlfir.yield %x_element_addr : !fir.ref<f32>
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @test_element_addr(
+// CHECK-SAME:                                 %[[VAL_0:.*]]: !fir.ref<!fir.array<100xf32>>,
+// CHECK-SAME:                                 %[[VAL_1:.*]]: !fir.ref<!fir.array<20xi32>>,
+// CHECK-SAME:                                 %[[VAL_2:.*]]: !fir.ref<!fir.array<20xf32>>) {
+// CHECK:           %[[VAL_3:.*]] = arith.constant 20 : index
+// CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
+// CHECK:           hlfir.region_assign {
+// CHECK:             hlfir.yield %[[VAL_2]] : !fir.ref<!fir.array<20xf32>>
+// CHECK:           } to {
+// CHECK:             hlfir.elemental_addr %[[VAL_4]] : !fir.shape<1> {
+// CHECK:             ^bb0(%[[VAL_5:.*]]: index):
+// CHECK:               %[[VAL_6:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_5]])  : (!fir.ref<!fir.array<20xi32>>, index) -> !fir.ref<i32>
+// CHECK:               %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref<i32>
+// CHECK:               %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_7]])  : (!fir.ref<!fir.array<100xf32>>, i32) -> !fir.ref<f32>
+// CHECK:               hlfir.yield %[[VAL_8]] : !fir.ref<f32>
+// CHECK:             }
+// CHECK:           }
+
+
+
+func.func @test_element_addr_cleanup(%x: !fir.box<!fir.array<?x!fir.char<1,?>>>, %y: !fir.box<!fir.array<?x!fir.char<1,?>>>) {
+  hlfir.region_assign {
+    hlfir.yield %y : !fir.box<!fir.array<?x!fir.char<1,?>>>
+  } to {
+    %c0 = arith.constant 0 : index
+    %len = fir.box_elesize %x : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
+    %vector = fir.call @returns_alloc() : () -> !fir.box<!fir.heap<!fir.array<?xi32>>>
+    %vector_dim:3 = fir.box_dims %x, %c0 : (!fir.box<!fir.array<?x!fir.char<1,?>>>, index) -> (index, index, index)
+    %vector_shape = fir.shape %vector_dim#1 : (index) -> !fir.shape<1>
+    hlfir.elemental_addr %vector_shape  typeparams %len : !fir.shape<1>, index {
+      ^bb0(%i: index):
+      %0 = hlfir.designate %vector (%i)  : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> !fir.ref<i32>
+      %1 = fir.load %0 : !fir.ref<i32>
+      %x_element_addr = hlfir.designate %x (%1) typeparams %len  : (!fir.box<!fir.array<?x!fir.char<1,?>>>, i32, index) -> !fir.boxchar<1>
+      hlfir.yield %x_element_addr : !fir.boxchar<1>
+    } cleanup {
+      %addr = fir.box_addr %vector : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
+      fir.freemem %addr : !fir.heap<!fir.array<?xi32>>
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @test_element_addr_cleanup(
+// CHECK-SAME:                                         %[[VAL_0:[^:]*]]: !fir.box<!fir.array<?x!fir.char<1,?>>>,
+// CHECK-SAME:                                         %[[VAL_1:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>>) {
+// CHECK:           hlfir.region_assign {
+// CHECK:             hlfir.yield %[[VAL_1]] : !fir.box<!fir.array<?x!fir.char<1,?>>>
+// CHECK:           } to {
+// CHECK:             %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK:             %[[VAL_3:.*]] = fir.box_elesize %[[VAL_0]] : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
+// CHECK:             %[[VAL_4:.*]] = fir.call @returns_alloc() : () -> !fir.box<!fir.heap<!fir.array<?xi32>>>
+// CHECK:             %[[VAL_5:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x!fir.char<1,?>>>, index) -> (index, index, index)
+// CHECK:             %[[VAL_6:.*]] = fir.shape %[[VAL_5]]#1 : (index) -> !fir.shape<1>
+// CHECK:             hlfir.elemental_addr %[[VAL_6]] typeparams %[[VAL_3]] : !fir.shape<1>, index {
+// CHECK:             ^bb0(%[[VAL_7:.*]]: index):
+// CHECK:               %[[VAL_8:.*]] = hlfir.designate %[[VAL_4]] (%[[VAL_7]])  : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> !fir.ref<i32>
+// CHECK:               %[[VAL_9:.*]] = fir.load %[[VAL_8]] : !fir.ref<i32>
+// CHECK:               %[[VAL_10:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_9]])  typeparams %[[VAL_3]] : (!fir.box<!fir.array<?x!fir.char<1,?>>>, i32, index) -> !fir.boxchar<1>
+// CHECK:               hlfir.yield %[[VAL_10]] : !fir.boxchar<1>
+// CHECK:             } cleanup {
+// CHECK:               %[[VAL_11:.*]] = fir.box_addr %[[VAL_4]] : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
+// CHECK:               fir.freemem %[[VAL_11]] : !fir.heap<!fir.array<?xi32>>
+// CHECK:             }
+// CHECK:           }

flang/test/HLFIR/invalid.fir

// 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
 }
+
+// -----
+func.func @bad_element_addr_1(%x: !fir.ref<!fir.array<20xf32>>) {
+  %c20 = arith.constant 20 : index
+  %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
+  hlfir.region_assign {
+    hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
+  } to {
+    // expected-error@+1 {{'hlfir.elemental_addr' op body must compute the address of a scalar entity}}
+    hlfir.elemental_addr %vector_shape  : !fir.shape<1> {
+      ^bb0(%i: index):
+      %c42 = arith.constant 42.0 : f32
+      hlfir.yield %c42 : f32
+    }
+  }
+  return
+}
+
+// -----
+func.func @bad_element_addr_2(%x: !fir.ref<!fir.array<20xf32>>) {
+  %c20 = arith.constant 20 : index
+  %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
+  hlfir.region_assign {
+    hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
+  } to {
+    // expected-error@+1 {{'hlfir.elemental_addr' op body must compute the address of a scalar entity}}
+    hlfir.elemental_addr %vector_shape  : !fir.shape<1> {
+      ^bb0(%i: index):
+      hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
+    }
+  }
+  return
+}
+
+// -----
+func.func @bad_element_addr_3(%x: !fir.ref<!fir.array<20xf32>>) {
+  %c20 = arith.constant 20 : index
+  %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
+  hlfir.region_assign {
+    hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
+  } to {
+    // expected-error@+1 {{'hlfir.elemental_addr' op body region must be terminated by an hlfir.yield}}
+    hlfir.elemental_addr %vector_shape  : !fir.shape<1> {
+      ^bb0(%i: index):
+      %c42 = arith.constant 42.0 : f32
+    }
+  }
+  return
+}
+
+// -----
+func.func @bad_element_addr_4(%x: !fir.ref<!fir.array<20xf32>>, %y: !fir.ref<!fir.array<20x20xf32>>) {
+  %c20 = arith.constant 20 : index
+  %vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
+  hlfir.region_assign {
+    hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
+  } to {
+    // expected-error@+1 {{'hlfir.elemental_addr' op body number of indices must match shape rank}}
+    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
+}