[fir] Add FIR CSE specific pass

[clementval]

Summary

Perform common subexpression elimination on FIR operations. This pass
differs from the MLIR CSE pass in that it is FIR/Fortran semantics aware.

It takes advantages of additional Fortran semantic such as pure function call.
pure procedures and functions are free of side effects. In some cases,
pure calls can be safely removed.

fir.load are also optimized by this dedicated CSE pass. Volatile loads will
prevent the optimization to happen.

Running the same test (flang/test/Fir/cse.fir) through the MLIR CSE
pass will not give a similar results.

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

[kiranchandramohan]

flang/test/Fir/cse.fir

31

Not asking for a change here. I made the following change and it retains a fir.call in the second block. I guess this is because we process blocks separately.

-  %9 = arith.addi %7, %8 : i64
+  br ^bb2
+ ^bb2:  %9 = arith.addi %7, %8 : i64

37

I think COUNT will not check that the addis are consecutive with no fir.calls in between.
See test in the comment above.

[mehdi_amini]

This pass differs from the MLIR CSE pass in that it is FIR/Fortran semantics aware.

Can you expand and provide details about this?

(marking as "request changes" because @clementval told me to do so when discussing a revision, I haven't looked at the code at this point)

[clementval]

• split input file for the test
• make the checks more explicit

[clementval]

In D112711#3095207, @mehdi_amini wrote:

This pass differs from the MLIR CSE pass in that it is FIR/Fortran semantics aware.

Can you expand and provide details about this?

(marking as "request changes" because @clementval told me to do so when discussing a revision, I haven't looked at the code at this point)

I updated the description about some information. This pass mainly use semantic information on fir.call and fir.load and optimized them out when possible. This is relying on the pure semantic for the calls.
The optimization on fir.load will be prevented if they are flagged as volatile loads.

I think we should use "request changes" if there is any blocker change requested in the code. I guess this is also fair usage if you want to block this patch for discussion.

[mehdi_amini]

Right so I'm interested in what it'll take to converge this with the existing CSE pass. What you mentioned here is interesting because it seems all in the domain of memory effects I believe: does it point mostly to the lack of support for the memory effects op interface in the CSE pass?

[clementval]

In D112711#3098053, @mehdi_amini wrote:

Right so I'm interested in what it'll take to converge this with the existing CSE pass. What you mentioned here is interesting because it seems all in the domain of memory effects I believe: does it point mostly to the lack of support for the memory effects op interface in the CSE pass?

I think what difference we have right now with the memory effects op interface is that in the case of the memory effects op interface it applies to a op all the time. In our case, the fir.call can be pure or not. I'm not sure how this would work with the memory effects. It would probably requires two distincts operations or I'm missing something. Sorry, I'm not be up to date with the current memory effects op interface (will have a look at it today to refresh my knowledge :-) ).

[mehdi_amini]

In D112711#3102388, @clementval wrote:

In D112711#3098053, @mehdi_amini wrote:

Right so I'm interested in what it'll take to converge this with the existing CSE pass. What you mentioned here is interesting because it seems all in the domain of memory effects I believe: does it point mostly to the lack of support for the memory effects op interface in the CSE pass?

I think what difference we have right now with the memory effects op interface is that in the case of the memory effects op interface it applies to a op all the time. In our case, the fir.call can be pure or not. I'm not sure how this would work with the memory effects.

When an an operation implements the interface, it can write any kind of C++, including looking up attributes on the operation. This is very much a query made on a specific instance of the operation and not a query that has to be static on the op class I believe.

See for example:

• The op definition here: https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/Dialect/Test/TestOps.td#L2058 (the line after shows the support for custom non-memory effects by the way)
• The interface implementation for the op: https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/Dialect/Test/TestDialect.cpp#L887-L921
• The associated test that shows how the same op can have different effects based on the attributes: https://github.com/llvm/llvm-project/blob/main/mlir/test/IR/test-side-effects.mlir

[clementval]

In D112711#3104689, @mehdi_amini wrote:

In D112711#3102388, @clementval wrote:

In D112711#3098053, @mehdi_amini wrote:

Right so I'm interested in what it'll take to converge this with the existing CSE pass. What you mentioned here is interesting because it seems all in the domain of memory effects I believe: does it point mostly to the lack of support for the memory effects op interface in the CSE pass?

I think what difference we have right now with the memory effects op interface is that in the case of the memory effects op interface it applies to a op all the time. In our case, the fir.call can be pure or not. I'm not sure how this would work with the memory effects.

When an an operation implements the interface, it can write any kind of C++, including looking up attributes on the operation. This is very much a query made on a specific instance of the operation and not a query that has to be static on the op class I believe.

See for example:

• The op definition here: https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/Dialect/Test/TestOps.td#L2058 (the line after shows the support for custom non-memory effects by the way)
• The interface implementation for the op: https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/Dialect/Test/TestDialect.cpp#L887-L921
• The associated test that shows how the same op can have different effects based on the attributes: https://github.com/llvm/llvm-project/blob/main/mlir/test/IR/test-side-effects.mlir

Thanks for the pointers. I'll have a look at that.

[clementval]

Coming back to this. Looks like the existing CSE is not able to deal with simple use case without analysis. Given the example below I would expect the pass to remove one of the redundant load but this is not happening with the existing MLIR CSE pass but we have such feature in the FIR CSE pass. Maybe I'm missing something but I guess that without a proper analysis the MLIR CSE pass would not be able to deal with such cases in FIR. Should we go ahead and have our dedicated CSE for specific FIR operations or is there smth I'm missing?

func @fun(%arg0: !fir.ref<i64>) -> i64 {
    %0 = fir.load %arg0 : !fir.ref<i64>
    %1 = fir.load %arg0 : !fir.ref<i64>
    %2 = arith.addi %0, %1 : i64
    return %2 : i64
}

[mehdi_amini]

I think it would be more productive as a community to improve the existing CSE to figure out how to make it work for these cases instead.

[clementval]

In D112711#3406661, @mehdi_amini wrote:

I think it would be more productive as a community to improve the existing CSE to figure out how to make it work for these cases instead.

Sure. I'll look how we can make it work.

[clementval]

First patch in this direction here: https://reviews.llvm.org/D122801

[clementval]

Adding CSE for commutative ops - https://reviews.llvm.org/D123492

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

 namespace fir {
 
 //===----------------------------------------------------------------------===//
 // Passes defined in Passes.td
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<mlir::Pass> createAbstractResultOptPass();
 std::unique_ptr<mlir::Pass> createAffineDemotionPass();
+std::unique_ptr<mlir::Pass> createCSEPass();
 std::unique_ptr<mlir::Pass> createFirToCfgPass();
 std::unique_ptr<mlir::Pass> createCharacterConversionPass();
 std::unique_ptr<mlir::Pass> createExternalNameConversionPass();
 std::unique_ptr<mlir::Pass> createPromoteToAffinePass();
 
 /// 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

let description = [{
     `fir.load` and `fir.store` with `!fir.ref<?>` types is required.
   }];
   let constructor = "::fir::createAffineDemotionPass()";
   let dependentDialects = [
     "fir::FIROpsDialect", "mlir::StandardOpsDialect", "mlir::AffineDialect"
   ];
 }
 
+def BasicCSE : FunctionPass<"basic-cse"> {
+  let summary = "Basic common sub-expression elimination.";
+  let description = [{
+     Perform common subexpression elimination on FIR operations. This pass
+     differs from the MLIR CSE pass in that it is FIR/Fortran semantics aware.
+  }];
+  let constructor = "::fir::createCSEPass()";
+  let dependentDialects = [
+    "fir::FIROpsDialect", "mlir::StandardOpsDialect"
+  ];
+}
+
 def CharacterConversion : Pass<"character-conversion"> {
   let summary = "Convert CHARACTER entities with different KINDs";
   let description = [{
     Translates entities of one CHARACTER KIND to another.
 
     By default the translation is to naively zero-extend or truncate a code
     point to fit the destination size.
   }];

flang/lib/Optimizer/Transforms/CMakeLists.txt

 add_flang_library(FIRTransforms
   AbstractResult.cpp
   AffinePromotion.cpp
   AffineDemotion.cpp
   CharacterConversion.cpp
+  CSE.cpp
   Inliner.cpp
   ExternalNameConversion.cpp
   RewriteLoop.cpp
 
   DEPENDS
   FIRDialect
   FIRSupport
   FIROptTransformsPassIncGen

flang/lib/Optimizer/Transforms/CSE.cpp

+//===-- CSE.cpp -- common subexpression elimination -----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This transformation pass performs a simple common sub-expression elimination
+/// algorithm on operations within a function.
+///
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/Utils.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/ScopedHashTable.h"
+#include "llvm/Support/Allocator.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/RecyclingAllocator.h"
+#include <deque>
+
+using namespace mlir;
+
+static llvm::cl::opt<bool>
+    leaveEffects("keep-effects",
+                 llvm::cl::desc("disable cleaning up effects attributes"),
+                 llvm::cl::init(false), llvm::cl::Hidden);
+
+namespace {
+
+struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
+
+  /// Compute the hash value of an Operation
+  static unsigned getHashValue(const Operation *opC) {
+    auto *op = const_cast<Operation *>(opC);
+    // Hash the operations based upon their:
+    //   - Operation Name
+    //   - Attributes
+    //   - Result Types
+    //   - Operands
+    unsigned hashOps;
+    if (op->hasTrait<mlir::OpTrait::IsCommutative>()) {
+      std::vector<void *> vec;
+      for (auto i = op->operand_begin(), e = op->operand_end(); i != e; ++i)
+        vec.push_back((*i).getAsOpaquePointer());
+      llvm::sort(vec.begin(), vec.end());
+      hashOps = llvm::hash_combine_range(vec.begin(), vec.end());
+    } else {
+      hashOps =
+          llvm::hash_combine_range(op->operand_begin(), op->operand_end());
+    }
+    auto hashResTys{llvm::hash_combine_range(op->result_type_begin(),
+                                             op->result_type_end())};
+    return llvm::hash_combine(op->getName(), op->getAttrs(), hashResTys,
+                              hashOps);
+  }
+
+  static bool isEqual(const Operation *lhsC, const Operation *rhsC) {
+    auto *lhs = const_cast<Operation *>(lhsC);
+    auto *rhs = const_cast<Operation *>(rhsC);
+    if (lhs == rhs)
+      return true;
+    if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
+        rhs == getTombstoneKey() || rhs == getEmptyKey())
+      return false;
+
+    // Compare the operation name.
+    if (lhs->getName() != rhs->getName())
+      return false;
+    // Check operand and result type counts.
+    if (lhs->getNumOperands() != rhs->getNumOperands() ||
+        lhs->getNumResults() != rhs->getNumResults())
+      return false;
+    // Compare attributes.
+    if (lhs->getAttrs() != rhs->getAttrs())
+      return false;
+    // Compare operands.
+    if (lhs->hasTrait<mlir::OpTrait::IsCommutative>()) {
+      SmallVector<void *> lops;
+      for (auto lod : lhs->getOperands())
+        lops.push_back(lod.getAsOpaquePointer());
+      llvm::sort(lops.begin(), lops.end());
+      SmallVector<void *> rops;
+      for (auto rod : rhs->getOperands())
+        rops.push_back(rod.getAsOpaquePointer());
+      llvm::sort(rops.begin(), rops.end());
+      if (!std::equal(lops.begin(), lops.end(), rops.begin()))
+        return false;
+    } else {
+      if (!std::equal(lhs->operand_begin(), lhs->operand_end(),
+                      rhs->operand_begin()))
+        return false;
+    }
+    // Compare result types.
+    return std::equal(lhs->result_type_begin(), lhs->result_type_end(),
+                      rhs->result_type_begin());
+  }
+};
+
+/// Basic common sub-expression elimination.
+struct BasicCSE : public fir::BasicCSEBase<BasicCSE> {
+  /// Shared implementation of operation elimination and scoped map definitions.
+  using AllocatorTy = llvm::RecyclingAllocator<
+      llvm::BumpPtrAllocator,
+      llvm::ScopedHashTableVal<Operation *, Operation *>>;
+  using ScopedMapTy = llvm::ScopedHashTable<Operation *, Operation *,
+                                            SimpleOperationInfo, AllocatorTy>;
+
+  /// Represents a single entry in the depth first traversal of a CFG.
+  struct CFGStackNode {
+    CFGStackNode(ScopedMapTy &knownValues, DominanceInfoNode *node)
+        : scope(knownValues), node(node), childIterator(node->begin()),
+          processed(false) {}
+
+    /// Scope for the known values.
+    ScopedMapTy::ScopeTy scope;
+
+    DominanceInfoNode *node;
+    DominanceInfoNode::iterator childIterator;
+
+    /// If this node has been fully processed yet or not.
+    bool processed;
+  };
+
+  /// Attempt to eliminate a redundant operation. Returns success if the
+  /// operation was marked for removal, failure otherwise.
+  LogicalResult simplifyOperation(ScopedMapTy &knownValues, Operation *op);
+
+  void simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                     Block *bb);
+  void simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                      Region &region);
+
+  void cleanupBlock(Block *bb) {
+    for (auto &inst : *bb) {
+      if (fir::nonVolatileLoad(&inst) || fir::pureCall(&inst)) {
+        inst.removeAttr(Identifier::get("effects_token", inst.getContext()));
+      } else if (inst.getNumRegions()) {
+        for (auto &region : inst.getRegions())
+          cleanupRegion(region);
+      }
+    }
+  }
+  void cleanupRegion(Region &region) {
+    for (auto &block : region)
+      cleanupBlock(&block);
+  }
+
+  void runOnFunction() override final;
+
+private:
+  /// Operations marked as dead and to be erased.
+  std::vector<Operation *> opsToErase;
+};
+
+/// Attempt to eliminate a redundant operation.
+LogicalResult BasicCSE::simplifyOperation(ScopedMapTy &knownValues,
+                                          Operation *op) {
+  if (op->hasTrait<mlir::OpTrait::IsTerminator>())
+    return failure();
+
+  if (isOpTriviallyDead(op)) {
+    opsToErase.push_back(op);
+    return success();
+  }
+
+  // Don't simplify operations with nested blocks. We don't currently model
+  // equality comparisons correctly among other things. It is also unclear
+  // whether we would want to CSE such operations.
+  if (op->getNumRegions() != 0)
+    return failure();
+
+  if (!MemoryEffectOpInterface::hasNoEffect(op) && !fir::nonVolatileLoad(op) &&
+      !fir::pureCall(op))
+    return failure();
+
+  // Look for an existing definition for the operation.
+  if (auto *existing = knownValues.lookup(op)) {
+    // If we find one then replace all uses of the current operation with the
+    // existing one and mark it for deletion.
+    op->replaceAllUsesWith(existing);
+    if (!op->hasTrait<mlir::OpTrait::IsTerminator>())
+      opsToErase.push_back(op);
+
+    // If the existing operation has an unknown location and the current
+    // operation doesn't, then set the existing op's location to that of the
+    // current op.
+    if (existing->getLoc().isa<UnknownLoc>() &&
+        !op->getLoc().isa<UnknownLoc>()) {
+      existing->setLoc(op->getLoc());
+    }
+    return success();
+  }
+
+  // Otherwise, we add this operation to the known values map.
+  knownValues.insert(op, op);
+  return failure();
+}
+
+void BasicCSE::simplifyBlock(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                             Block *bb) {
+  std::intptr_t token = reinterpret_cast<std::intptr_t>(bb);
+  for (auto &inst : *bb) {
+    if (fir::nonVolatileLoad(&inst) || fir::pureCall(&inst))
+      inst.setAttr("effects_token",
+                   IntegerAttr::get(IndexType::get(inst.getContext()), token));
+    if (isa<fir::StoreOp>(&inst) || fir::impureCall(&inst) ||
+        inst.getNumRegions() != 0)
+      token = reinterpret_cast<std::intptr_t>(&inst);
+  }
+  for (auto &inst : *bb) {
+    // If the operation is simplified, we don't process any held regions.
+    if (succeeded(simplifyOperation(knownValues, &inst)))
+      continue;
+
+    // If this operation is isolated above, we can't process nested regions with
+    // the given 'knownValues' map. This would cause the insertion of implicit
+    // captures in explicit capture only regions.
+    if (!inst.isRegistered() ||
+        inst.hasTrait<mlir::OpTrait::IsIsolatedFromAbove>()) {
+      ScopedMapTy nestedKnownValues;
+      for (auto &region : inst.getRegions())
+        simplifyRegion(nestedKnownValues, domInfo, region);
+      continue;
+    }
+
+    // Otherwise, process nested regions normally.
+    for (auto &region : inst.getRegions())
+      simplifyRegion(knownValues, domInfo, region);
+  }
+}
+
+void BasicCSE::simplifyRegion(ScopedMapTy &knownValues, DominanceInfo &domInfo,
+                              Region &region) {
+  // If the region is empty there is nothing to do.
+  if (region.empty())
+    return;
+
+  // If the region only contains one block, then simplify it directly.
+  if (std::next(region.begin()) == region.end()) {
+    ScopedMapTy::ScopeTy scope(knownValues);
+    simplifyBlock(knownValues, domInfo, &region.front());
+    return;
+  }
+
+  // Note, deque is being used here because there was significant performance
+  // gains over vector when the container becomes very large due to the
+  // specific access patterns. If/when these performance issues are no
+  // longer a problem we can change this to vector. For more information see
+  // the llvm mailing list discussion on this:
+  // http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20120116/135228.html
+  std::deque<std::unique_ptr<CFGStackNode>> stack;
+
+  // Process the nodes of the dom tree for this region.
+  stack.emplace_back(std::make_unique<CFGStackNode>(
+      knownValues, domInfo.getRootNode(&region)));
+
+  while (!stack.empty()) {
+    auto &currentNode = stack.back();
+
+    // Check to see if we need to process this node.
+    if (!currentNode->processed) {
+      currentNode->processed = true;
+      simplifyBlock(knownValues, domInfo, currentNode->node->getBlock());
+    }
+
+    // Otherwise, check to see if we need to process a child node.
+    if (currentNode->childIterator != currentNode->node->end()) {
+      auto *childNode = *(currentNode->childIterator++);
+      stack.emplace_back(
+          std::make_unique<CFGStackNode>(knownValues, childNode));
+    } else {
+      // Finally, if the node and all of its children have been processed
+      // then we delete the node.
+      stack.pop_back();
+    }
+  }
+}
+
+void BasicCSE::runOnFunction() {
+  /// A scoped hash table of defining operations within a function.
+  {
+    ScopedMapTy knownValues;
+    simplifyRegion(knownValues, getAnalysis<DominanceInfo>(),
+                   getFunction().getBody());
+  }
+  if (!leaveEffects)
+    cleanupRegion(getFunction().getBody());
+
+  // If no operations were erased, then we mark all analyses as preserved.
+  if (opsToErase.empty())
+    return markAllAnalysesPreserved();
+
+  /// Erase any operations that were marked as dead during simplification.
+  for (auto *op : opsToErase)
+    op->erase();
+  opsToErase.clear();
+
+  // We currently don't remove region operations, so mark dominance as
+  // preserved.
+  markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();
+}
+
+} // end anonymous namespace
+
+std::unique_ptr<mlir::Pass> fir::createCSEPass() {
+  return std::make_unique<BasicCSE>();
+}

flang/test/Fir/cse.fir

+// Test CSE pass
+
+// RUN: fir-opt --split-input-file --basic-cse %s | FileCheck %s
+
+func @fun(%a : !fir.ref<i64>) -> i64 {
+  %1 = fir.load %a : !fir.ref<i64>
+  %2 = fir.load %a : !fir.ref<i64>
+  %3 = arith.addi %1, %2 : i64
+  %4 = fir.load %a : !fir.ref<i64>
+  %5 = arith.addi %3, %4 : i64
+  %6 = fir.load %a : !fir.ref<i64>
+  %7 = arith.addi %5, %6 : i64
+  %8 = fir.load %a : !fir.ref<i64>
+  %9 = arith.addi %7, %8 : i64
+  %10 = fir.load %a : !fir.ref<i64>
+  %11 = arith.addi %10, %9 : i64
+  %12 = fir.load %a : !fir.ref<i64>
+  %13 = arith.addi %11, %12 : i64
+  return %13 : i64
+}
+
+// CHECK-LABEL: @fun
+// CHECK:         %{{.*}} = fir.load %{{.*}} : !fir.ref<i64>
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     fir.load %{{.*}} : !fir.ref<i64>
+// CHECK-NEXT:    return %{{.*}} i64
+
+// -----
+
+func private @bar(%a : !fir.ref<i64>) -> i64
+
+func @fun2(%a : !fir.ref<i64>) -> i64 {
+  %1 = fir.load %a : !fir.ref<i64>
+  %2 = fir.call @bar(%a) { pure = true } : (!fir.ref<i64>) -> i64
+  %3 = arith.addi %1, %2 : i64
+  %4 = fir.call @bar(%a) { pure = true } : (!fir.ref<i64>) -> i64
+  %5 = arith.addi %3, %4 : i64
+  %6 = fir.call @bar(%a) { pure = true } : (!fir.ref<i64>) -> i64
+  %7 = arith.addi %5, %6 : i64
+  %8 = fir.call @bar(%a) { pure = true } : (!fir.ref<i64>) -> i64
+  %9 = arith.addi %7, %8 : i64
+  %10 = fir.call @bar(%a) { pure = true } : (!fir.ref<i64>) -> i64
+  %11 = arith.addi %10, %9 : i64
+  %12 = fir.call @bar(%a) { pure = true } : (!fir.ref<i64>) -> i64
+  %13 = arith.addi %11, %12 : i64
+  return %13 : i64
+}
+
+// CHECK-LABEL: @bar
+
+// CHECK-LABEL: @fun2
+// CHECK:         %{{.*}} = fir.load %{{.*}} : !fir.ref<i64>
+// CHECK-NEXT:    fir.call @bar
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     %{{.*}} = fir.call
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     %{{.*}} = fir.call
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     %{{.*}} = fir.call
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     %{{.*}} = fir.call
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NOT:     %{{.*}} = fir.call
+// CHECK:         arith.addi %{{.*}}, %{{.*}} : i64
+// CHECK-NEXT:    return %{{.*}} : i64
+
+// -----
+
+// Negative test: do not merge loads when an op with regions is between
+
+func @foo(%arg0: !fir.ref<f32>) -> f32 {
+  %0 = fir.alloca f32 {name = "x"}
+  %1 = fir.load %arg0 : !fir.ref<f32>
+  fir.store %1 to %0 : !fir.ref<f32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %2 = fir.load %0 : !fir.ref<f32>
+  %3 = arith.cmpf olt, %2, %cst : f32
+  fir.if %3 {
+    %7 = fir.load %0 : !fir.ref<f32>
+    %8 = arith.negf %7 : f32
+    fir.store %8 to %0 : !fir.ref<f32>
+  }
+  %cst_0 = arith.constant 1.000000e+00 : f32
+  %4 = fir.load %0 : !fir.ref<f32>
+  %5 = arith.addf %4, %cst_0 : f32
+  fir.store %5 to %0 : !fir.ref<f32>
+  %6 = fir.load %0 : !fir.ref<f32>
+  return %6 : f32
+}
+
+// CHECK-LABEL: @foo
+// CHECK:         %{{.*}} = fir.alloca f32 {name = "x"}
+// CHECK-COUNT-4: %{{.*}} = fir.load %{{.*}} : !fir.ref<f32>