[flang] hlfir.apply codegen

[jeanPerier]

Summary

Lower hlfir.apply to hlfir.designate in bufferization.

Depends on D140099

[PeteSteinfeld]

I don't understand this code very well, so it would be good to have someone else approve. But all builds and tests correctly and looks good.

[clementval]

LGTM

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

matchAndRewrite(hlfir::AsExprOp asExpr, OpAdaptor adaptor,
     builder.create<hlfir::AssignOp>(loc, source, temp);
     mlir::Value bufferizedExpr =
         packageBufferizedExpr(loc, builder, temp, cleanup);
     rewriter.replaceOp(asExpr, bufferizedExpr);
     return mlir::success();
   }
 };
 
+struct ApplyOpConversion : public mlir::OpConversionPattern<hlfir::ApplyOp> {
+  using mlir::OpConversionPattern<hlfir::ApplyOp>::OpConversionPattern;
+  explicit ApplyOpConversion(mlir::MLIRContext *ctx)
+      : mlir::OpConversionPattern<hlfir::ApplyOp>{ctx} {}
+  mlir::LogicalResult
+  matchAndRewrite(hlfir::ApplyOp apply, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Location loc = apply->getLoc();
+    hlfir::Entity bufferizedExpr{getBufferizedExprStorage(adaptor.getExpr())};
+    mlir::Type resultType = hlfir::getVariableElementType(bufferizedExpr);
+    mlir::Value result = rewriter.create<hlfir::DesignateOp>(
+        loc, resultType, bufferizedExpr, adaptor.getIndices(),
+        adaptor.getTypeparams());
+    if (fir::isa_trivial(apply.getType()))
+      result = rewriter.create<fir::LoadOp>(loc, result);
+    rewriter.replaceOp(apply, result);
+    return mlir::success();
+  }
+};
+
 struct AssignOpConversion : public mlir::OpConversionPattern<hlfir::AssignOp> {
   using mlir::OpConversionPattern<hlfir::AssignOp>::OpConversionPattern;
   explicit AssignOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::AssignOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::AssignOp assign, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<hlfir::AssignOp>(
void runOnOperation() override {
     // intrinsics calls creation. This may create race conflict if the pass is
     // scheduled on FuncOp. A solution could be to provide an optional mutex
     // when building a FirOpBuilder and locking around FuncOp and GlobalOp
     // creation, but this needs a bit more thinking, so at this point the pass
     // is scheduled on the moduleOp.
     auto module = this->getOperation();
     auto *context = &getContext();
     mlir::RewritePatternSet patterns(context);
-    patterns
-        .insert<AsExprOpConversion, AssignOpConversion, AssociateOpConversion,
-                ConcatOpConversion, ElementalOpConversion,
-                EndAssociateOpConversion, NoReassocOpConversion>(context);
+    patterns.insert<ApplyOpConversion, AsExprOpConversion, AssignOpConversion,
+                    AssociateOpConversion, ConcatOpConversion,
+                    ElementalOpConversion, EndAssociateOpConversion,
+                    NoReassocOpConversion>(context);
     mlir::ConversionTarget target(*context);
-    target.addIllegalOp<hlfir::AssociateOp, hlfir::ElementalOp,
+    target.addIllegalOp<hlfir::ApplyOp, hlfir::AssociateOp, hlfir::ElementalOp,
                         hlfir::EndAssociateOp, hlfir::YieldElementOp>();
     target.markUnknownOpDynamicallyLegal([](mlir::Operation *op) {
       return llvm::all_of(
                  op->getResultTypes(),
                  [](mlir::Type ty) { return !ty.isa<hlfir::ExprType>(); }) &&
              llvm::all_of(op->getOperandTypes(), [](mlir::Type ty) {
                return !ty.isa<hlfir::ExprType>();
              });

flang/test/HLFIR/apply-codegen.fir

+// Test hlfir.apply code generation
+// RUN: fir-opt %s -bufferize-hlfir | FileCheck %s
+
+func.func @numeric_apply(%arg0 : !fir.ref<!fir.array<100xi32>>) {
+  %expr = hlfir.as_expr %arg0 : (!fir.ref<!fir.array<100xi32>>) -> !hlfir.expr<100xi32>
+  %c42 = arith.constant 42 : index
+  %elem = hlfir.apply %expr, %c42 : (!hlfir.expr<100xi32>, index) -> i32
+  return
+}
+// CHECK-LABEL:   func.func @numeric_apply(
+// CHECK:    %[[VAL_5:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = ".tmp"}
+// CHECK:    %[[VAL_9:.*]] = arith.constant 42 : index
+// CHECK:    %[[VAL_10:.*]] = hlfir.designate %[[VAL_5]]#0 (%[[VAL_9]])  : (!fir.heap<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+// CHECK:    %[[VAL_11:.*]] = fir.load %[[VAL_10]] : !fir.ref<i32>
+
+func.func @character_apply(%arg0 : !fir.ref<!fir.array<?x!fir.char<2,?>>>, %l: index, %n: index) {
+  %shape = fir.shape %n : (index) -> !fir.shape<1>
+  %c:2 = hlfir.declare %arg0(%shape) typeparams %l {uniq_name = "c"} : (!fir.ref<!fir.array<?x!fir.char<2,?>>>, !fir.shape<1>, index) -> (!fir.box<!fir.array<?x!fir.char<2,?>>>, !fir.ref<!fir.array<?x!fir.char<2,?>>>)
+  %expr = hlfir.as_expr %c#0 : (!fir.box<!fir.array<?x!fir.char<2,?>>>) -> !hlfir.expr<?x!fir.char<2,?>>
+  %c42 = arith.constant 42 : index
+  %elem = hlfir.apply %expr, %c42 typeparams %l: (!hlfir.expr<?x!fir.char<2,?>>, index, index) -> !hlfir.expr<!fir.char<2,?>>
+  return
+}
+// CHECK-LABEL:   func.func @character_apply(
+// CHECK:    %[[VAL_7:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) typeparams %[[VAL_1:.*]] {uniq_name = ".tmp"}
+// CHECK:    %[[VAL_11:.*]] = arith.constant 42 : index
+// CHECK:    %[[VAL_12:.*]] = hlfir.designate %[[VAL_7]]#0 (%[[VAL_11]])  typeparams %[[VAL_1]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>, index, index) -> !fir.boxchar<2>