[flang] Add hlfir.apply definition

hlfir.apply allows retrieving element values from an array expression
value. See https://github.com/llvm/llvm-project/blob/main/flang/docs/HighLevelFIR.md
for more detail.

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

flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td

@@ -67,6 +67,13 @@ def hlfir_ExprType : TypeDef<hlfir_Dialect, "Expr"> {
     bool isArray() const { return !isScalar(); }
     bool isPolymorphic() const { return getPolymorphic(); }
     unsigned getRank() const {return getShape().size();}
+    mlir::Type getElementExprType() const {
+      mlir::Type eleTy = getElementType();
+      if (fir::isa_trivial(eleTy))
+        return eleTy;
+      return hlfir::ExprType::get(eleTy.getContext(), Shape{}, eleTy,
+                isPolymorphic());
+    }
   }];
 
   let hasCustomAssemblyFormat = 1;

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

@@ -403,4 +403,31 @@ def hlfir_YieldElementOp : hlfir_Op<"yield_element", [Terminator, HasParent<"Ele
   let assemblyFormat = "$element_value attr-dict `:` type($element_value)";
 }
 
+def hlfir_ApplyOp : hlfir_Op<"apply", [NoMemoryEffect, AttrSizedOperandSegments]> {
+  let summary = "get the element value of an expression";
+  let description = [{
+    Given an hlfir.expr array value, hlfir.apply allow retrieving
+    the value for an element given one based indices.
+    When hlfir.apply is used on an hlfir.elemental, and if the hlfir.elemental
+    operation evaluation can be moved to the location of the hlfir.apply, it is
+    as if the hlfir.elemental body was evaluated given the hlfir.apply indices.
+  }];
+
+  let arguments = (ins hlfir_ExprType:$expr,
+                   Variadic<Index>:$indices,
+                   Variadic<AnyIntegerType>:$typeparams
+                  );
+  let results = (outs AnyFortranValue:$element_value);
+
+  let assemblyFormat = [{
+    $expr `,` $indices (`typeparams` $typeparams^)?
+    attr-dict `:` functional-type(operands, results)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "mlir::Value":$expr, "mlir::ValueRange":$indices,
+      "mlir::ValueRange":$typeparams)>
+  ];
+}
+
 #endif // FORTRAN_DIALECT_HLFIR_OPS

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

@@ -454,5 +454,19 @@ void hlfir::ElementalOp::build(mlir::OpBuilder &builder,
   }
 }
 
+//===----------------------------------------------------------------------===//
+// ApplyOp
+//===----------------------------------------------------------------------===//
+
+void hlfir::ApplyOp::build(mlir::OpBuilder &builder,
+                           mlir::OperationState &odsState, mlir::Value expr,
+                           mlir::ValueRange indices,
+                           mlir::ValueRange typeparams) {
+  mlir::Type resultType = expr.getType();
+  if (auto exprType = resultType.dyn_cast<hlfir::ExprType>())
+    resultType = exprType.getElementExprType();
+  build(builder, odsState, resultType, expr, indices, typeparams);
+}
+
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

flang/test/HLFIR/apply.fir

@@ -0,0 +1,42 @@
+// Test hlfir.apply operation parse, verify (no errors), and unparse.
+
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+func.func @numeric(%expr: !hlfir.expr<?x?xf32>) {
+  %c9 = arith.constant 9 : index
+  %c2 = arith.constant 2 : index
+  %0 = hlfir.apply %expr, %c9, %c2 : (!hlfir.expr<?x?xf32>, index, index) -> f32
+  return
+}
+// CHECK-LABEL:   func.func @numeric(
+// CHECK-SAME:    %[[VAL_0:.*]]: !hlfir.expr<?x?xf32>) {
+// CHECK:  %[[VAL_1:.*]] = arith.constant 9 : index
+// CHECK:  %[[VAL_2:.*]] = arith.constant 2 : index
+// CHECK:  %[[VAL_3:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_1]], %[[VAL_2]] : (!hlfir.expr<?x?xf32>, index, index) -> f32
+
+func.func @char(%expr: !hlfir.expr<?x?x!fir.char<1,?>>, %l: index) {
+  %c9 = arith.constant 9 : index
+  %c2 = arith.constant 2 : index
+  %0 = hlfir.apply %expr, %c9, %c2 typeparams %l: (!hlfir.expr<?x?x!fir.char<1,?>>, index, index, index) -> !hlfir.expr<!fir.char<1,?>>
+  return
+}
+// CHECK-LABEL:   func.func @char(
+// CHECK-SAME:    %[[VAL_0:.*]]: !hlfir.expr<?x?x!fir.char<1,?>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: index) {
+// CHECK:  %[[VAL_2:.*]] = arith.constant 9 : index
+// CHECK:  %[[VAL_3:.*]] = arith.constant 2 : index
+// CHECK:  %[[VAL_4:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_2]], %[[VAL_3]] typeparams %[[VAL_1]] : (!hlfir.expr<?x?x!fir.char<1,?>>, index, index, index) -> !hlfir.expr<!fir.char<1,?>>
+
+!pdt = !fir.type<pdt(param:i32){field:f32}>
+func.func @derived(%expr: !hlfir.expr<?x?x!pdt>, %l: i32) {
+  %c9 = arith.constant 9 : index
+  %c2 = arith.constant 2 : index
+  %0 = hlfir.apply %expr, %c9, %c2 typeparams %l: (!hlfir.expr<?x?x!pdt>, index, index, i32) -> !hlfir.expr<!pdt>
+  return
+}
+// CHECK-LABEL:   func.func @derived(
+// CHECK-SAME:    %[[VAL_0:.*]]: !hlfir.expr<?x?x!fir.type<pdt(param:i32){field:f32}>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: i32) {
+// CHECK:  %[[VAL_2:.*]] = arith.constant 9 : index
+// CHECK:  %[[VAL_3:.*]] = arith.constant 2 : index
+// CHECK:  %[[VAL_4:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_2]], %[[VAL_3]] typeparams %[[VAL_1]] : (!hlfir.expr<?x?x!fir.type<pdt(param:i32){field:f32}>>, index, index, i32) -> !hlfir.expr<!fir.type<pdt(param:i32){field:f32}>>