[flang-commits] [flang] [flang][hlfir] Optimize MINLOC/MAXLOC with equality mask (PR #211722)
Slava Zakharin via flang-commits
flang-commits at lists.llvm.org
Thu Sep 17 18:54:40 PDT 2026
================
@@ -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 &&
+ isInvariant(cmpOp.getRhs()))
+ return cmpOp.getRhs();
+ if (rhsBase && rhsBase == canonSearch && !lhsBase &&
+ isInvariant(cmpOp.getLhs()))
+ return cmpOp.getLhs();
+ return nullptr;
+}
+
+/// Drop the redundant minMax iter arg from an equality-mask location-search
+/// loop nest.
+static bool transformEqualityMaskMinMaxLoop(fir::DoLoopOp rootLoop,
+ fir::FirOpBuilder &builder) {
+ unsigned numArgs = rootLoop.getInitArgs().size();
+ if (numArgs < 3)
+ return false;
+ if (rootLoop.getFinalValue().value_or(false))
+ return false;
+
+ unsigned isFirstIdx = numArgs - 1;
+ unsigned minMaxIdx = numArgs - 2;
+
+ if (!rootLoop.getResult(minMaxIdx).use_empty())
+ return false;
+
+ // Last iter arg must be i1 initialised to true.
+ if (!rootLoop.getInitArgs()[isFirstIdx].getType().isInteger(1))
+ return false;
+ auto isFirstConst = rootLoop.getInitArgs()[isFirstIdx]
+ .getDefiningOp<mlir::arith::ConstantOp>();
+ if (!isFirstConst)
+ return false;
+ auto boolAttr = mlir::dyn_cast<mlir::BoolAttr>(isFirstConst.getValue());
+ if (!boolAttr || !boolAttr.getValue())
+ return false;
+
+ // Skip if the direct parent is the outer loop of the same nest.
+ mlir::Operation *parent = rootLoop->getParentOp();
+ if (auto parentLoop = mlir::dyn_cast_or_null<fir::DoLoopOp>(parent)) {
+ if (parentLoop.getInitArgs().size() == numArgs) {
+ auto parentIsFirstConst = parentLoop.getInitArgs()[isFirstIdx]
+ .getDefiningOp<mlir::arith::ConstantOp>();
+ if (parentIsFirstConst) {
+ auto parentBoolAttr =
+ mlir::dyn_cast<mlir::BoolAttr>(parentIsFirstConst.getValue());
+ if (parentBoolAttr && parentBoolAttr.getValue())
+ return false;
+ }
+ }
+ }
+
+ // Collect the contiguous loop nest.
+ llvm::SmallVector<fir::DoLoopOp> nest;
+ fir::DoLoopOp curr = rootLoop;
+ while (curr) {
+ nest.push_back(curr);
+ mlir::Block *body = curr.getBody();
+ fir::DoLoopOp inner;
+ for (auto &op : *body) {
+ if (auto nested = mlir::dyn_cast<fir::DoLoopOp>(op)) {
+ if (nested.getInitArgs().size() == numArgs) {
+ inner = nested;
+ break;
+ }
+ }
+ }
+ curr = inner;
+ }
+
+ fir::DoLoopOp innerLoop = nest.back();
+
+ // Find the fir.if inside the innermost loop body.
+ fir::IfOp ifOp;
+ for (auto &op : *innerLoop.getBody()) {
+ if ((ifOp = mlir::dyn_cast<fir::IfOp>(op)))
+ break;
+ }
+ if (!ifOp || ifOp.getThenRegion().empty() || ifOp.getElseRegion().empty())
+ return false;
+
+ mlir::Block *thenBlock = &ifOp.getThenRegion().front();
+ auto thenYield = mlir::dyn_cast<fir::ResultOp>(thenBlock->getTerminator());
+ if (!thenYield || thenYield.getNumOperands() != numArgs)
+ return false;
+
+ // Then-branch must yield false at isFirstIdx.
+ auto thenIsFirstConst =
+ thenYield.getOperand(isFirstIdx).getDefiningOp<mlir::arith::ConstantOp>();
+ if (!thenIsFirstConst)
+ return false;
+ auto thenBoolAttr =
+ mlir::dyn_cast<mlir::BoolAttr>(thenIsFirstConst.getValue());
+ if (!thenBoolAttr || thenBoolAttr.getValue())
+ return false;
+
+ mlir::Block *elseBlock = &ifOp.getElseRegion().front();
+ auto elseYield = mlir::dyn_cast<fir::ResultOp>(elseBlock->getTerminator());
+ if (!elseYield || elseYield.getNumOperands() != numArgs)
+ return false;
+
+ // Else-branch must pass all iter args through unchanged.
+ mlir::ValueRange elseArgs = innerLoop.getRegionIterArgs();
+ for (unsigned i = 0; i < numArgs; ++i)
+ if (elseYield.getOperand(i) != elseArgs[i])
+ return false;
+
+ mlir::Value searchArray = getSearchArray(ifOp, minMaxIdx);
+ if (!searchArray)
+ return false;
+
+ mlir::Value target =
+ getFIREqualityMaskTarget(ifOp.getCondition(), searchArray, rootLoop);
+ if (!target)
+ return false;
+
+ mlir::Location loc = rootLoop.getLoc();
+
+ // Drop minMax from then/else yields; unwrap arith.select for location args.
+ builder.setInsertionPoint(thenYield);
+ llvm::SmallVector<mlir::Value> newThenResults;
+ for (unsigned i = 0; i < numArgs; ++i) {
+ if (i == minMaxIdx)
+ continue;
+ if (i == isFirstIdx) {
+ newThenResults.push_back(thenYield.getOperand(isFirstIdx));
+ } else {
+ mlir::Value val = thenYield.getOperand(i);
+ while (auto sel = val.getDefiningOp<mlir::arith::SelectOp>())
+ val = sel.getTrueValue();
+ newThenResults.push_back(val);
+ }
+ }
+ fir::ResultOp::create(builder, loc, newThenResults);
+ thenYield.erase();
+
+ builder.setInsertionPoint(elseYield);
+ llvm::SmallVector<mlir::Value> newElseResults;
+ for (unsigned i = 0; i < numArgs; ++i) {
+ if (i == minMaxIdx)
+ continue;
+ newElseResults.push_back(elseYield->getOperand(i));
+ }
+ fir::ResultOp::create(builder, loc, newElseResults);
+ elseYield.erase();
+
+ // Replace fir.if condition with (mask AND isFirst).
+ mlir::Value origMaskCond = ifOp.getCondition();
+ if (origMaskCond.getType() != builder.getI1Type()) {
+ while (auto conv = origMaskCond.getDefiningOp<fir::ConvertOp>())
+ origMaskCond = conv.getOperand();
+ if (origMaskCond.getType() != builder.getI1Type()) {
+ builder.setInsertionPoint(ifOp);
+ origMaskCond = fir::ConvertOp::create(builder, loc, builder.getI1Type(),
+ origMaskCond);
+ }
+ }
+ builder.setInsertionPoint(ifOp);
+ mlir::Value isFirstIterArg = innerLoop.getRegionIterArgs()[isFirstIdx];
+ mlir::Value newIfCond =
+ mlir::arith::AndIOp::create(builder, loc, origMaskCond, isFirstIterArg);
+ // Replace fir.if with one that returns reduced result types (no minMax).
+ llvm::SmallVector<mlir::Type> newResultTypes;
+ for (unsigned i = 0; i < numArgs; ++i) {
+ if (i != minMaxIdx)
+ newResultTypes.push_back(ifOp.getResultTypes()[i]);
+ }
+ auto newIfOp = fir::IfOp::create(builder, loc, newResultTypes, newIfCond,
+ /*withElseRegion=*/true);
+ newIfOp.getThenRegion().takeBody(ifOp.getThenRegion());
----------------
vzakhari wrote:
What are the guarantees that then/else regions do not contain something that must be executed on each iteration of the loop? By short-circuiting the whole IF operation via is-first we change the number of executions of those "other" things.
I think to guarantee that this recognition kicks in more general case, we need to create a new IF that clones the min/macloc idiom part and leaves the original IF untouched. DCE should be able to remove any dead code that is left in the original IF.
I guess it is worth adding a test case where there is a call inside then/else, for example.
https://github.com/llvm/llvm-project/pull/211722
More information about the flang-commits
mailing list