[flang-commits] [flang] [flang][HLFIR] compute elemental function result length parameters (PR #93983)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/93983

Prepare the argument and map them to their corresponding dummy symbol in order to lower the specification expression of the function result.

Extract the preparation of arguments according to the interface to its own function to be reused.

It seems there is no need to conditionally compute the length on the input since all the information comes from the CharBoxValue or the descriptor for cases where the number of element could be 0. 


>From 1719b271ee8f4927a220623f305168e5a6b71d61 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 29 May 2024 08:15:44 -0700
Subject: [PATCH] [flang][HLFIR] compute elemental function result length
 parameters

Prepare the argument and map them to their corresponding dummy symbol
in order to lower the specification expression of the function result.
---
 flang/lib/Lower/ConvertCall.cpp               | 77 ++++++++++++---
 .../Lower/HLFIR/elemental-result-length.f90   | 95 +++++++++++++++++++
 2 files changed, 159 insertions(+), 13 deletions(-)
 create mode 100644 flang/test/Lower/HLFIR/elemental-result-length.f90

diff --git a/flang/lib/Lower/ConvertCall.cpp b/flang/lib/Lower/ConvertCall.cpp
index 7ec719a2cb9ec..f8442b8e61e46 100644
--- a/flang/lib/Lower/ConvertCall.cpp
+++ b/flang/lib/Lower/ConvertCall.cpp
@@ -1454,21 +1454,16 @@ static PreparedDummyArgument prepareProcedurePointerActualArgument(
   return PreparedDummyArgument{tempBoxProc, /*cleanups=*/{}};
 }
 
-/// Lower calls to user procedures with actual arguments that have been
-/// pre-lowered but not yet prepared according to the interface.
-/// This can be called for elemental procedures, but only with scalar
-/// arguments: if there are array arguments, it must be provided with
-/// the array argument elements value and will return the corresponding
-/// scalar result value.
-static std::optional<hlfir::EntityWithAttributes>
-genUserCall(Fortran::lower::PreparedActualArguments &loweredActuals,
-            Fortran::lower::CallerInterface &caller,
-            mlir::FunctionType callSiteType, CallContext &callContext) {
+/// Prepare arguments of calls to user procedures with actual arguments that
+/// have been pre-lowered but not yet prepared according to the interface.
+void prepareUserCallArguments(
+    Fortran::lower::PreparedActualArguments &loweredActuals,
+    Fortran::lower::CallerInterface &caller, mlir::FunctionType callSiteType,
+    CallContext &callContext, llvm::SmallVector<CallCleanUp> &callCleanUps) {
   using PassBy = Fortran::lower::CallerInterface::PassEntityBy;
   mlir::Location loc = callContext.loc;
   bool mustRemapActualToDummyDescriptors = false;
   fir::FirOpBuilder &builder = callContext.getBuilder();
-  llvm::SmallVector<CallCleanUp> callCleanUps;
   for (auto [preparedActual, arg] :
        llvm::zip(loweredActuals, caller.getPassedArguments())) {
     mlir::Type argTy = callSiteType.getInput(arg.firArgument);
@@ -1626,11 +1621,30 @@ genUserCall(Fortran::lower::PreparedActualArguments &loweredActuals,
     } break;
     }
   }
+
   // Handle cases where caller must allocate the result or a fir.box for it.
   if (mustRemapActualToDummyDescriptors)
     remapActualToDummyDescriptors(loc, callContext.converter,
                                   callContext.symMap, loweredActuals, caller,
                                   callContext.isBindcCall());
+}
+
+/// Lower calls to user procedures with actual arguments that have been
+/// pre-lowered but not yet prepared according to the interface.
+/// This can be called for elemental procedures, but only with scalar
+/// arguments: if there are array arguments, it must be provided with
+/// the array argument elements value and will return the corresponding
+/// scalar result value.
+static std::optional<hlfir::EntityWithAttributes>
+genUserCall(Fortran::lower::PreparedActualArguments &loweredActuals,
+            Fortran::lower::CallerInterface &caller,
+            mlir::FunctionType callSiteType, CallContext &callContext) {
+  mlir::Location loc = callContext.loc;
+  llvm::SmallVector<CallCleanUp> callCleanUps;
+  fir::FirOpBuilder &builder = callContext.getBuilder();
+
+  prepareUserCallArguments(loweredActuals, caller, callSiteType, callContext,
+                           callCleanUps);
 
   // Prepare lowered arguments according to the interface
   // and map the lowered values to the dummy
@@ -2204,8 +2218,45 @@ class ElementalUserCallBuilder
   mlir::Value computeDynamicCharacterResultLength(
       Fortran::lower::PreparedActualArguments &loweredActuals,
       CallContext &callContext) {
-    TODO(callContext.loc,
-         "compute elemental function result length parameters in HLFIR");
+    fir::FirOpBuilder &builder = callContext.getBuilder();
+    mlir::Location loc = callContext.loc;
+    auto &converter = callContext.converter;
+    mlir::Type idxTy = builder.getIndexType();
+    llvm::SmallVector<CallCleanUp> callCleanUps;
+
+    prepareUserCallArguments(loweredActuals, caller, callSiteType, callContext,
+                             callCleanUps);
+
+    callContext.symMap.pushScope();
+
+    // Map prepared argument to dummy symbol to be able to lower spec expr.
+    for (const auto &arg : caller.getPassedArguments()) {
+      const Fortran::semantics::Symbol *sym = caller.getDummySymbol(arg);
+      assert(sym && "expect symbol for dummy argument");
+      auto input = caller.getInput(arg);
+      fir::ExtendedValue exv = Fortran::lower::translateToExtendedValue(
+          loc, builder, hlfir::Entity{input}, callContext.stmtCtx);
+      fir::FortranVariableOpInterface variableIface = hlfir::genDeclare(
+          loc, builder, exv, "dummy.tmp", fir::FortranVariableFlagsAttr{});
+      callContext.symMap.addVariableDefinition(*sym, variableIface);
+    }
+
+    auto lowerSpecExpr = [&](const auto &expr) -> mlir::Value {
+      mlir::Value convertExpr = builder.createConvert(
+          loc, idxTy,
+          fir::getBase(converter.genExprValue(expr, callContext.stmtCtx)));
+      return fir::factory::genMaxWithZero(builder, loc, convertExpr);
+    };
+
+    llvm::SmallVector<mlir::Value> lengths;
+    caller.walkResultLengths(
+        [&](const Fortran::lower::SomeExpr &e, bool isAssumedSizeExtent) {
+          assert(!isAssumedSizeExtent && "result cannot be assumed-size");
+          lengths.emplace_back(lowerSpecExpr(e));
+        });
+    callContext.symMap.popScope();
+    assert(lengths.size() == 1 && "expect 1 length parameter for the result");
+    return lengths[0];
   }
 
   mlir::Value getPolymorphicResultMold(
diff --git a/flang/test/Lower/HLFIR/elemental-result-length.f90 b/flang/test/Lower/HLFIR/elemental-result-length.f90
new file mode 100644
index 0000000000000..0aaf7c93770c9
--- /dev/null
+++ b/flang/test/Lower/HLFIR/elemental-result-length.f90
@@ -0,0 +1,95 @@
+! RUN: bbc -emit-hlfir -o - %s | fir-opt --canonicalize | FileCheck %s
+
+module m1
+contains
+elemental function fct1(a, b) result(t)
+  character(*), intent(in) :: a, b
+  character(len(a) + len(b)) :: t
+  t = a // b
+end function
+
+elemental function fct2(c) result(t)
+  integer, intent(in) :: c
+  character(c) :: t
+
+end function
+
+subroutine sub2(a,b,c)
+  character(*), intent(inout) :: c
+  character(*), intent(in) :: a, b
+
+  c = fct1(a,b)
+end subroutine
+
+! CHECK-LABEL: func.func @_QMm1Psub2(
+! CHECK-SAME: %[[ARG0:.*]]: !fir.boxchar<1> {fir.bindc_name = "a"}, %[[ARG1:.*]]: !fir.boxchar<1> {fir.bindc_name = "b"}, %[[ARG2:.*]]: !fir.boxchar<1> {fir.bindc_name = "c"}) {
+! CHECK: %[[UNBOX_ARG0:.*]]:2 = fir.unboxchar %[[ARG0]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+! CHECK: %[[A:.*]]:2 = hlfir.declare %[[UNBOX_ARG0]]#0 typeparams %[[UNBOX_ARG0]]#1 dummy_scope %0 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Fsub2Ea"} : (!fir.ref<!fir.char<1,?>>, index, !fir.dscope) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
+! CHECK: %[[UNBOX_ARG1:.*]]:2 = fir.unboxchar %[[ARG1]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+! CHECK: %[[B:.*]]:2 = hlfir.declare %[[UNBOX_ARG1]]#0 typeparams %[[UNBOX_ARG1]]#1 dummy_scope %{{.*}} {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Fsub2Eb"} : (!fir.ref<!fir.char<1,?>>, index, !fir.dscope) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
+! CHECK: %[[UNBOX_ARG2:.*]]:2 = fir.unboxchar %[[ARG2]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+! CHECK: %[[C:.*]]:2 = hlfir.declare %[[UNBOX_ARG2]]#0 typeparams %[[UNBOX_ARG2]]#1 dummy_scope %{{.*}} {fortran_attrs = #fir.var_attrs<intent_inout>, uniq_name = "_QMm1Fsub2Ec"} : (!fir.ref<!fir.char<1,?>>, index, !fir.dscope) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
+! CHECK: %[[UNBOX_A:.*]]:2 = fir.unboxchar %[[A]]#0 : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+! CHECK: %[[DUMMYA:.*]]:2 = hlfir.declare %[[UNBOX_A]]#0 typeparams %[[UNBOX_A]]#1 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Ffct1Ea"} : (!fir.ref<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
+! CHECK: %[[UNBOX_B:.*]]:2 = fir.unboxchar %[[B]]#0 : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+! CHECK: %[[DUMMYB:.*]]:2 = hlfir.declare %[[UNBOX_B]]#0 typeparams %[[UNBOX_B]]#1 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Ffct1Eb"} : (!fir.ref<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
+! CHECK: %[[LEN_A:.*]] = fir.convert %[[UNBOX_A]]#1 : (index) -> i32
+! CHECK: %[[LEN_B:.*]] = fir.convert %[[UNBOX_B]]#1 : (index) -> i32
+! CHECK: %[[LEN_LEN:.*]] = arith.addi %[[LEN_A]], %[[LEN_B]] : i32
+! CHECK: %[[LEN_LEN_IDX:.*]] = fir.convert %[[LEN_LEN]] : (i32) -> index
+! CHECK: %[[CMPI:.*]] = arith.cmpi sgt, %[[LEN_LEN_IDX]], %c0{{.*}} : index
+! CHECK: %[[RES_LENGTH:.*]] = arith.select %[[CMPI]], %[[LEN_LEN_IDX]], %c0{{.*}} : index
+! CHECK: %[[RES:.*]] = fir.alloca !fir.char<1,?>(%[[RES_LENGTH]] : index) {bindc_name = ".result"}
+! CHECK: fir.call @_QMm1Pfct1
+
+subroutine sub3(c)
+  character(*), intent(inout) :: c(:)
+
+  c = fct2(10)
+end subroutine
+
+! CHECK-LABEL: func.func @_QMm1Psub3(
+! CHECK-SAME: %[[ARG0:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>> {fir.bindc_name = "c"}) {
+! CHECK: %[[C10:.*]] = arith.constant 10 : i32
+! CHECK: %[[C:.*]]:2 = hlfir.declare %[[ARG0]] dummy_scope %{{.*}} {fortran_attrs = #fir.var_attrs<intent_inout>, uniq_name = "_QMm1Fsub3Ec"} : (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.dscope) -> (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.box<!fir.array<?x!fir.char<1,?>>>)
+! CHECK: %[[ASSOC:.*]]:3 = hlfir.associate %[[C10]] {adapt.valuebyref} : (i32) -> (!fir.ref<i32>, !fir.ref<i32>, i1)
+! CHECK: %[[INPUT_ARG0:.*]]:2 = hlfir.declare %[[ASSOC]]#1 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Ffct2Ec"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[LOAD_INPUT_ARG0:.*]] = fir.load %[[INPUT_ARG0]]#0 : !fir.ref<i32>
+! CHECK: %[[LOAD_INPUT_ARG0_IDX:.*]] = fir.convert %[[LOAD_INPUT_ARG0]] : (i32) -> index
+! CHECK: %[[CMPI:.*]] = arith.cmpi sgt, %[[LOAD_INPUT_ARG0_IDX]], %c0{{.*}} : index
+! CHECK: %[[LENGTH:.*]] = arith.select %[[CMPI]], %[[LOAD_INPUT_ARG0_IDX]], %c0{{.*}} : index
+! CHECK: %[[RES:.*]] = fir.alloca !fir.char<1,?>(%[[LENGTH]] : index) {bindc_name = ".result"}
+! CHECK: fir.call @_QMm1Pfct2
+
+subroutine sub4(a,b,c)
+  character(*), intent(inout) :: c(:)
+  character(*), intent(in) :: a(:), b(:)
+
+  c = fct1(a,b)
+end subroutine
+
+! CHECK-LABEL: func.func @_QMm1Psub4(
+! CHECK-SAME: %[[ARG0:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>> {fir.bindc_name = "a"}, %[[ARG1:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>> {fir.bindc_name = "b"}, %[[ARG2:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>> {fir.bindc_name = "c"}) {
+! CHECK: %[[A:.*]]:2 = hlfir.declare %[[ARG0]] dummy_scope %{{.*}} {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Fsub4Ea"} : (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.dscope) -> (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.box<!fir.array<?x!fir.char<1,?>>>)
+! CHECK: %[[B:.*]]:2 = hlfir.declare %[[ARG1]] dummy_scope %{{.*}} {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QMm1Fsub4Eb"} : (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.dscope) -> (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.box<!fir.array<?x!fir.char<1,?>>>)
+! CHECK: %[[C:.*]]:2 = hlfir.declare %[[ARG2]] dummy_scope %{{.*}} {fortran_attrs = #fir.var_attrs<intent_inout>, uniq_name = "_QMm1Fsub4Ec"} : (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.dscope) -> (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.box<!fir.array<?x!fir.char<1,?>>>)
+! CHECK: %[[LEN_A:.*]] = fir.box_elesize %[[A]]#1 : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
+! CHECK: %[[LEN_B:.*]] = fir.box_elesize %[[B]]#1 : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
+! CHECK: %[[LEN_A_I32:.*]] = fir.convert %[[LEN_A]] : (index) -> i32
+! CHECK: %[[LEN_B_I32:.*]] = fir.convert %[[LEN_B]] : (index) -> i32
+! CHECK: %[[LEN_LEN:.*]] = arith.addi %[[LEN_A_I32]], %[[LEN_B_I32]] : i32
+! CHECK: %[[LEN_LEN_IDX:.*]] = fir.convert %[[LEN_LEN]] : (i32) -> index
+! CHECK: %[[CMPI:.*]] = arith.cmpi sgt, %[[LEN_LEN_IDX]], %c0{{.*}} : index
+! CHECK: %[[LENGTH:.*]] = arith.select %[[CMPI]], %17, %c0{{.*}} : index
+! CHECK: %{{.*}} = hlfir.elemental %{{.*}} typeparams %[[LENGTH]] unordered : (!fir.shape<1>, index) -> !hlfir.expr<?x!fir.char<1,?>>
+
+end module
+
+program test
+  use m1
+  character(5) :: a(2) = ['abcde', 'klmnop'], b(2) = ['fghij', 'qrstu']
+  character(10) :: c(2)
+
+  call sub2(a(1), b(1), c(1))
+  print*, c(1)
+end