[flang] Lower elemental calls

[clementval]

Summary

This patch adds more lowering of operations sub-expression inside elemental call arguments.
It tests array contexts where an address is needed for each element (for
the argument), but part of the array sub-expression must be lowered by value
(for the operation)

This patch is part of the upstreaming effort from fir-dev branch.

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Lower/ConvertExpr.cpp

return [=](IterSpace iters) -> ExtValue {
           !(fir::isa_ref_type(valTy) && fir::isa_integer(ty)))
         return convertAdjustedType(builder, loc, ty, exv);
       return builder.createConvert(loc, ty, val);
     };
   }
 
   template <int KIND>
   CC genarr(const Fortran::evaluate::ComplexComponent<KIND> &x) {
-    TODO(getLoc(), "");
+    TODO(getLoc(), "ComplexComponent<KIND>");
   }
 
   template <typename T>
   CC genarr(const Fortran::evaluate::Parentheses<T> &x) {
-    TODO(getLoc(), "");
+    mlir::Location loc = getLoc();
+    if (isReferentiallyOpaque()) {
+      // Context is a call argument in, for example, an elemental procedure
+      // call. TODO: all array arguments should use array_load, array_access,
+      // array_amend, and INTENT(OUT), INTENT(INOUT) arguments should have
+      // array_merge_store ops.
+      TODO(loc, "parentheses on argument in elemental call");
+    }
+    auto f = genarr(x.left());
+    return [=](IterSpace iters) -> ExtValue {
+      auto val = f(iters);
+      mlir::Value base = fir::getBase(val);
+      auto newBase =
+          builder.create<fir::NoReassocOp>(loc, base.getType(), base);
+      return fir::substBase(val, newBase);
+    };
   }
-
   template <int KIND>
   CC genarr(const Fortran::evaluate::Negate<Fortran::evaluate::Type<
                 Fortran::common::TypeCategory::Integer, KIND>> &x) {
-    TODO(getLoc(), "");
+    mlir::Location loc = getLoc();
+    auto f = genarr(x.left());
+    return [=](IterSpace iters) -> ExtValue {
+      mlir::Value val = fir::getBase(f(iters));
+      mlir::Type ty =
+          converter.genType(Fortran::common::TypeCategory::Integer, KIND);
+      mlir::Value zero = builder.createIntegerConstant(loc, ty, 0);
+      return builder.create<mlir::arith::SubIOp>(loc, zero, val);
+    };
   }
-
   template <int KIND>
   CC genarr(const Fortran::evaluate::Negate<Fortran::evaluate::Type<
                 Fortran::common::TypeCategory::Real, KIND>> &x) {
     mlir::Location loc = getLoc();
     auto f = genarr(x.left());
     return [=](IterSpace iters) -> ExtValue {
       return builder.create<mlir::arith::NegFOp>(loc, fir::getBase(f(iters)));
     };
   }
   template <int KIND>
   CC genarr(const Fortran::evaluate::Negate<Fortran::evaluate::Type<
                 Fortran::common::TypeCategory::Complex, KIND>> &x) {
-    TODO(getLoc(), "");
+    mlir::Location loc = getLoc();
+    auto f = genarr(x.left());
+    return [=](IterSpace iters) -> ExtValue {
+      return builder.create<fir::NegcOp>(loc, fir::getBase(f(iters)));
+    };
   }
 
   //===--------------------------------------------------------------------===//
   // Binary elemental ops
   //===--------------------------------------------------------------------===//
 
   template <typename OP, typename A>
   CC createBinaryOp(const A &evEx) {
#define GENBIN(GenBinEvOp, GenBinTyCat, GenBinFirOp)                           \
   GENBIN(Multiply, Complex, fir::MulcOp)
   GENBIN(Divide, Integer, mlir::arith::DivSIOp)
   GENBIN(Divide, Real, mlir::arith::DivFOp)
   GENBIN(Divide, Complex, fir::DivcOp)
 
   template <Fortran::common::TypeCategory TC, int KIND>
   CC genarr(
       const Fortran::evaluate::Power<Fortran::evaluate::Type<TC, KIND>> &x) {
-    TODO(getLoc(), "genarr Power<Fortran::evaluate::Type<TC, KIND>>");
+    mlir::Location loc = getLoc();
+    mlir::Type ty = converter.genType(TC, KIND);
+    auto lf = genarr(x.left());
+    auto rf = genarr(x.right());
+    return [=](IterSpace iters) -> ExtValue {
+      mlir::Value lhs = fir::getBase(lf(iters));
+      mlir::Value rhs = fir::getBase(rf(iters));
+      return Fortran::lower::genPow(builder, loc, ty, lhs, rhs);
+    };
   }
   template <Fortran::common::TypeCategory TC, int KIND>
   CC genarr(
       const Fortran::evaluate::Extremum<Fortran::evaluate::Type<TC, KIND>> &x) {
     TODO(getLoc(), "genarr Extremum<Fortran::evaluate::Type<TC, KIND>>");
   }
   template <Fortran::common::TypeCategory TC, int KIND>
   CC genarr(
#define GENBIN(GenBinEvOp, GenBinTyCat, GenBinFirOp)                           \
   CC genarr(const Fortran::evaluate::DescriptorInquiry &x) {
     TODO(getLoc(), "genarr DescriptorInquiry");
   }
 
   CC genarr(const Fortran::evaluate::StructureConstructor &x) {
     TODO(getLoc(), "genarr StructureConstructor");
   }
 
+  //===--------------------------------------------------------------------===//
+  // LOCICAL operators (.NOT., .AND., .EQV., etc.)
+  //===--------------------------------------------------------------------===//
+
   template <int KIND>
   CC genarr(const Fortran::evaluate::Not<KIND> &x) {
-    TODO(getLoc(), "genarr Not");
+    mlir::Location loc = getLoc();
+    mlir::IntegerType i1Ty = builder.getI1Type();
+    auto lambda = genarr(x.left());
+    mlir::Value truth = builder.createBool(loc, true);
+    return [=](IterSpace iters) -> ExtValue {
+      mlir::Value logical = fir::getBase(lambda(iters));
+      mlir::Value val = builder.createConvert(loc, i1Ty, logical);
+      return builder.create<mlir::arith::XOrIOp>(loc, val, truth);
+    };
+  }
+  template <typename OP, typename A>
+  CC createBinaryBoolOp(const A &x) {
+    mlir::Location loc = getLoc();
+    mlir::IntegerType i1Ty = builder.getI1Type();
+    auto lf = genarr(x.left());
+    auto rf = genarr(x.right());
+    return [=](IterSpace iters) -> ExtValue {
+      mlir::Value left = fir::getBase(lf(iters));
+      mlir::Value right = fir::getBase(rf(iters));
+      mlir::Value lhs = builder.createConvert(loc, i1Ty, left);
+      mlir::Value rhs = builder.createConvert(loc, i1Ty, right);
+      return builder.create<OP>(loc, lhs, rhs);
+    };
+  }
+  template <typename OP, typename A>
+  CC createCompareBoolOp(mlir::arith::CmpIPredicate pred, const A &x) {
+    mlir::Location loc = getLoc();
+    mlir::IntegerType i1Ty = builder.getI1Type();
+    auto lf = genarr(x.left());
+    auto rf = genarr(x.right());
+    return [=](IterSpace iters) -> ExtValue {
+      mlir::Value left = fir::getBase(lf(iters));
+      mlir::Value right = fir::getBase(rf(iters));
+      mlir::Value lhs = builder.createConvert(loc, i1Ty, left);
+      mlir::Value rhs = builder.createConvert(loc, i1Ty, right);
+      return builder.create<OP>(loc, pred, lhs, rhs);
+    };
   }
-
   template <int KIND>
   CC genarr(const Fortran::evaluate::LogicalOperation<KIND> &x) {
-    TODO(getLoc(), "genarr LogicalOperation");
+    switch (x.logicalOperator) {
+    case Fortran::evaluate::LogicalOperator::And:
+      return createBinaryBoolOp<mlir::arith::AndIOp>(x);
+    case Fortran::evaluate::LogicalOperator::Or:
+      return createBinaryBoolOp<mlir::arith::OrIOp>(x);
+    case Fortran::evaluate::LogicalOperator::Eqv:
+      return createCompareBoolOp<mlir::arith::CmpIOp>(
+          mlir::arith::CmpIPredicate::eq, x);
+    case Fortran::evaluate::LogicalOperator::Neqv:
+      return createCompareBoolOp<mlir::arith::CmpIOp>(
+          mlir::arith::CmpIPredicate::ne, x);
+    case Fortran::evaluate::LogicalOperator::Not:
+      llvm_unreachable(".NOT. handled elsewhere");
+    }
+    llvm_unreachable("unhandled case");
   }
 
   //===--------------------------------------------------------------------===//
   // Relational operators (<, <=, ==, etc.)
   //===--------------------------------------------------------------------===//
 
   template <typename OP, typename PRED, typename A>
   CC createCompareOp(PRED pred, const A &x) {

flang/lib/Lower/IntrinsicCall.cpp

struct IntrinsicLibrary {
   getRuntimeCallGenerator(llvm::StringRef name,
                           mlir::FunctionType soughtFuncType);
 
   /// Lowering for the ABS intrinsic. The ABS intrinsic expects one argument in
   /// the llvm::ArrayRef. The ABS intrinsic is lowered into MLIR/FIR operation
   /// if the argument is an integer, into llvm intrinsics if the argument is
   /// real and to the `hypot` math routine if the argument is of complex type.
   mlir::Value genAbs(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  mlir::Value genAimag(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genAssociated(mlir::Type,
                                    llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genChar(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   template <Extremum, ExtremumBehavior>
   mlir::Value genExtremum(mlir::Type, llvm::ArrayRef<mlir::Value>);
   /// Lowering for the IAND intrinsic. The IAND intrinsic expects two arguments
   /// in the llvm::ArrayRef.
   mlir::Value genIand(mlir::Type, llvm::ArrayRef<mlir::Value>);

 /// one to one mapping with Fortran arguments. If no mapping is
 /// defined here for a generic intrinsic, genRuntimeCall will be called
 /// to look for a match in the runtime a emit a call. Note that the argument
 /// lowering rules for an intrinsic need to be provided only if at least one
 /// argument must not be lowered by value. In which case, the lowering rules
 /// should be provided for all the intrinsic arguments for completeness.
 static constexpr IntrinsicHandler handlers[]{
     {"abs", &I::genAbs},
+    {"aimag", &I::genAimag},
     {"associated",
      &I::genAssociated,
      {{{"pointer", asInquired}, {"target", asInquired}}},
      /*isElemental=*/false},
     {"char", &I::genChar},
     {"iand", &I::genIand},
+    {"min", &I::genExtremum<Extremum::Min, ExtremumBehavior::MinMaxss>},
     {"sum",
      &I::genSum,
      {{{"array", asBox},
        {"dim", asValue},
        {"mask", asBox, handleDynamicOptional}}},
      /*isElemental=*/false},
     {"ubound",
      &I::genUbound,
if (fir::isa_complex(type)) {
     // Use HYPOT to fulfill the no underflow/overflow requirement.
     auto parts = fir::factory::Complex{builder, loc}.extractParts(arg);
     llvm::SmallVector<mlir::Value> args = {parts.first, parts.second};
     return genRuntimeCall("hypot", resultType, args);
   }
   llvm_unreachable("unexpected type in ABS argument");
 }
 
+// AIMAG
+mlir::Value IntrinsicLibrary::genAimag(mlir::Type resultType,
+                                       llvm::ArrayRef<mlir::Value> args) {
+  assert(args.size() == 1);
+  return fir::factory::Complex{builder, loc}.extractComplexPart(
+      args[0], true /* isImagPart */);
+}
+
 // ASSOCIATED
 fir::ExtendedValue
 IntrinsicLibrary::genAssociated(mlir::Type resultType,
                                 llvm::ArrayRef<fir::ExtendedValue> args) {
   assert(args.size() == 2);
   auto *pointer =
       args[0].match([&](const fir::MutableBoxValue &x) { return &x; },
                     [&](const auto &) -> const fir::MutableBoxValue * {

flang/test/Lower/array-elemental-calls-2.f90

module test_ops
   ! CHECK-LABEL: func @_QMtest_opsPcheck_exteremum() {
   subroutine check_exteremum()
     print *,  elem_func(min(i, j))
   ! CHECK:  %[[VAL_0:.*]] = fir.alloca i32
   ! CHECK:  fir.do_loop
   ! CHECK:  %[[VAL_25:.*]] = fir.array_fetch %{{.*}}, %{{.*}} : (!fir.array<10xi32>, index) -> i32
   ! CHECK:  %[[VAL_26:.*]] = fir.array_fetch %{{.*}}, %{{.*}} : (!fir.array<10xi32>, index) -> i32
   ! CHECK:  %[[VAL_27:.*]] = arith.cmpi slt, %[[VAL_25]], %[[VAL_26]] : i32
-  ! CHECK:  %[[VAL_28:.*]] = select %[[VAL_27]], %[[VAL_25]], %[[VAL_26]] : i32
+  ! CHECK:  %[[VAL_28:.*]] = arith.select %[[VAL_27]], %[[VAL_25]], %[[VAL_26]] : i32
   ! CHECK:  fir.store %[[VAL_28]] to %[[VAL_0]] : !fir.ref<i32>
   ! CHECK:  fir.call @_QPelem_func(%[[VAL_0]]) : (!fir.ref<i32>) -> i32
   end subroutine
   
   ! CHECK-LABEL: func @_QMtest_opsPcheck_logical_unary_ops() {
   subroutine check_logical_unary_ops()
     print *,  elem_func_logical(.not.b)
   ! CHECK:  %[[VAL_0:.*]] = fir.alloca !fir.logical<8>