[flang] Lower F08 FINDLOC intrinsic

The implementation follows the patterns established by the lowering of other
similar intrinsics.

In addition to the code for lowering, the DoTotalReduction template had to be
fixed to correctly break when signaled to do so by the accumulator function.

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

Author: tarunprabhu

flang/include/flang/Optimizer/Builder/Runtime/Reduction.h

@@ -75,6 +75,19 @@ mlir::Value genDotProduct(fir::FirOpBuilder &builder, mlir::Location loc,
                           mlir::Value vectorABox, mlir::Value vectorBBox,
                           mlir::Value resultBox);
 
+/// Generate call to `Findloc` intrinsic runtime routine. This is the version
+/// that does not take a dim argument.
+void genFindloc(fir::FirOpBuilder &builder, mlir::Location loc,
+                mlir::Value resultBox, mlir::Value arrayBox, mlir::Value val,
+                mlir::Value maskBox, mlir::Value kind, mlir::Value back);
+
+/// Generate call to `FindlocDim` intrinsic runtime routine. This is the version
+/// that takes a dim argument.
+void genFindlocDim(fir::FirOpBuilder &builder, mlir::Location loc,
+                   mlir::Value resultBox, mlir::Value arrayBox, mlir::Value val,
+                   mlir::Value dim, mlir::Value maskBox, mlir::Value kind,
+                   mlir::Value back);
+
 /// Generate call to `Maxloc` intrinsic runtime routine. This is the version
 /// that does not take a dim argument.
 void genMaxloc(fir::FirOpBuilder &builder, mlir::Location loc,

flang/lib/Lower/IntrinsicCall.cpp

@@ -516,6 +516,7 @@ struct IntrinsicLibrary {
   mlir::Value genIbits(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIbset(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genIchar(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
+  fir::ExtendedValue genFindloc(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genIeeeIsFinite(mlir::Type, llvm::ArrayRef<mlir::Value>);
   template <mlir::arith::CmpIPredicate pred>
   fir::ExtendedValue genIeeeTypeCompare(mlir::Type,
@@ -801,6 +802,15 @@ static constexpr IntrinsicHandler handlers[]{
      {{{"status", asValue, handleDynamicOptional}}},
      /*isElemental=*/false},
     {"exponent", &I::genExponent},
+    {"findloc",
+     &I::genFindloc,
+     {{{"array", asBox},
+       {"value", asAddr},
+       {"dim", asValue},
+       {"mask", asBox, handleDynamicOptional},
+       {"kind", asValue},
+       {"back", asValue, handleDynamicOptional}}},
+     /*isElemental=*/false},
     {"floor", &I::genFloor},
     {"fraction", &I::genFraction},
     {"get_command",
@@ -3181,6 +3191,98 @@ mlir::Value IntrinsicLibrary::genExponent(mlir::Type resultType,
                                 fir::getBase(args[0])));
 }
 
+// FINDLOC
+fir::ExtendedValue
+IntrinsicLibrary::genFindloc(mlir::Type resultType,
+                             llvm::ArrayRef<fir::ExtendedValue> args) {
+  assert(args.size() == 6);
+
+  llvm::StringRef errMsg = "unexpected result for Findloc";
+
+  // Handle required array argument
+  mlir::Value array = builder.createBox(loc, args[0]);
+  unsigned rank = fir::BoxValue(array).rank();
+  assert(rank >= 1);
+
+  // Handle required value argument
+  mlir::Value val = builder.createBox(loc, args[1]);
+
+  // Check if dim argument is present
+  bool absentDim = isStaticallyAbsent(args[2]);
+
+  // Handle optional mask argument
+  auto mask = isStaticallyAbsent(args[3])
+                  ? builder.create<fir::AbsentOp>(
+                        loc, fir::BoxType::get(builder.getI1Type()))
+                  : builder.createBox(loc, args[3]);
+
+  // Handle optional kind argument
+  auto kind = isStaticallyAbsent(args[4])
+                  ? builder.createIntegerConstant(
+                        loc, builder.getIndexType(),
+                        builder.getKindMap().defaultIntegerKind())
+                  : fir::getBase(args[4]);
+
+  // Handle optional back argument
+  auto back = isStaticallyAbsent(args[5]) ? builder.createBool(loc, false)
+                                          : fir::getBase(args[5]);
+
+  if (!absentDim && rank == 1) {
+    // If dim argument is present and the array is rank 1, then the result is
+    // a scalar (since the the result is rank-1 or 0).
+    // Therefore, we use a scalar result descriptor with FindlocDim().
+    // Create mutable fir.box to be passed to the runtime for the result.
+    fir::MutableBoxValue resultMutableBox =
+        fir::factory::createTempMutableBox(builder, loc, resultType);
+    mlir::Value resultIrBox =
+        fir::factory::getMutableIRBox(builder, loc, resultMutableBox);
+    mlir::Value dim = fir::getBase(args[2]);
+
+    fir::runtime::genFindlocDim(builder, loc, resultIrBox, array, val, dim,
+                                mask, kind, back);
+
+    // Handle cleanup of allocatable result descriptor and return
+    fir::ExtendedValue res =
+        fir::factory::genMutableBoxRead(builder, loc, resultMutableBox);
+    return res.match(
+        [&](const mlir::Value &addr) -> fir::ExtendedValue {
+          addCleanUpForTemp(loc, addr);
+          return builder.create<fir::LoadOp>(loc, resultType, addr);
+        },
+        [&](const auto &) -> fir::ExtendedValue {
+          fir::emitFatalError(loc, errMsg);
+        });
+  }
+
+  // The result will be an array. Create mutable fir.box to be passed to the
+  // runtime for the result.
+  mlir::Type resultArrayType =
+      builder.getVarLenSeqTy(resultType, absentDim ? 1 : rank - 1);
+  fir::MutableBoxValue resultMutableBox =
+      fir::factory::createTempMutableBox(builder, loc, resultArrayType);
+  mlir::Value resultIrBox =
+      fir::factory::getMutableIRBox(builder, loc, resultMutableBox);
+
+  if (absentDim) {
+    fir::runtime::genFindloc(builder, loc, resultIrBox, array, val, mask, kind,
+                             back);
+  } else {
+    mlir::Value dim = fir::getBase(args[2]);
+    fir::runtime::genFindlocDim(builder, loc, resultIrBox, array, val, dim,
+                                mask, kind, back);
+  }
+
+  return fir::factory::genMutableBoxRead(builder, loc, resultMutableBox)
+      .match(
+          [&](const fir::ArrayBoxValue &box) -> fir::ExtendedValue {
+            addCleanUpForTemp(loc, box.getAddr());
+            return box;
+          },
+          [&](const auto &) -> fir::ExtendedValue {
+            fir::emitFatalError(loc, errMsg);
+          });
+}
+
 // FLOOR
 mlir::Value IntrinsicLibrary::genFloor(mlir::Type resultType,
                                        llvm::ArrayRef<mlir::Value> args) {

flang/lib/Optimizer/Builder/Runtime/Reduction.cpp

@@ -576,6 +576,41 @@ void fir::runtime::genCountDim(fir::FirOpBuilder &builder, mlir::Location loc,
   builder.create<fir::CallOp>(loc, func, args);
 }
 
+/// Generate call to `Findloc` intrinsic runtime routine. This is the version
+/// that does not take a dim argument.
+void fir::runtime::genFindloc(fir::FirOpBuilder &builder, mlir::Location loc,
+                              mlir::Value resultBox, mlir::Value arrayBox,
+                              mlir::Value valBox, mlir::Value maskBox,
+                              mlir::Value kind, mlir::Value back) {
+  auto func = fir::runtime::getRuntimeFunc<mkRTKey(Findloc)>(loc, builder);
+  auto fTy = func.getFunctionType();
+  auto sourceFile = fir::factory::locationToFilename(builder, loc);
+  auto sourceLine =
+      fir::factory::locationToLineNo(builder, loc, fTy.getInput(5));
+  auto args = fir::runtime::createArguments(builder, loc, fTy, resultBox,
+                                            arrayBox, valBox, kind, sourceFile,
+                                            sourceLine, maskBox, back);
+  builder.create<fir::CallOp>(loc, func, args);
+}
+
+/// Generate call to `FindlocDim` intrinsic runtime routine. This is the version
+/// that takes a dim argument.
+void fir::runtime::genFindlocDim(fir::FirOpBuilder &builder, mlir::Location loc,
+                                 mlir::Value resultBox, mlir::Value arrayBox,
+                                 mlir::Value valBox, mlir::Value dim,
+                                 mlir::Value maskBox, mlir::Value kind,
+                                 mlir::Value back) {
+  auto func = fir::runtime::getRuntimeFunc<mkRTKey(FindlocDim)>(loc, builder);
+  auto fTy = func.getFunctionType();
+  auto sourceFile = fir::factory::locationToFilename(builder, loc);
+  auto sourceLine =
+      fir::factory::locationToLineNo(builder, loc, fTy.getInput(6));
+  auto args = fir::runtime::createArguments(
+      builder, loc, fTy, resultBox, arrayBox, valBox, kind, dim, sourceFile,
+      sourceLine, maskBox, back);
+  builder.create<fir::CallOp>(loc, func, args);
+}
+
 /// Generate call to `Maxloc` intrinsic runtime routine. This is the version
 /// that does not take a dim argument.
 void fir::runtime::genMaxloc(fir::FirOpBuilder &builder, mlir::Location loc,

flang/runtime/reduction-templates.h

@@ -57,7 +57,8 @@ inline void DoTotalReduction(const Descriptor &x, int dim,
       for (auto elements{x.Elements()}; elements--;
            x.IncrementSubscripts(xAt), mask->IncrementSubscripts(maskAt)) {
         if (IsLogicalElementTrue(*mask, maskAt)) {
-          accumulator.template AccumulateAt<TYPE>(xAt);
+          if (!accumulator.template AccumulateAt<TYPE>(xAt))
+            break;
         }
       }
       return;