[flang-commits] [flang] [flang][hlfir] Optimize MINLOC/MAXLOC with equality mask (PR #211722)

[Slava Zakharin via flang-commits]

================
@@ -0,0 +1,393 @@
+//===- LoopIdiomRecognize.cpp ---------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Eliminate redundant minMax tracking in equality-mask
+// location-search loops after HLFIR-to-FIR lowering and LICM.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/HLFIRTools.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Pass/Pass.h"
+
+namespace fir {
+#define GEN_PASS_DEF_LOOPIDIOMRECOGNIZE
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+namespace {
+
+/// Return the array base reached via array_coor, designate, or apply, or
+/// nullptr.
+static mlir::Value getBase(mlir::Value v) {
+  while (v) {
+    while (auto conv = v.getDefiningOp<fir::ConvertOp>())
+      v = conv.getOperand();
+
+    if (!v)
+      break;
+
+    mlir::Operation *def = v.getDefiningOp();
+    if (!def)
+      break;
+
+    if (auto coor = mlir::dyn_cast<fir::ArrayCoorOp>(def))
+      return coor.getMemref();
+    if (auto desig = mlir::dyn_cast<hlfir::DesignateOp>(def))
+      return desig.getMemref();
+    if (auto apply = mlir::dyn_cast<hlfir::ApplyOp>(def))
+      return apply.getExpr();
+
+    if (auto load = mlir::dyn_cast<fir::LoadOp>(def)) {
+      if (fir::isa_volatile_type(load.getMemref().getType()) ||
+          fir::isa_volatile_type(load.getType()))
+        break;
+      v = load.getMemref();
+      continue;
+    }
+
+    break;
+  }
+  return nullptr;
+}
+
+/// Return the array being searched, extracted from the minMax then-yield.
+static mlir::Value getSearchArray(fir::IfOp ifOp, unsigned minMaxIdx) {
+  mlir::Block &thenBlock = ifOp.getThenRegion().front();
+  auto thenYield = mlir::cast<fir::ResultOp>(thenBlock.getTerminator());
+  mlir::Value minMax = thenYield.getOperand(minMaxIdx);
+
+  while (auto sel = minMax.getDefiningOp<mlir::arith::SelectOp>())
+    minMax = sel.getTrueValue();
+
+  return getBase(minMax);
+}
+
+/// Match arith.cmpi eq where one side loads from searchArray and the other is
+/// loop-invariant.
+static mlir::Value getFIREqualityMaskTarget(mlir::Value cond,
+                                            mlir::Value searchArray,
+                                            fir::DoLoopOp rootLoop) {
+  while (auto conv = cond.getDefiningOp<fir::ConvertOp>())
+    cond = conv.getOperand();
+
+  auto cmpOp = cond.getDefiningOp<mlir::arith::CmpIOp>();
+  if (!cmpOp || cmpOp.getPredicate() != mlir::arith::CmpIPredicate::eq)
+    return nullptr;
+
+  auto canonicalizeBase = [](mlir::Value v) -> mlir::Value {
+    while (v) {
+      mlir::Operation *def = v.getDefiningOp();
+      if (!def)
+        break;
+      if (auto decl = mlir::dyn_cast<hlfir::DeclareOp>(def))
+        v = decl.getMemref();
+      else if (mlir::isa<fir::ConvertOp, hlfir::AsExprOp>(def))
+        v = def->getOperand(0);
+      else
+        break;
+    }
+    return v;
+  };
+
+  mlir::Value canonSearch = canonicalizeBase(searchArray);
+  mlir::Value lhsBase = canonicalizeBase(getBase(cmpOp.getLhs()));
+  mlir::Value rhsBase = canonicalizeBase(getBase(cmpOp.getRhs()));
+
+  auto isInvariant = [&](mlir::Value v) -> bool {
+    if (mlir::Operation *def = v.getDefiningOp())
+      return !rootLoop->isAncestor(def);
+    if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(v)) {
+      mlir::Operation *ownerOp = blockArg.getOwner()->getParentOp();
+      return !rootLoop->isAncestor(ownerOp) &&
+             (ownerOp != rootLoop.getOperation());
+    }
+    return true;
+  };
+
+  if (lhsBase && lhsBase == canonSearch && !rhsBase &&
----------------
vzakhari wrote:

I do not think this guarantees that the min/max select and the comparison are actually using the same value loaded from the array. The base might be the same, but indices might be different. I am not sure why you need all this complexity with handling the array base. All we care is the value produced by the load.

https://github.com/llvm/llvm-project/pull/211722