[flang] add hlfir.any intrinsic

Adds a HLFIR operation for the ANY intrinsic according to the
design set out in flang/docs/HighLevel.md

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

7 files changed, 209 insertions, 0 deletions

diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h b/flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h
index 0386b37dbda3..0ccab339fa95 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h
@@ -73,6 +73,7 @@ bool isFortranScalarNumericalType(mlir::Type);
 bool isFortranNumericalArrayObject(mlir::Type);
 bool isFortranNumericalOrLogicalArrayObject(mlir::Type);
 bool isFortranArrayObject(mlir::Type);
+bool isFortranLogicalArrayObject(mlir::Type);
 bool isPassByRefOrIntegerType(mlir::Type);
 bool isI1Type(mlir::Type);
 // scalar i1 or logical, or sequence of logical (via (boxed?) array or expr)
diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td b/flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td
index 2c620af80bfe..5d37bd50c4b5 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td
@@ -133,4 +133,9 @@ def IsMaskArgumentPred
 def AnyFortranLogicalOrI1ArrayObject : Type<IsMaskArgumentPred,
     "A scalar i1 or logical or an array-like object containing logicals">;
 
+def IsFortranLogicalArrayPred
+        : CPred<"::hlfir::isFortranLogicalArrayObject($_self)">;
+def AnyFortranLogicalArrayObject : Type<IsFortranLogicalArrayPred,
+    "any array-like object containing logicals">;
+
 #endif // FORTRAN_DIALECT_HLFIR_OP_BASE
diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index 87136197851a..924e868d32af 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
@@ -317,6 +317,28 @@ def hlfir_ConcatOp : hlfir_Op<"concat", []> {
   let hasVerifier = 1;
 }
 
+def hlfir_AnyOp : hlfir_Op<"any", []> {
+  let summary = "ANY transformational intrinsic";
+  let description = [{
+    Takes a logical array MASK as argument, optionally along a particular dimension,
+    and returns true if any element of MASK is true.
+  }];
+
+  let arguments = (ins
+    AnyFortranLogicalArrayObject:$mask,
+    Optional<AnyIntegerType>:$dim
+  );
+
+  let results = (outs hlfir_ExprType);
+
+  let assemblyFormat = [{
+    $mask  (`dim` $dim^)?  attr-dict `:` functional-type(operands, results)
+  }];
+
+  let hasVerifier = 1;
+}
+
+
 def hlfir_ProductOp : hlfir_Op<"product", [AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<ArithFastMathInterface>]> {
   let summary = "PRODUCT transformational intrinsic";
diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIRDialect.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIRDialect.cpp
index cf6b332028c7..fda12d1ea861 100644
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIRDialect.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIRDialect.cpp
@@ -148,6 +148,17 @@ bool hlfir::isI1Type(mlir::Type type) {
   return false;
 }
 
+bool hlfir::isFortranLogicalArrayObject(mlir::Type type) {
+  if (isBoxAddressType(type))
+    return false;
+  if (auto arrayTy =
+          getFortranElementOrSequenceType(type).dyn_cast<fir::SequenceType>()) {
+    mlir::Type eleTy = arrayTy.getEleTy();
+    return mlir::isa<fir::LogicalType>(eleTy);
+  }
+  return false;
+}
+
 bool hlfir::isMaskArgument(mlir::Type type) {
   if (isBoxAddressType(type))
     return false;
diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 7599b8d2756e..b399a457d13d 100644
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -442,6 +442,44 @@ mlir::LogicalResult hlfir::ParentComponentOp::verify() {
 }
 
 //===----------------------------------------------------------------------===//
+// AnyOp
+//===----------------------------------------------------------------------===//
+mlir::LogicalResult hlfir::AnyOp::verify() {
+  mlir::Operation *op = getOperation();
+
+  auto results = op->getResultTypes();
+  assert(results.size() == 1);
+
+  mlir::Value mask = getMask();
+  fir::SequenceType maskTy =
+      hlfir::getFortranElementOrSequenceType(mask.getType())
+          .cast<fir::SequenceType>();
+  mlir::Type logicalTy = maskTy.getEleTy();
+  llvm::ArrayRef<int64_t> maskShape = maskTy.getShape();
+  hlfir::ExprType resultTy = results[0].cast<hlfir::ExprType>();
+
+  // Result is of the same type as MASK
+  if (resultTy.getElementType() != logicalTy)
+    return emitOpError(
+        "result must have the same element type as MASK argument");
+
+  if (resultTy.isArray()) {
+    // Result is of the same type as MASK
+    if (resultTy.getEleTy() != logicalTy)
+      return emitOpError(
+          "result must have the same element type as MASK argument");
+
+    llvm::ArrayRef<int64_t> resultShape = resultTy.getShape();
+
+    // Result has rank n-1
+    if (resultShape.size() != (maskShape.size() - 1))
+      return emitOpError("result rank must be one less than MASK");
+  }
+
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 
diff --git a/flang/test/HLFIR/any.fir b/flang/test/HLFIR/any.fir
new file mode 100644
index 000000000000..cccbf831f306
--- /dev/null
+++ b/flang/test/HLFIR/any.fir
@@ -0,0 +1,113 @@
+// Test hlfir.product operation parse, verify (no errors), and unparse
+
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+// mask is an expression of known shape
+func.func @any0(%arg0: !hlfir.expr<2x!fir.logical<4>>) {
+  %any = hlfir.any %arg0 : (!hlfir.expr<2x!fir.logical<4>>) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any0(%[[ARRAY:.*]]: !hlfir.expr<2x!fir.logical<4>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!hlfir.expr<2x!fir.logical<4>>) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// mask is an expression of assumed shape
+func.func @any1(%arg0: !hlfir.expr<?x!fir.logical<4>>) {
+  %any = hlfir.any %arg0 : (!hlfir.expr<?x!fir.logical<4>>) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any1(%[[ARRAY:.*]]: !hlfir.expr<?x!fir.logical<4>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!hlfir.expr<?x!fir.logical<4>>) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// mask is a boxed array
+func.func @any2(%arg0: !fir.box<!fir.array<2x!fir.logical<4>>>) {
+  %any = hlfir.any %arg0 : (!fir.box<!fir.array<2x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any2(%[[ARRAY:.*]]: !fir.box<!fir.array<2x!fir.logical<4>>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!fir.box<!fir.array<2x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// mask is an assumed shape boxed array
+func.func @any3(%arg0: !fir.box<!fir.array<?x!fir.logical<4>>>){
+  %any = hlfir.any %arg0 : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any3(%[[ARRAY:.*]]: !fir.box<!fir.array<?x!fir.logical<4>>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// mask is a 2-dimensional array
+func.func @any4(%arg0: !fir.box<!fir.array<?x?x!fir.logical<4>>>){
+  %any = hlfir.any %arg0 : (!fir.box<!fir.array<?x?x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any4(%[[ARRAY:.*]]: !fir.box<!fir.array<?x?x!fir.logical<4>>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!fir.box<!fir.array<?x?x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// mask and dim argument
+func.func @any5(%arg0: !fir.box<!fir.array<?x!fir.logical<4>>>, %arg1: i32) {
+  %any = hlfir.any %arg0 dim %arg1 : (!fir.box<!fir.array<?x!fir.logical<4>>>, i32) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any5(%[[ARRAY:.*]]: !fir.box<!fir.array<?x!fir.logical<4>>>, %[[DIM:.*]]: i32) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] dim %[[DIM]] : (!fir.box<!fir.array<?x!fir.logical<4>>>, i32) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// hlfir.any with dim argument with an unusual type
+func.func @any6(%arg0: !fir.box<!fir.array<?x!fir.logical<4>>>, %arg1: index) {
+  %any = hlfir.any %arg0 dim %arg1 : (!fir.box<!fir.array<?x!fir.logical<4>>>, index) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any6(%[[ARRAY:.*]]: !fir.box<!fir.array<?x!fir.logical<4>>>, %[[DIM:.*]]: index) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] dim %[[DIM]] : (!fir.box<!fir.array<?x!fir.logical<4>>>, index) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// mask is a 2 dimensional array with dim
+func.func @any7(%arg0: !fir.box<!fir.array<?x?x!fir.logical<4>>>, %arg1: i32) {
+  %any = hlfir.any %arg0 dim %arg1 : (!fir.box<!fir.array<?x?x!fir.logical<4>>>, i32) -> !hlfir.expr<?x!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any7(%[[ARRAY:.*]]: !fir.box<!fir.array<?x?x!fir.logical<4>>>, %[[DIM:.*]]: i32) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] dim %[[DIM]] : (!fir.box<!fir.array<?x?x!fir.logical<4>>>, i32) -> !hlfir.expr<?x!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// known shape expr return
+func.func @any8(%arg0: !fir.box<!fir.array<2x2x!fir.logical<4>>>, %arg1: i32) {
+  %any = hlfir.any %arg0 dim %arg1 : (!fir.box<!fir.array<2x2x!fir.logical<4>>>, i32) -> !hlfir.expr<2x!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any8(%[[ARRAY:.*]]: !fir.box<!fir.array<2x2x!fir.logical<4>>>, %[[DIM:.*]]: i32) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] dim %[[DIM]] : (!fir.box<!fir.array<2x2x!fir.logical<4>>>, i32) -> !hlfir.expr<2x!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// hlfir.any with mask argument of ref<array<>> type
+func.func @any9(%arg0: !fir.ref<!fir.array<?x!fir.logical<4>>>) {
+  %any = hlfir.any %arg0 : (!fir.ref<!fir.array<?x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+  return
+}
+// CHECK:      func.func @any9(%[[ARRAY:.*]]: !fir.ref<!fir.array<?x!fir.logical<4>>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!fir.ref<!fir.array<?x!fir.logical<4>>>) -> !hlfir.expr<!fir.logical<4>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+
+// hlfir.any with fir.logical<8> type
+func.func @any10(%arg0: !fir.box<!fir.array<?x!fir.logical<8>>>) {
+  %any = hlfir.any %arg0 : (!fir.box<!fir.array<?x!fir.logical<8>>>) -> !hlfir.expr<!fir.logical<8>>
+  return
+}
+// CHECK:      func.func @any10(%[[ARRAY:.*]]: !fir.box<!fir.array<?x!fir.logical<8>>>) {
+// CHECK-NEXT:   %[[ANY:.*]] = hlfir.any %[[ARRAY]] : (!fir.box<!fir.array<?x!fir.logical<8>>>) -> !hlfir.expr<!fir.logical<8>>
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir
index 6b3ed77eb1aa..e4e9cd00b41b 100644
--- a/flang/test/HLFIR/invalid.fir
+++ b/flang/test/HLFIR/invalid.fir
@@ -297,6 +297,25 @@ func.func @bad_concat_4(%arg0: !fir.ref<!fir.char<1,30>>) {
 }
 
 // -----
+func.func @bad_any1(%arg0: !hlfir.expr<?x!fir.logical<4>>) {
+  // expected-error@+1 {{'hlfir.any' op result must have the same element type as MASK argument}}
+  %0 = hlfir.any %arg0 : (!hlfir.expr<?x!fir.logical<4>>) -> !hlfir.expr<!fir.logical<8>>
+}
+
+// -----
+func.func @bad_any2(%arg0: !hlfir.expr<?x?x!fir.logical<4>>) {
+  // expected-error@+1 {{'hlfir.any' op result must have the same element type as MASK argument}}
+  %0 = hlfir.any %arg0 : (!hlfir.expr<?x?x!fir.logical<4>>) -> !hlfir.expr<?x!fir.logical<8>>
+}
+
+// -----
+func.func @bad_any3(%arg0: !hlfir.expr<?x?x!fir.logical<4>>, %arg1: i32){
+  // expected-error@+1 {{'hlfir.any' op result rank must be one less than MASK}}
+  %0 = hlfir.any %arg0 dim %arg1 : (!hlfir.expr<?x?x!fir.logical<4>>, i32) -> !hlfir.expr<?x?x!fir.logical<4>>
+}
+
+
+// -----
 func.func @bad_product1(%arg0: !hlfir.expr<?xi32>, %arg1: i32, %arg2: !fir.box<!fir.logical<4>>) {
   // expected-error@+1 {{'hlfir.product' op result must have the same element type as ARRAY argument}}
   %0 = hlfir.product %arg0 dim %arg1 mask %arg2 : (!hlfir.expr<?xi32>, i32, !fir.box<!fir.logical<4>>) -> !hlfir.expr<f32>