[fir] Add affine promotion pass

[clementval]

Summary

Convert fir operations which satisfy affine constraints to the affine
dialect.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Jean Perier <[e-mail]>
Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: V Donaldson <[e-mail]>
Co-authored-by: Rajan Walia <[e-mail]>
Co-authored-by: Sourabh Singh Tomar <[e-mail]>

[awarzynski]

Could you add a test for affine.if? How about some negative tests - not mission critical, just nice-to-have.

[clementval]

In D111155#3043356, @awarzynski wrote:

Could you add a test for affine.if? How about some negative tests - not mission critical, just nice-to-have.

I'm not sure what you have in mind for negative test in a transformation?

[awarzynski]

In D111155#3043512, @clementval wrote:

In D111155#3043356, @awarzynski wrote:

Could you add a test for affine.if? How about some negative tests - not mission critical, just nice-to-have.

I'm not sure what you have in mind for negative test in a transformation?

E.g. a test with a fir.do_loop that should not be transformed and indeed isn't.

[clementval]

In D111155#3043586, @awarzynski wrote:

In D111155#3043512, @clementval wrote:

In D111155#3043356, @awarzynski wrote:

Could you add a test for affine.if? How about some negative tests - not mission critical, just nice-to-have.

I'm not sure what you have in mind for negative test in a transformation?

E.g. a test with a fir.do_loop that should not be transformed and indeed isn't.

Ok. Yeah I can add some.

[clementval]

Add tests

[awarzynski]

flang/test/Fir/affine-promotion.fir

2

This comment suggests that tests for fir.if and fir.load should go to a dedicated file.

94

Would a test with fir.if _outside of_ fir.do_loop be possible? I think that it would nicely highlight _what_ is being tested. Currently both versions of @calc look like tests for fir.do_loop specifically (that's the first thing that you notice). fir.if/affine.if` or fir.load/affine.store are not obvious to spot.

[clementval]

flang/test/Fir/affine-promotion.fir

2

I'll update the comment so both tests can be in the same file.

94

a fir.if outside of loop would never produce affine operations. affine.if has meaning in a loop only.

[awarzynski]

LGTM. Thanks for adding the tests!

flang/test/Fir/affine-promotion.fir

9

[nit] Could you rename calc as loop_with_load_and_store or something similar? So that it is clear which case is being tested?

82

[nit] Could you rename calc as loop_with_if or something similar? So that it is clear which case is being tested?

94

Ah, indeed:

The affine.if operation restricts execution to a subset of the loop iteration space defined by an integer set (a conjunction of affine constraints).

From : https://mlir.llvm.org/docs/Dialects/Affine/#affineif-mliraffineifop. I missed that, thanks!

[clementval]

Address nit comment

[clementval]

Rebase

[jdoerfert]

Again, glancing at these transformations makes me doubt how well they are tested. The documentation doesn't seem to indicate this is a WIP or not ready yet.
I only read over parts and found 3 things that I suspect are "broken", for some definition of the word. (The multi symbol stuff is technically not a correctness issue.)

I'm also guessing we ignore overflows and such completely when we do these things, right?

flang/lib/Optimizer/Transforms/AffinePromotion.cpp

204

This way of translating to "affine" should break as soon as lhs or rhs is an induction variable and the other one is not a constant.
Same for other binary operations. I didn't see a test for this but I would hope something in affine will complain.

215

This assumes perfectly nested loops, right? Is this checked?
I doubt the symCount stuff works if a symbol is used multiple times.

[clementval]

In D111155#3053414, @jdoerfert wrote:

Again, glancing at these transformations makes me doubt how well they are tested. The documentation doesn't seem to indicate this is a WIP or not ready yet.
I only read over parts and found 3 things that I suspect are "broken", for some definition of the word. (The multi symbol stuff is technically not a correctness issue.)

I'm also guessing we ignore overflows and such completely when we do these things, right?

As for the patch related to affine promotion/demotion these are prototype. I would be happy to add a mention in the comment about that. Alternatively, we can also not upstream these transformations.

[clementval]

Comment added with patch D111629

flang/include/flang/Optimizer/Transforms/Passes.h

 class BlockAndValueMapping;
 class Operation;
 class Pass;
 class Region;
 } // namespace mlir
 
 namespace fir {
 
+std::unique_ptr<mlir::Pass> createPromoteToAffinePass();
 std::unique_ptr<mlir::Pass> createExternalNameConversionPass();
 
 /// Support for inlining on FIR.
 bool canLegallyInline(mlir::Operation *op, mlir::Region *reg,
                       mlir::BlockAndValueMapping &map);
 
 // declarative passes
 #define GEN_PASS_REGISTRATION
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 
 } // namespace fir
 
 #endif // OPTIMIZER_TRANSFORMS_PASSES_H

flang/include/flang/Optimizer/Transforms/Passes.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FLANG_OPTIMIZER_TRANSFORMS_PASSES
 #define FLANG_OPTIMIZER_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
+def AffineDialectPromotion : FunctionPass<"promote-to-affine"> {
+  let summary = "Promotes `fir.{do_loop,if}` to `affine.{for,if}`.";
+  let description = [{
+    Convert fir operations which satisfy affine constraints to the affine
+    dialect.
+
+    `fir.do_loop` will be converted to `affine.for` if the loops inside the body
+    can be converted and the indices for memory loads and stores satisfy
+    `affine.apply` criteria for symbols and dimensions.
+
+    `fir.if` will be converted to `affine.if` where possible. `affine.if`'s
+    condition uses an integer set (==, >=) and an analysis is done to determine
+    the fir condition's parent operations to construct the integer set.
+
+    `fir.load` (`fir.store`) will be converted to `affine.load` (`affine.store`)
+    where possible. This conversion includes adding a dummy `fir.convert` cast
+    to adapt values of type `!fir.ref<!fir.array>` to `memref`. This is done
+    because the affine dialect presently only understands the `memref` type.
+  }];
+  let constructor = "::fir::createPromoteToAffinePass()";
+  let dependentDialects = [
+    "fir::FIROpsDialect", "mlir::StandardOpsDialect", "mlir::AffineDialect"
+  ];
+}
+
 def ExternalNameConversion : Pass<"external-name-interop", "mlir::ModuleOp"> {
   let summary = "Convert name for external interoperability";
   let description = [{
     Demangle FIR internal name and mangle them for external interoperability.
   }];
   let constructor = "::fir::createExternalNameConversionPass()";
 }
 
 #endif // FLANG_OPTIMIZER_TRANSFORMS_PASSES

flang/lib/Optimizer/Transforms/AffinePromotion.cpp

+//===-- AffinePromotion.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "flang-affine-promotion"
+
+using namespace fir;
+
+namespace {
+struct AffineLoopAnalysis;
+struct AffineIfAnalysis;
+
+/// Stores analysis objects for all loops and if operations inside a function
+///  these analysis are used twice, first for marking operations for rewrite and
+///  second when doing rewrite.
+struct AffineFunctionAnalysis {
+  explicit AffineFunctionAnalysis(mlir::FuncOp funcOp) {
+    for (fir::DoLoopOp op : funcOp.getOps<fir::DoLoopOp>())
+      loopAnalysisMap.try_emplace(op, op, *this);
+  }
+
+  AffineLoopAnalysis getChildLoopAnalysis(fir::DoLoopOp op) const;
+
+  AffineIfAnalysis getChildIfAnalysis(fir::IfOp op) const;
+
+  llvm::DenseMap<mlir::Operation *, AffineLoopAnalysis> loopAnalysisMap;
+  llvm::DenseMap<mlir::Operation *, AffineIfAnalysis> ifAnalysisMap;
+};
+} // namespace
+
+static bool analyzeCoordinate(mlir::Value coordinate, mlir::Operation *op) {
+  if (auto blockArg = coordinate.dyn_cast<mlir::BlockArgument>()) {
+    if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()))
+      return true;
+    LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: array coordinate is not a "
+                               "loop induction variable (owner not loopOp)\n";
+               op->dump());
+    return false;
+  }
+  LLVM_DEBUG(
+      llvm::dbgs() << "AffineLoopAnalysis: array coordinate is not a loop "
+                      "induction variable (not a block argument)\n";
+      op->dump(); coordinate.getDefiningOp()->dump());
+  return false;
+}
+
+namespace {
+struct AffineLoopAnalysis {
+  AffineLoopAnalysis() = default;
+
+  explicit AffineLoopAnalysis(fir::DoLoopOp op, AffineFunctionAnalysis &afa)
+      : legality(analyzeLoop(op, afa)) {}
+
+  bool canPromoteToAffine() { return legality; }
+
+private:
+  bool analyzeBody(fir::DoLoopOp loopOperation,
+                   AffineFunctionAnalysis &functionAnalysis) {
+    for (auto loopOp : loopOperation.getOps<fir::DoLoopOp>()) {
+      auto analysis = functionAnalysis.loopAnalysisMap
+                          .try_emplace(loopOp, loopOp, functionAnalysis)
+                          .first->getSecond();
+      if (!analysis.canPromoteToAffine())
+        return false;
+    }
+    for (auto ifOp : loopOperation.getOps<fir::IfOp>())
+      functionAnalysis.ifAnalysisMap.try_emplace(ifOp, ifOp, functionAnalysis);
+    return true;
+  }
+
+  bool analyzeLoop(fir::DoLoopOp loopOperation,
+                   AffineFunctionAnalysis &functionAnalysis) {
+    LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: \n"; loopOperation.dump(););
+    return analyzeMemoryAccess(loopOperation) &&
+           analyzeBody(loopOperation, functionAnalysis);
+  }
+
+  bool analyzeReference(mlir::Value memref, mlir::Operation *op) {
+    if (auto acoOp = memref.getDefiningOp<ArrayCoorOp>()) {
+      if (acoOp.memref().getType().isa<fir::BoxType>()) {
+        // TODO: Look if and how fir.box can be promoted to affine.
+        LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: cannot promote loop, "
+                                   "array memory operation uses fir.box\n";
+                   op->dump(); acoOp.dump(););
+        return false;
+      }
+      bool canPromote = true;
+      for (auto coordinate : acoOp.indices())
+        canPromote = canPromote && analyzeCoordinate(coordinate, op);
+      return canPromote;
+    }
+    if (auto coOp = memref.getDefiningOp<CoordinateOp>()) {
+      LLVM_DEBUG(llvm::dbgs()
+                     << "AffineLoopAnalysis: cannot promote loop, "
+                        "array memory operation uses non ArrayCoorOp\n";
+                 op->dump(); coOp.dump(););
+
+      return false;
+    }
+    LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: unknown type of memory "
+                               "reference for array load\n";
+               op->dump(););
+    return false;
+  }
+
+  bool analyzeMemoryAccess(fir::DoLoopOp loopOperation) {
+    for (auto loadOp : loopOperation.getOps<fir::LoadOp>())
+      if (!analyzeReference(loadOp.memref(), loadOp))
+        return false;
+    for (auto storeOp : loopOperation.getOps<fir::StoreOp>())
+      if (!analyzeReference(storeOp.memref(), storeOp))
+        return false;
+    return true;
+  }
+
+  bool legality{};
+};
+} // namespace
+
+AffineLoopAnalysis
+AffineFunctionAnalysis::getChildLoopAnalysis(fir::DoLoopOp op) const {
+  auto it = loopAnalysisMap.find_as(op);
+  if (it == loopAnalysisMap.end()) {
+    LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n";
+               op.dump(););
+    op.emitError("error in fetching loop analysis in AffineFunctionAnalysis\n");
+    return {};
+  }
+  return it->getSecond();
+}
+
+namespace {
+/// Calculates arguments for creating an IntegerSet. symCount, dimCount are the
+/// final number of symbols and dimensions of the affine map. Integer set if
+/// possible is in Optional IntegerSet.
+struct AffineIfCondition {
+  using MaybeAffineExpr = llvm::Optional<mlir::AffineExpr>;
+
+  explicit AffineIfCondition(mlir::Value fc) : firCondition(fc) {
+    if (auto condDef = firCondition.getDefiningOp<mlir::CmpIOp>())
+      fromCmpIOp(condDef);
+  }
+
+  bool hasIntegerSet() const { return integerSet.hasValue(); }
+
+  mlir::IntegerSet getIntegerSet() const {
+    assert(hasIntegerSet() && "integer set is missing");
+    return integerSet.getValue();
+  }
+
+  mlir::ValueRange getAffineArgs() const { return affineArgs; }
+
+private:
+  MaybeAffineExpr affineBinaryOp(mlir::AffineExprKind kind, mlir::Value lhs,
+                                 mlir::Value rhs) {
+    return affineBinaryOp(kind, toAffineExpr(lhs), toAffineExpr(rhs));
+  }
+
+  MaybeAffineExpr affineBinaryOp(mlir::AffineExprKind kind, MaybeAffineExpr lhs,
+                                 MaybeAffineExpr rhs) {
+    if (lhs.hasValue() && rhs.hasValue())
+      return mlir::getAffineBinaryOpExpr(kind, lhs.getValue(), rhs.getValue());
+    return {};
+  }
+
+  MaybeAffineExpr toAffineExpr(MaybeAffineExpr e) { return e; }
+
+  MaybeAffineExpr toAffineExpr(int64_t value) {
+    return {mlir::getAffineConstantExpr(value, firCondition.getContext())};
+  }
+
+  /// Returns an AffineExpr if it is a result of operations that can be done
+  /// in an affine expression, this includes -, +, *, rem, constant.
+  /// block arguments of a loopOp or forOp are used as dimensions
+  MaybeAffineExpr toAffineExpr(mlir::Value value) {
+    if (auto op = value.getDefiningOp<mlir::SubIOp>())
+      return affineBinaryOp(mlir::AffineExprKind::Add, toAffineExpr(op.lhs()),
+                            affineBinaryOp(mlir::AffineExprKind::Mul,
+                                           toAffineExpr(op.rhs()),
+                                           toAffineExpr(-1)));
+    if (auto op = value.getDefiningOp<mlir::AddIOp>())
+      return affineBinaryOp(mlir::AffineExprKind::Add, op.lhs(), op.rhs());
+    if (auto op = value.getDefiningOp<mlir::MulIOp>())
+      return affineBinaryOp(mlir::AffineExprKind::Mul, op.lhs(), op.rhs());
+    if (auto op = value.getDefiningOp<mlir::UnsignedRemIOp>())
+      return affineBinaryOp(mlir::AffineExprKind::Mod, op.lhs(), op.rhs());
+    if (auto op = value.getDefiningOp<mlir::ConstantOp>())
+      if (auto intConstant = op.getValue().dyn_cast<IntegerAttr>())
+        return toAffineExpr(intConstant.getInt());
+    if (auto blockArg = value.dyn_cast<mlir::BlockArgument>()) {
+      affineArgs.push_back(value);
+      if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()) ||
+          isa<mlir::AffineForOp>(blockArg.getOwner()->getParentOp()))
+        return {mlir::getAffineDimExpr(dimCount++, value.getContext())};
+      return {mlir::getAffineSymbolExpr(symCount++, value.getContext())};
+    }
+    return {};
+  }
+
+  void fromCmpIOp(mlir::CmpIOp cmpOp) {
+    auto lhsAffine = toAffineExpr(cmpOp.lhs());
+    auto rhsAffine = toAffineExpr(cmpOp.rhs());
+    if (!lhsAffine.hasValue() || !rhsAffine.hasValue())
+      return;
+    auto constraintPair = constraint(
+        cmpOp.predicate(), rhsAffine.getValue() - lhsAffine.getValue());
+    if (!constraintPair)
+      return;
+    integerSet = mlir::IntegerSet::get(dimCount, symCount,
+                                       {constraintPair.getValue().first},
+                                       {constraintPair.getValue().second});
+    return;
+  }
+
+  llvm::Optional<std::pair<AffineExpr, bool>>
+  constraint(mlir::CmpIPredicate predicate, mlir::AffineExpr basic) {
+    switch (predicate) {
+    case mlir::CmpIPredicate::slt:
+      return {std::make_pair(basic - 1, false)};
+    case mlir::CmpIPredicate::sle:
+      return {std::make_pair(basic, false)};
+    case mlir::CmpIPredicate::sgt:
+      return {std::make_pair(1 - basic, false)};
+    case mlir::CmpIPredicate::sge:
+      return {std::make_pair(0 - basic, false)};
+    case mlir::CmpIPredicate::eq:
+      return {std::make_pair(basic, true)};
+    default:
+      return {};
+    }
+  }
+
+  llvm::SmallVector<mlir::Value> affineArgs;
+  llvm::Optional<mlir::IntegerSet> integerSet;
+  mlir::Value firCondition;
+  unsigned symCount{0u};
+  unsigned dimCount{0u};
+};
+} // namespace
+
+namespace {
+/// Analysis for affine promotion of fir.if
+struct AffineIfAnalysis {
+  AffineIfAnalysis() = default;
+
+  explicit AffineIfAnalysis(fir::IfOp op, AffineFunctionAnalysis &afa)
+      : legality(analyzeIf(op, afa)) {}
+
+  bool canPromoteToAffine() { return legality; }
+
+private:
+  bool analyzeIf(fir::IfOp op, AffineFunctionAnalysis &afa) {
+    if (op.getNumResults() == 0)
+      return true;
+    LLVM_DEBUG(llvm::dbgs()
+                   << "AffineIfAnalysis: not promoting as op has results\n";);
+    return false;
+  }
+
+  bool legality{};
+};
+} // namespace
+
+AffineIfAnalysis
+AffineFunctionAnalysis::getChildIfAnalysis(fir::IfOp op) const {
+  auto it = ifAnalysisMap.find_as(op);
+  if (it == ifAnalysisMap.end()) {
+    LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n";
+               op.dump(););
+    op.emitError("error in fetching if analysis in AffineFunctionAnalysis\n");
+    return {};
+  }
+  return it->getSecond();
+}
+
+/// AffineMap rewriting fir.array_coor operation to affine apply,
+/// %dim = fir.gendim %lowerBound, %upperBound, %stride
+/// %a = fir.array_coor %arr(%dim) %i
+/// returning affineMap = affine_map<(i)[lb, ub, st] -> (i*st - lb)>
+static mlir::AffineMap createArrayIndexAffineMap(unsigned dimensions,
+                                                 MLIRContext *context) {
+  auto index = mlir::getAffineConstantExpr(0, context);
+  auto accuExtent = mlir::getAffineConstantExpr(1, context);
+  for (unsigned i = 0; i < dimensions; ++i) {
+    mlir::AffineExpr idx = mlir::getAffineDimExpr(i, context),
+                     lowerBound = mlir::getAffineSymbolExpr(i * 3, context),
+                     currentExtent =
+                         mlir::getAffineSymbolExpr(i * 3 + 1, context),
+                     stride = mlir::getAffineSymbolExpr(i * 3 + 2, context),
+                     currentPart = (idx * stride - lowerBound) * accuExtent;
+    index = currentPart + index;
+    accuExtent = accuExtent * currentExtent;
+  }
+  return mlir::AffineMap::get(dimensions, dimensions * 3, index);
+}
+
+static Optional<int64_t> constantIntegerLike(const mlir::Value value) {
+  if (auto definition = value.getDefiningOp<ConstantOp>())
+    if (auto stepAttr = definition.getValue().dyn_cast<IntegerAttr>())
+      return stepAttr.getInt();
+  return {};
+}
+
+static mlir::Type coordinateArrayElement(fir::ArrayCoorOp op) {
+  if (auto refType = op.memref().getType().dyn_cast_or_null<ReferenceType>()) {
+    if (auto seqType = refType.getEleTy().dyn_cast_or_null<SequenceType>()) {
+      return seqType.getEleTy();
+    }
+  }
+  op.emitError(
+      "AffineLoopConversion: array type in coordinate operation not valid\n");
+  return mlir::Type();
+}
+
+static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeOp shape,
+                              SmallVectorImpl<mlir::Value> &indexArgs,
+                              mlir::PatternRewriter &rewriter) {
+  auto one = rewriter.create<mlir::ConstantOp>(
+      acoOp.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(1));
+  auto extents = shape.extents();
+  for (auto i = extents.begin(); i < extents.end(); i++) {
+    indexArgs.push_back(one);
+    indexArgs.push_back(*i);
+    indexArgs.push_back(one);
+  }
+}
+
+static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeShiftOp shape,
+                              SmallVectorImpl<mlir::Value> &indexArgs,
+                              mlir::PatternRewriter &rewriter) {
+  auto one = rewriter.create<mlir::ConstantOp>(
+      acoOp.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(1));
+  auto extents = shape.pairs();
+  for (auto i = extents.begin(); i < extents.end();) {
+    indexArgs.push_back(*i++);
+    indexArgs.push_back(*i++);
+    indexArgs.push_back(one);
+  }
+}
+
+static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::SliceOp slice,
+                              SmallVectorImpl<mlir::Value> &indexArgs,
+                              mlir::PatternRewriter &rewriter) {
+  auto extents = slice.triples();
+  for (auto i = extents.begin(); i < extents.end();) {
+    indexArgs.push_back(*i++);
+    indexArgs.push_back(*i++);
+    indexArgs.push_back(*i++);
+  }
+}
+
+static void populateIndexArgs(fir::ArrayCoorOp acoOp,
+                              SmallVectorImpl<mlir::Value> &indexArgs,
+                              mlir::PatternRewriter &rewriter) {
+  if (auto shape = acoOp.shape().getDefiningOp<ShapeOp>())
+    return populateIndexArgs(acoOp, shape, indexArgs, rewriter);
+  if (auto shapeShift = acoOp.shape().getDefiningOp<ShapeShiftOp>())
+    return populateIndexArgs(acoOp, shapeShift, indexArgs, rewriter);
+  if (auto slice = acoOp.shape().getDefiningOp<SliceOp>())
+    return populateIndexArgs(acoOp, slice, indexArgs, rewriter);
+  return;
+}
+
+/// Returns affine.apply and fir.convert from array_coor and gendims
+static std::pair<mlir::AffineApplyOp, fir::ConvertOp>
+createAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter) {
+  auto acoOp = arrayRef.getDefiningOp<ArrayCoorOp>();
+  auto affineMap =
+      createArrayIndexAffineMap(acoOp.indices().size(), acoOp.getContext());
+  SmallVector<mlir::Value> indexArgs;
+  indexArgs.append(acoOp.indices().begin(), acoOp.indices().end());
+
+  populateIndexArgs(acoOp, indexArgs, rewriter);
+
+  auto affineApply = rewriter.create<mlir::AffineApplyOp>(acoOp.getLoc(),
+                                                          affineMap, indexArgs);
+  auto arrayElementType = coordinateArrayElement(acoOp);
+  auto newType = mlir::MemRefType::get({-1}, arrayElementType);
+  auto arrayConvert =
+      rewriter.create<fir::ConvertOp>(acoOp.getLoc(), newType, acoOp.memref());
+  return std::make_pair(affineApply, arrayConvert);
+}
+
+static void rewriteLoad(fir::LoadOp loadOp, mlir::PatternRewriter &rewriter) {
+  rewriter.setInsertionPoint(loadOp);
+  auto affineOps = createAffineOps(loadOp.memref(), rewriter);
+  rewriter.replaceOpWithNewOp<mlir::AffineLoadOp>(
+      loadOp, affineOps.second.getResult(), affineOps.first.getResult());
+}
+
+static void rewriteStore(fir::StoreOp storeOp,
+                         mlir::PatternRewriter &rewriter) {
+  rewriter.setInsertionPoint(storeOp);
+  auto affineOps = createAffineOps(storeOp.memref(), rewriter);
+  rewriter.replaceOpWithNewOp<mlir::AffineStoreOp>(storeOp, storeOp.value(),
+                                                   affineOps.second.getResult(),
+                                                   affineOps.first.getResult());
+}
+
+static void rewriteMemoryOps(Block *block, mlir::PatternRewriter &rewriter) {
+  for (auto &bodyOp : block->getOperations()) {
+    if (isa<fir::LoadOp>(bodyOp))
+      rewriteLoad(cast<fir::LoadOp>(bodyOp), rewriter);
+    if (isa<fir::StoreOp>(bodyOp))
+      rewriteStore(cast<fir::StoreOp>(bodyOp), rewriter);
+  }
+}
+
+namespace {
+/// Convert `fir.do_loop` to `affine.for`, creates fir.convert for arrays to
+/// memref, rewrites array_coor to affine.apply with affine_map. Rewrites fir
+/// loads and stores to affine.
+class AffineLoopConversion : public mlir::OpRewritePattern<fir::DoLoopOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  AffineLoopConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa)
+      : OpRewritePattern(context), functionAnalysis(afa) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::DoLoopOp loop,
+                  mlir::PatternRewriter &rewriter) const override {
+    LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: rewriting loop:\n";
+               loop.dump(););
+    LLVM_ATTRIBUTE_UNUSED auto loopAnalysis =
+        functionAnalysis.getChildLoopAnalysis(loop);
+    auto &loopOps = loop.getBody()->getOperations();
+    auto loopAndIndex = createAffineFor(loop, rewriter);
+    auto affineFor = loopAndIndex.first;
+    auto inductionVar = loopAndIndex.second;
+
+    rewriter.startRootUpdate(affineFor.getOperation());
+    affineFor.getBody()->getOperations().splice(
+        std::prev(affineFor.getBody()->end()), loopOps, loopOps.begin(),
+        std::prev(loopOps.end()));
+    rewriter.finalizeRootUpdate(affineFor.getOperation());
+
+    rewriter.startRootUpdate(loop.getOperation());
+    loop.getInductionVar().replaceAllUsesWith(inductionVar);
+    rewriter.finalizeRootUpdate(loop.getOperation());
+
+    rewriteMemoryOps(affineFor.getBody(), rewriter);
+
+    LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: loop rewriten to:\n";
+               affineFor.dump(););
+    rewriter.replaceOp(loop, affineFor.getOperation()->getResults());
+    return success();
+  }
+
+private:
+  std::pair<mlir::AffineForOp, mlir::Value>
+  createAffineFor(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {
+    if (auto constantStep = constantIntegerLike(op.step()))
+      if (constantStep.getValue() > 0)
+        return positiveConstantStep(op, constantStep.getValue(), rewriter);
+    return genericBounds(op, rewriter);
+  }
+
+  // when step for the loop is positive compile time constant
+  std::pair<mlir::AffineForOp, mlir::Value>
+  positiveConstantStep(fir::DoLoopOp op, int64_t step,
+                       mlir::PatternRewriter &rewriter) const {
+    auto affineFor = rewriter.create<mlir::AffineForOp>(
+        op.getLoc(), ValueRange(op.lowerBound()),
+        mlir::AffineMap::get(0, 1,
+                             mlir::getAffineSymbolExpr(0, op.getContext())),
+        ValueRange(op.upperBound()),
+        mlir::AffineMap::get(0, 1,
+                             1 + mlir::getAffineSymbolExpr(0, op.getContext())),
+        step);
+    return std::make_pair(affineFor, affineFor.getInductionVar());
+  }
+
+  std::pair<mlir::AffineForOp, mlir::Value>
+  genericBounds(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {
+    auto lowerBound = mlir::getAffineSymbolExpr(0, op.getContext());
+    auto upperBound = mlir::getAffineSymbolExpr(1, op.getContext());
+    auto step = mlir::getAffineSymbolExpr(2, op.getContext());
+    mlir::AffineMap upperBoundMap = mlir::AffineMap::get(
+        0, 3, (upperBound - lowerBound + step).floorDiv(step));
+    auto genericUpperBound = rewriter.create<mlir::AffineApplyOp>(
+        op.getLoc(), upperBoundMap,
+        ValueRange({op.lowerBound(), op.upperBound(), op.step()}));
+    auto actualIndexMap = mlir::AffineMap::get(
+        1, 2,
+        (lowerBound + mlir::getAffineDimExpr(0, op.getContext())) *
+            mlir::getAffineSymbolExpr(1, op.getContext()));
+
+    auto affineFor = rewriter.create<mlir::AffineForOp>(
+        op.getLoc(), ValueRange(),
+        AffineMap::getConstantMap(0, op.getContext()),
+        genericUpperBound.getResult(),
+        mlir::AffineMap::get(0, 1,
+                             1 + mlir::getAffineSymbolExpr(0, op.getContext())),
+        1);
+    rewriter.setInsertionPointToStart(affineFor.getBody());
+    auto actualIndex = rewriter.create<mlir::AffineApplyOp>(
+        op.getLoc(), actualIndexMap,
+        ValueRange({affineFor.getInductionVar(), op.lowerBound(), op.step()}));
+    return std::make_pair(affineFor, actualIndex.getResult());
+  }
+
+  AffineFunctionAnalysis &functionAnalysis;
+};
+
+/// Convert `fir.if` to `affine.if`.
+class AffineIfConversion : public mlir::OpRewritePattern<fir::IfOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  AffineIfConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa)
+      : OpRewritePattern(context) {}
+  mlir::LogicalResult
+  matchAndRewrite(fir::IfOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: rewriting if:\n";
+               op.dump(););
+    auto &ifOps = op.thenRegion().front().getOperations();
+    auto affineCondition = AffineIfCondition(op.condition());
+    if (!affineCondition.hasIntegerSet()) {
+      LLVM_DEBUG(
+          llvm::dbgs()
+              << "AffineIfConversion: couldn't calculate affine condition\n";);
+      return failure();
+    }
+    auto affineIf = rewriter.create<mlir::AffineIfOp>(
+        op.getLoc(), affineCondition.getIntegerSet(),
+        affineCondition.getAffineArgs(), !op.elseRegion().empty());
+    rewriter.startRootUpdate(affineIf);
+    affineIf.getThenBlock()->getOperations().splice(
+        std::prev(affineIf.getThenBlock()->end()), ifOps, ifOps.begin(),
+        std::prev(ifOps.end()));
+    if (!op.elseRegion().empty()) {
+      auto &otherOps = op.elseRegion().front().getOperations();
+      affineIf.getElseBlock()->getOperations().splice(
+          std::prev(affineIf.getElseBlock()->end()), otherOps, otherOps.begin(),
+          std::prev(otherOps.end()));
+    }
+    rewriter.finalizeRootUpdate(affineIf);
+    rewriteMemoryOps(affineIf.getBody(), rewriter);
+
+    LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: if converted to:\n";
+               affineIf.dump(););
+    rewriter.replaceOp(op, affineIf.getOperation()->getResults());
+    return success();
+  }
+};
+
+/// Promote fir.do_loop and fir.if to affine.for and affine.if, in the cases
+/// where such a promotion is possible.
+class AffineDialectPromotion
+    : public AffineDialectPromotionBase<AffineDialectPromotion> {
+public:
+  void runOnFunction() override {
+
+    auto *context = &getContext();
+    auto function = getFunction();
+    markAllAnalysesPreserved();
+    auto functionAnalysis = AffineFunctionAnalysis(function);
+    mlir::OwningRewritePatternList patterns(context);
+    patterns.insert<AffineIfConversion>(context, functionAnalysis);
+    patterns.insert<AffineLoopConversion>(context, functionAnalysis);
+    mlir::ConversionTarget target = *context;
+    target.addLegalDialect<mlir::AffineDialect, FIROpsDialect,
+                           mlir::scf::SCFDialect, mlir::StandardOpsDialect>();
+    target.addDynamicallyLegalOp<IfOp>([&functionAnalysis](fir::IfOp op) {
+      return !(functionAnalysis.getChildIfAnalysis(op).canPromoteToAffine());
+    });
+    target.addDynamicallyLegalOp<DoLoopOp>([&functionAnalysis](
+                                               fir::DoLoopOp op) {
+      return !(functionAnalysis.getChildLoopAnalysis(op).canPromoteToAffine());
+    });
+
+    LLVM_DEBUG(llvm::dbgs()
+                   << "AffineDialectPromotion: running promotion on: \n";
+               function.print(llvm::dbgs()););
+    // apply the patterns
+    if (mlir::failed(mlir::applyPartialConversion(function, target,
+                                                  std::move(patterns)))) {
+      mlir::emitError(mlir::UnknownLoc::get(context),
+                      "error in converting to affine dialect\n");
+      signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+/// Convert FIR loop constructs to the Affine dialect
+std::unique_ptr<mlir::Pass> fir::createPromoteToAffinePass() {
+  return std::make_unique<AffineDialectPromotion>();
+}

flang/lib/Optimizer/Transforms/CMakeLists.txt

 add_flang_library(FIRTransforms
+  AffinePromotion.cpp
   Inliner.cpp
   ExternalNameConversion.cpp
 
   DEPENDS
   FIRDialect
   FIRSupport
   FIROptTransformsPassIncGen
 
   LINK_LIBS
   FIRDialect
   MLIRAffineToStandard
   MLIRLLVMIR
   MLIROpenACC
   MLIROpenMP
   FIRSupport
 )

flang/lib/Optimizer/Transforms/PassDetail.h

 //===- PassDetail.h - Optimizer Transforms Pass class details ---*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 #ifndef FORTRAN_OPTMIZER_TRANSFORMS_PASSDETAIL_H
 #define FORTRAN_OPTMIZER_TRANSFORMS_PASSDETAIL_H
 
 #include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
 
 namespace fir {
 
 #define GEN_PASS_CLASSES
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 
 } // namespace fir
 
 #endif // FORTRAN_OPTMIZER_TRANSFORMS_PASSDETAIL_H

flang/test/Fir/affine-promotion.fir

+// Test affine promotion pass
+
+// RUN: fir-opt --split-input-file --promote-to-affine --affine-loop-invariant-code-motion --cse %s | FileCheck %s
+
+!arr_d1 = type !fir.ref<!fir.array<?xf32>>
+#arr_len = affine_map<()[j1,k1] -> (k1 - j1 + 1)>
+
+func @loop_with_load_and_store(%a1: !arr_d1, %a2: !arr_d1, %a3: !arr_d1) {
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %len = constant 100 : index
+  %dims = fir.shape %len : (index) -> !fir.shape<1>
+  %siz = affine.apply #arr_len()[%c1,%len]
+  %t1 = fir.alloca !fir.array<?xf32>, %siz
+
+  fir.do_loop %i = %c1 to %len step %c1 {
+    %a1_idx = fir.array_coor %a1(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+    %a1_v = fir.load %a1_idx : !fir.ref<f32>
+
+    %a2_idx = fir.array_coor %a2(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+    %a2_v = fir.load %a2_idx : !fir.ref<f32>
+
+    %v = addf %a1_v, %a2_v : f32
+    %t1_idx = fir.array_coor %t1(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+
+    fir.store %v to %t1_idx : !fir.ref<f32>
+  }
+  fir.do_loop %i = %c1 to %len step %c1 {
+    %t1_idx = fir.array_coor %t1(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+    %t1_v = fir.load %t1_idx : !fir.ref<f32>
+
+    %a2_idx = fir.array_coor %a2(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+    %a2_v = fir.load %a2_idx : !fir.ref<f32>
+
+    %v = mulf %t1_v, %a2_v : f32
+    %a3_idx = fir.array_coor %a3(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+
+    fir.store %v to %a3_idx : !fir.ref<f32>
+  }
+  return
+}
+
+// CHECK:  func @loop_with_load_and_store(%[[VAL_0:.*]]: !fir.ref<!fir.array<?xf32>>, %[[VAL_1:.*]]: !fir.ref<!fir.array<?xf32>>, %[[VAL_2:.*]]: !fir.ref<!fir.array<?xf32>>) {
+// CHECK:    %[[VAL_3:.*]] = constant 1 : index
+// CHECK:    %[[VAL_4:.*]] = constant 100 : index
+// CHECK:    %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
+// CHECK:    %[[VAL_6:.*]] = affine.apply #map0(){{\[}}%[[VAL_3]], %[[VAL_4]]]
+// CHECK:    %[[VAL_7:.*]] = fir.alloca !fir.array<?xf32>, %[[VAL_6]]
+// CHECK:    %[[VAL_8:.*]] = fir.convert %[[VAL_0]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
+// CHECK:    %[[VAL_9:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
+// CHECK:    %[[VAL_10:.*]] = fir.convert %[[VAL_7]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
+// CHECK:    affine.for %[[VAL_11:.*]] = %[[VAL_3]] to #map1(){{\[}}%[[VAL_4]]] {
+// CHECK:      %[[VAL_12:.*]] = affine.apply #map2(%[[VAL_11]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]]
+// CHECK:      %[[VAL_13:.*]] = affine.load %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<?xf32>
+// CHECK:      %[[VAL_14:.*]] = affine.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<?xf32>
+// CHECK:      %[[VAL_15:.*]] = addf %[[VAL_13]], %[[VAL_14]] : f32
+// CHECK:      affine.store %[[VAL_15]], %[[VAL_10]]{{\[}}%[[VAL_12]]] : memref<?xf32>
+// CHECK:    }
+// CHECK:    %[[VAL_16:.*]] = fir.convert %[[VAL_2]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
+// CHECK:    affine.for %[[VAL_17:.*]] = %[[VAL_3]] to #map1(){{\[}}%[[VAL_4]]] {
+// CHECK:      %[[VAL_18:.*]] = affine.apply #map2(%[[VAL_17]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]]
+// CHECK:      %[[VAL_19:.*]] = affine.load %[[VAL_10]]{{\[}}%[[VAL_18]]] : memref<?xf32>
+// CHECK:      %[[VAL_20:.*]] = affine.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xf32>
+// CHECK:      %[[VAL_21:.*]] = mulf %[[VAL_19]], %[[VAL_20]] : f32
+// CHECK:      affine.store %[[VAL_21]], %[[VAL_16]]{{\[}}%[[VAL_18]]] : memref<?xf32>
+// CHECK:    }
+// CHECK:    return
+// CHECK:  }
+
+// -----
+
+!arr_d1 = type !fir.ref<!fir.array<?xf32>>
+#arr_len = affine_map<()[j1,k1] -> (k1 - j1 + 1)>
+
+func @loop_with_if(%a: !arr_d1, %v: f32) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %len = constant 100 : index
+  %dims = fir.shape %len : (index) -> !fir.shape<1>
+
+  fir.do_loop %i = %c1 to %len step %c1 {
+    fir.do_loop %j = %c1 to %len step %c1 {
+      fir.do_loop %k = %c1 to %len step %c1 {
+        %im2 = subi %i, %c2 : index
+        %cond = cmpi "sgt", %im2, %c0 : index
+        fir.if %cond {
+          %a_idx = fir.array_coor %a(%dims) %i
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+          fir.store %v to %a_idx : !fir.ref<f32>
+        }
+          %aj_idx = fir.array_coor %a(%dims) %j
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+          fir.store %v to %aj_idx : !fir.ref<f32>
+          %ak_idx = fir.array_coor %a(%dims) %k
+            : (!arr_d1, !fir.shape<1>, index) -> !fir.ref<f32>
+          fir.store %v to %ak_idx : !fir.ref<f32>
+      }
+    }
+  }
+  return
+}
+
+// CHECK: func @loop_with_if(%[[VAL_0:.*]]: !fir.ref<!fir.array<?xf32>>, %[[VAL_1:.*]]: f32) {
+// CHECK:   %[[VAL_2:.*]] = constant 0 : index
+// CHECK:   %[[VAL_3:.*]] = constant 1 : index
+// CHECK:   %[[VAL_4:.*]] = constant 2 : index
+// CHECK:   %[[VAL_5:.*]] = constant 100 : index
+// CHECK:   %[[VAL_6:.*]] = fir.shape %[[VAL_5]] : (index) -> !fir.shape<1>
+// CHECK:   %[[VAL_7:.*]] = fir.convert %[[VAL_0]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
+// CHECK:   affine.for %[[VAL_8:.*]] = %[[VAL_3]] to #map0(){{\[}}%[[VAL_5]]] {
+// CHECK:     %[[VAL_9:.*]] = affine.apply #map1(%[[VAL_8]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
+// CHECK:     affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_9]]] : memref<?xf32>
+// CHECK:   }
+// CHECK:   affine.for %[[VAL_10:.*]] = %[[VAL_3]] to #map0(){{\[}}%[[VAL_5]]] {
+// CHECK:     %[[VAL_11:.*]] = affine.apply #map1(%[[VAL_10]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
+// CHECK:     affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xf32>
+// CHECK:   }
+// CHECK:   affine.for %[[VAL_12:.*]] = %[[VAL_3]] to #map0(){{\[}}%[[VAL_5]]] {
+// CHECK:     %[[VAL_13:.*]] = subi %[[VAL_12]], %[[VAL_4]] : index
+// CHECK:     affine.if #set(%[[VAL_12]]) {
+// CHECK:       %[[VAL_14:.*]] = affine.apply #map1(%[[VAL_12]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
+// CHECK:       affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xf32>
+// CHECK:     }
+// CHECK:   }
+// CHECK:   return
+// CHECK: }