[flang] Lower elemental calls

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.

Reviewed By: PeteSteinfeld

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

Co-authored-by: Jean Perier <jperier@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

flang/lib/Lower/ConvertExpr.cpp

@@ -4031,20 +4031,41 @@ public:
 
   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) {
@@ -4057,7 +4078,11 @@ public:
   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)));
+    };
   }
 
   //===--------------------------------------------------------------------===//
@@ -4100,7 +4125,15 @@ public:
   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(
@@ -4909,14 +4942,67 @@ public:
     TODO(getLoc(), "genarr StructureConstructor");
   }
 
-  template <int KIND>
+  //===--------------------------------------------------------------------===//
-  CC genarr(const Fortran::evaluate::Not<KIND> &x) {
+  // LOCICAL operators (.NOT., .AND., .EQV., etc.)
-    TODO(getLoc(), "genarr Not");
+  //===--------------------------------------------------------------------===//
-  }
 
+  template <int KIND>
+  CC genarr(const Fortran::evaluate::Not<KIND> &x) {
+    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");
   }
 
   //===--------------------------------------------------------------------===//

flang/lib/Lower/IntrinsicCall.cpp

@@ -234,6 +234,7 @@ struct IntrinsicLibrary {
   /// 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>);
@@ -333,12 +334,14 @@ static constexpr bool handleDynamicOptional = true;
 /// 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},
@@ -1056,6 +1059,14 @@ mlir::Value IntrinsicLibrary::genAbs(mlir::Type resultType,
   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,

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

@@ -81,7 +81,7 @@ module test_ops
   ! 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