mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp

//===- ParallelLoopFusion.cpp - Code to perform loop fusion ---------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements loop fusion on parallel loops.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/SCF/Transforms/Passes.h"

#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/InterleavedRange.h"

#include "llvm/Support/DebugLog.h"
#include <numeric>
#include <optional>
#include <tuple>
#define DEBUG_TYPE "parallel-loop-fusion"

namespace mlir {
#define GEN_PASS_DEF_SCFPARALLELLOOPFUSION
#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
} // namespace mlir

using namespace mlir;
using namespace mlir::scf;

/// Verify there are no nested ParallelOps.
static bool hasNestedParallelOp(ParallelOp ploop) {
  auto walkResult =
      ploop.getBody()->walk([](ParallelOp) { return WalkResult::interrupt(); });
  return walkResult.wasInterrupted();
}

/// Verify equal iteration spaces.
static bool equalIterationSpaces(ParallelOp firstPloop,
                                 ParallelOp secondPloop) {
  if (firstPloop.getNumLoops() != secondPloop.getNumLoops() ||
      firstPloop.getUnsignedCmp() != secondPloop.getUnsignedCmp())
    return false;

  // Two bounds match if they are the same value, or if both are constants
  // holding the same value. The latter matters because equivalent bounds are
  // often materialized by distinct `arith.constant` ops, which leaves the
  // iteration spaces equal even though the SSA values differ.
  auto matchOperands = [&](const OperandRange &lhs,
                           const OperandRange &rhs) -> bool {
    // TODO: Extend this to support aliases.
    return std::equal(lhs.begin(), lhs.end(), rhs.begin(),
                      [](Value lhsValue, Value rhsValue) {
                        if (lhsValue == rhsValue)
                          return true;
                        std::optional<int64_t> lhsConst =
                            getConstantIntValue(lhsValue);
                        std::optional<int64_t> rhsConst =
                            getConstantIntValue(rhsValue);
                        return lhsConst && rhsConst && *lhsConst == *rhsConst;
                      });
  };
  return matchOperands(firstPloop.getLowerBound(),
                       secondPloop.getLowerBound()) &&
         matchOperands(firstPloop.getUpperBound(),
                       secondPloop.getUpperBound()) &&
         matchOperands(firstPloop.getStep(), secondPloop.getStep());
}

/// Check if both operations are the same type of memory write op and
/// write to the same memory location (same buffer and same indices).
static bool opsWriteSameMemLocation(Operation *op1, Operation *op2) {
  if (!op1 || !op2 || op1->getName() != op2->getName())
    return false;
  if (op1 == op2)
    return true;
  // support only these memory-writing ops for now
  if (!isa<memref::StoreOp, vector::TransferWriteOp, vector::StoreOp>(op1))
    return false;
  bool opsAreIdentical =
      llvm::TypeSwitch<Operation *, bool>(op1)
          .Case([&](memref::StoreOp storeOp1) {
            auto storeOp2 = cast<memref::StoreOp>(op2);
            return (storeOp1.getMemRef() == storeOp2.getMemRef()) &&
                   (storeOp1.getIndices() == storeOp2.getIndices());
          })
          .Case([&](vector::TransferWriteOp writeOp1) {
            auto writeOp2 = cast<vector::TransferWriteOp>(op2);
            return (writeOp1.getBase() == writeOp2.getBase()) &&
                   (writeOp1.getIndices() == writeOp2.getIndices()) &&
                   (writeOp1.getMask() == writeOp2.getMask()) &&
                   (writeOp1.getValueToStore().getType() ==
                    writeOp2.getValueToStore().getType()) &&
                   (writeOp1.getInBounds() == writeOp2.getInBounds());
          })
          .Case([&](vector::StoreOp vecStoreOp1) {
            auto vecStoreOp2 = cast<vector::StoreOp>(op2);
            return (vecStoreOp1.getBase() == vecStoreOp2.getBase()) &&
                   (vecStoreOp1.getIndices() == vecStoreOp2.getIndices()) &&
                   (vecStoreOp1.getValueToStore().getType() ==
                    vecStoreOp2.getValueToStore().getType()) &&
                   (vecStoreOp1.getAlignment() == vecStoreOp2.getAlignment()) &&
                   (vecStoreOp1.getNontemporal() ==
                    vecStoreOp2.getNontemporal());
          })
          .Default([](Operation *) { return false; });
  return opsAreIdentical;
}

/// Check if val1 (from the first parallel loop) and val2 (from the
/// second) are equivalent, considering the mapping of induction variables from
/// the first to the second parallel loop.
static bool valsAreEquivalent(Value val1, Value val2,
                              const IRMapping &loopsIVsMap) {
  if (val1 == val2 || loopsIVsMap.lookupOrDefault(val1) == val2 ||
      loopsIVsMap.lookupOrDefault(val2) == val1)
    return true;
  Operation *val1DefOp = val1.getDefiningOp();
  Operation *val2DefOp = val2.getDefiningOp();
  if (!val1DefOp || !val2DefOp)
    return false;
  if (!isMemoryEffectFree(val1DefOp) || !isMemoryEffectFree(val2DefOp))
    return false;
  return OperationEquivalence::isEquivalentTo(
      val1DefOp, val2DefOp,
      [&](Value v1, Value v2) {
        return success(loopsIVsMap.lookupOrDefault(v1) == v2 ||
                       loopsIVsMap.lookupOrDefault(v2) == v1);
      },
      /*markEquivalent=*/nullptr, OperationEquivalence::Flags::IgnoreLocations);
}

/// If the `expr` value is the result of an integer addition of `base` and a
/// constant, return the constant.
static std::optional<int64_t> getAddConstant(Value expr, Value base,
                                             const IRMapping &loopsIVsMap) {
  if (auto addOp = expr.getDefiningOp<arith::AddIOp>()) {
    if (auto constOp = getConstantIntValue(addOp.getLhs());
        constOp && valsAreEquivalent(addOp.getRhs(), base, loopsIVsMap))
      return constOp.value();
    if (auto constOp = getConstantIntValue(addOp.getRhs());
        constOp && valsAreEquivalent(addOp.getLhs(), base, loopsIVsMap))
      return constOp.value();
    return std::nullopt;
  }

  if (auto addOp = expr.getDefiningOp<index::AddOp>()) {
    if (auto constOp = getConstantIntValue(addOp.getLhs());
        constOp && valsAreEquivalent(addOp.getRhs(), base, loopsIVsMap))
      return constOp.value();
    if (auto constOp = getConstantIntValue(addOp.getRhs());
        constOp && valsAreEquivalent(addOp.getLhs(), base, loopsIVsMap))
      return constOp.value();
    return std::nullopt;
  }

  if (auto applyOp = expr.getDefiningOp<affine::AffineApplyOp>()) {
    AffineMap map = applyOp.getAffineMap();
    if (map.getNumResults() != 1 || map.getNumDims() != 1 ||
        map.getNumSymbols() != 0)
      return std::nullopt;
    if (!valsAreEquivalent(applyOp.getOperand(0), base, loopsIVsMap))
      return std::nullopt;
    AffineExpr result = map.getResult(0);
    auto bin = dyn_cast<AffineBinaryOpExpr>(result);
    if (!bin || bin.getKind() != AffineExprKind::Add)
      return std::nullopt;
    auto lhsDim = dyn_cast<AffineDimExpr>(bin.getLHS());
    auto rhsDim = dyn_cast<AffineDimExpr>(bin.getRHS());
    auto lhsConst = dyn_cast<AffineConstantExpr>(bin.getLHS());
    auto rhsConst = dyn_cast<AffineConstantExpr>(bin.getRHS());
    if (lhsConst && rhsDim)
      return lhsConst.getValue();
    if (rhsConst && lhsDim)
      return rhsConst.getValue();
  }
  return std::nullopt;
}

// Return true if the scalar load index may hit any element covered by a
// vector.store/transfer_write along a single memref dimension. Supported cases:
//
// 1) Direct index match (with optional offset):
//    vector.transfer_write %v, %A[%i] : vector<4xf32>, memref<...>
//    %x = memref.load %A[%i] : memref<...>
//
// 2) Loop IV range intersects the write range:
//    vector.transfer_write %v, %A[%c0] : vector<4xf32>, memref<...>
//    scf.for %k = %c0 to %c4 step %c1 { %x = memref.load %A[%k] }
//
// 3) Constant index (or IV + constant) within the write range:
//    vector.transfer_write %v, %A[%c0] : vector<4xf32>, memref<...>
//    %x = memref.load %A[%c2] : memref<...>
//    %y = memref.load %A[%i + %c1] : memref<...>
//
// Args:
// - loadIndex: index used by the scalar load for this dimension.
// - offset: subview offset for the base memref dimension (if any).
// - writeIndex: index used by the transfer_write for this dimension. Can be
// null if the dim was dropped by a rank reducing subview, whose result is
// written by the vector.write.
// - extent: vector size along this dimension (number of elements written).
// - loopsIVsMap: IV equivalence map between fused loops.
static bool loadIndexWithinWriteRange(Value loadIndex, OpFoldResult offset,
                                      Value writeIndex, int64_t extent,
                                      const IRMapping &loopsIVsMap) {
  if (extent <= 0)
    return false;

  // Extract constant loop bounds for loop IVs (e.g. from scf.for).
  auto getConstLoopBoundsForIV =
      [](Value index) -> std::optional<std::tuple<int64_t, int64_t, int64_t>> {
    auto blockArg = dyn_cast<BlockArgument>(index);
    if (!blockArg)
      return std::nullopt;
    auto *parentOp = blockArg.getOwner()->getParentOp();
    auto loopLike = dyn_cast<LoopLikeOpInterface>(parentOp);
    if (!loopLike)
      return std::nullopt;
    auto ranges = getConstLoopBounds(loopLike);
    if (ranges.empty())
      return std::nullopt;

    auto ivs = loopLike.getLoopInductionVars();
    if (!ivs)
      return std::nullopt;
    auto it = llvm::find(*ivs, blockArg);
    if (it == ivs->end())
      return std::nullopt;
    unsigned pos = std::distance(ivs->begin(), it);
    if (pos >= ranges.size())
      return std::nullopt;
    auto [lb, ub, step] = ranges[pos];
    return std::make_tuple(lb, ub, step);
  };

  std::optional<int64_t> offsetConst = getConstantIntValue(offset);
  std::optional<int64_t> writeConst =
      writeIndex ? getConstantIntValue(writeIndex) : std::optional<int64_t>(0);
  if (!writeConst && writeIndex) {
    // Treat single-iteration IVs as constants for matching.
    if (auto bounds = getConstLoopBoundsForIV(writeIndex)) {
      auto [lb, ub, step] = *bounds;
      if (step > 0 && ub == lb + step)
        writeConst = lb;
    }
  }

  // Check whether a loop IV is fully contained in a constant write range.
  auto loopIVWithinRange = [](int64_t lb, int64_t ub, int64_t step,
                              int64_t rangeStart, int64_t rangeExtent) -> bool {
    if (rangeExtent <= 0 || step <= 0)
      return false;
    if (ub <= lb)
      return false;
    int64_t rangeEnd = rangeStart + rangeExtent;
    return lb >= rangeStart && ub <= rangeEnd;
  };

  if (offsetConst && writeConst) {
    // Constant start of the write range; check constant load or loop IV range.
    int64_t start = *offsetConst + *writeConst;
    if (auto loadConst = getConstantIntValue(loadIndex))
      return (*loadConst >= start && *loadConst < start + extent);
    if (auto bounds = getConstLoopBoundsForIV(loadIndex)) {
      auto [lb, ub, step] = *bounds;
      return loopIVWithinRange(lb, ub, step, start, extent);
    }
  }

  if (writeIndex) {
    // Direct IV match (or IV + constant) against the write index.
    if (offsetConst && *offsetConst == 0 &&
        valsAreEquivalent(loadIndex, writeIndex, loopsIVsMap))
      return true;
    if (auto addConst = getAddConstant(loadIndex, writeIndex, loopsIVsMap)) {
      // Match load index of the form writeIndex + C within the write extent.
      if (offsetConst) {
        int64_t start = *offsetConst;
        return (*addConst >= start && *addConst < start + extent);
      }
    }
    return false;
  }

  if (auto offsetVal = dyn_cast<Value>(offset)) {
    // Exact match when extent is 1 and the load hits the offset value.
    if (extent == 1 && valsAreEquivalent(loadIndex, offsetVal, loopsIVsMap))
      return true;
  }

  return false;
}

/// Return the base memref value used by the given memory op.
static Value getBaseMemref(Operation *op) {
  // TODO: use the common interface for memory ops once available.
  return llvm::TypeSwitch<Operation *, Value>(op)
      .Case([&](memref::LoadOp load) { return load.getMemRef(); })
      .Case([&](memref::StoreOp store) { return store.getMemRef(); })
      .Case([&](vector::TransferReadOp read) { return read.getBase(); })
      .Case([&](vector::TransferWriteOp write) { return write.getBase(); })
      .Case([&](vector::LoadOp load) { return load.getBase(); })
      .Case([&](vector::StoreOp store) { return store.getBase(); })
      .Default([](Operation *) { return Value(); });
}

/// Recognize scalar memref.load of an element produced by a vector write
/// (vector.transfer_write or vector.store, optionally through a rank-reducing
/// unit-stride subview) of the same buffer. This covers the pattern where a
/// vector write stores a full lane pack and a subsequent scalar load reads an
/// element from that lane pack. EXAMPLE:
///  vector.transfer_write %V, %arg[%x, %y, ..., 0] {in_bounds = [true]} :
///             vector<4xf32>, memref<4xf32, strided<[1], offset: ?>>
///  scf.for %iter = %c0 to %c4 step %c1 iter_args(...) -> (f32) {
///    %0 = memref.load %arg[%x, %y, ..., %iter] : memref<1x128x16x4xf32>
///    ...
///  }
///
static bool isLoadOnWrittenVector(memref::LoadOp loadOp, Value writeBase,
                                  ValueRange writeIndices, VectorType vecTy,
                                  ArrayRef<int64_t> vectorDimForWriteDim,
                                  const IRMapping &ivsMap) {
  if (!vecTy)
    return false;

  Value base = writeBase;
  // The write base if there is no subview, or the subview source otherwise.
  MemrefValue baseMemref = nullptr;
  SmallVector<OpFoldResult> offsets;
  llvm::SmallBitVector droppedDims;
  bool hasSubview = false;
  auto *ctx = loadOp.getContext();
  if (auto subView = base.getDefiningOp<memref::SubViewOp>()) {
    if (!subView.hasUnitStride())
      return false;
    baseMemref = cast<MemrefValue>(subView.getSource());
    offsets = llvm::to_vector(subView.getMixedOffsets());
    droppedDims = subView.getDroppedDims();
    hasSubview = true;
  } else {
    baseMemref = dyn_cast<MemrefValue>(base);
    if (!baseMemref)
      return false;
  }

  auto loadIndices = loadOp.getIndices();
  unsigned baseRank = baseMemref.getType().getRank();
  if ((loadOp.getMemref() != baseMemref) || (loadIndices.size() != baseRank))
    return false;

  unsigned writeRank = writeIndices.size();
  if ((!hasSubview && writeRank != baseRank) ||
      (hasSubview && offsets.size() != baseRank) ||
      (vectorDimForWriteDim.size() != writeRank))
    return false;

  auto zeroAttr = IntegerAttr::get(IndexType::get(ctx), 0);
  unsigned writeMemrefDim = 0;
  for (unsigned baseDim : llvm::seq(baseRank)) {
    bool wasDropped = (hasSubview && droppedDims.test(baseDim));
    int64_t vectorDim = !wasDropped ? vectorDimForWriteDim[writeMemrefDim] : -1;
    int64_t extent = 1;
    if (vectorDim >= 0) {
      int64_t dimSize = vecTy.getDimSize(vectorDim);
      if (dimSize == ShapedType::kDynamic)
        return false;
      extent = dimSize;
    }
    Value writeIndex = !wasDropped ? writeIndices[writeMemrefDim] : Value();
    OpFoldResult offset =
        hasSubview ? offsets[baseDim] : OpFoldResult(zeroAttr);
    if (!loadIndexWithinWriteRange(loadIndices[baseDim], offset, writeIndex,
                                   extent, ivsMap))
      return false;
    if (!wasDropped)
      ++writeMemrefDim;
  }

  return true;
}

/// Recognize scalar memref.load of an element produced by a
/// vector.transfer_write
static bool loadMatchesVectorWrite(memref::LoadOp loadOp,
                                   vector::TransferWriteOp writeOp,
                                   const IRMapping &ivsMap) {
  auto vecTy = dyn_cast<VectorType>(writeOp.getVector().getType());
  if (!vecTy)
    return false;

  unsigned writeRank = writeOp.getIndices().size();
  AffineMap permutationMap = writeOp.getPermutationMap();
  if (!permutationMap.isProjectedPermutation() ||
      permutationMap.getNumResults() != vecTy.getRank() ||
      permutationMap.getNumDims() != writeRank)
    return false;

  SmallVector<int64_t> vectorDimForWriteDim(writeRank, -1);
  for (unsigned vecDim = 0; vecDim < permutationMap.getNumResults(); ++vecDim) {
    auto dimExpr = dyn_cast<AffineDimExpr>(permutationMap.getResult(vecDim));
    if (!dimExpr)
      return false;
    unsigned writeDim = dimExpr.getPosition();
    if (writeDim >= writeRank || vectorDimForWriteDim[writeDim] != -1)
      return false;
    vectorDimForWriteDim[writeDim] = vecDim;
  }

  return isLoadOnWrittenVector(loadOp, writeOp.getBase(), writeOp.getIndices(),
                               vecTy, vectorDimForWriteDim, ivsMap);
}

/// Recognize scalar memref.load of an element produced by a vector.store
static bool loadMatchesVectorStore(memref::LoadOp loadOp,
                                   vector::StoreOp storeOp,
                                   const IRMapping &ivsMap) {
  auto vecTy = dyn_cast<VectorType>(storeOp.getValueToStore().getType());
  if (!vecTy)
    return false;

  unsigned writeRank = storeOp.getIndices().size();
  if (vecTy.getRank() > writeRank)
    return false;

  SmallVector<int64_t> vectorDimForWriteDim(writeRank, -1);
  unsigned vecRank = vecTy.getRank();
  for (unsigned i = 0; i < vecRank; ++i) {
    unsigned writeDim = writeRank - vecRank + i;
    vectorDimForWriteDim[writeDim] = i;
  }

  return isLoadOnWrittenVector(loadOp, storeOp.getBase(), storeOp.getIndices(),
                               vecTy, vectorDimForWriteDim, ivsMap);
}

/// Check if both operations access the same positions of the same
/// buffer, but one of the two does it through a rank-reducing full subview of
/// the buffer (the other's base). EXAMPLE:
///  memref.store %a, %buf[%c0, %i, %j] : memref<1x2x2xf32>
///  %alias = memref.subview %buf[0, 0, 0][1, 2, 2][1, 1, 1]: memref<1x2x2xf32>
///                           to memref<2x2xf32>
///  %val = memref.load %alias[%i, %j] : memref<2x2xf32>
template <typename OpTy1, typename OpTy2>
static bool opsAccessSameIndicesViaRankReducingSubview(
    OpTy1 op1, OpTy2 op2, const IRMapping &firstToSecondPloopIVsMap,
    OpBuilder &b) {
  auto base1 = cast<MemrefValue>(getBaseMemref(op1));
  auto base2 = cast<MemrefValue>(getBaseMemref(op2));
  if (!base1 || !base2)
    return false;

  auto accessThroughTrivialSubviewIsSame =
      [&b](memref::SubViewOp subView, ValueRange subViewAccess,
           ValueRange sourceAccess, const IRMapping &ivsMap) -> bool {
    SmallVector<Value> resolvedSubviewAccess;
    LogicalResult resolved = resolveSourceIndicesRankReducingSubview(
        subView.getLoc(), b, subView, subViewAccess, resolvedSubviewAccess);
    if (failed(resolved) ||
        (resolvedSubviewAccess.size() != sourceAccess.size()))
      return false;
    for (auto [dimIdx, resolvedIndex] :
         llvm::enumerate(resolvedSubviewAccess)) {
      if (!matchPattern(resolvedIndex, m_Zero()) &&
          !valsAreEquivalent(resolvedIndex, sourceAccess[dimIdx], ivsMap))
        return false;
    }
    return true;
  };

  // Case 1: op1 uses a subview of op2's base.
  if (auto subView = base1.template getDefiningOp<memref::SubViewOp>();
      subView &&
      memref::isSameViewOrTrivialAlias(
          base2, cast<MemrefValue>(subView.getSource())) &&
      accessThroughTrivialSubviewIsSame(subView, op1.getIndices(),
                                        op2.getIndices(),
                                        firstToSecondPloopIVsMap))
    return true;

  // Case 2: op2 uses a subview of op1's base.
  if (auto subView = base2.template getDefiningOp<memref::SubViewOp>();
      subView &&
      memref::isSameViewOrTrivialAlias(
          base1, cast<MemrefValue>(subView.getSource())) &&
      accessThroughTrivialSubviewIsSame(subView, op2.getIndices(),
                                        op1.getIndices(),
                                        firstToSecondPloopIVsMap))
    return true;

  return false;
}

/// Check if both memory read/write operations access the same indices
/// (considering also the mapping of induction variables from the first to the
/// second parallel loop).
template <typename OpTy1, typename OpTy2>
static bool opsAccessSameIndices(OpTy1 op1, OpTy2 op2,
                                 const IRMapping &loopsIVsMap, OpBuilder &b) {
  auto indices1 = op1.getIndices();
  auto indices2 = op2.getIndices();
  if (indices1.size() != indices2.size())
    return opsAccessSameIndicesViaRankReducingSubview(op1, op2, loopsIVsMap, b);
  for (auto [idx1, idx2] : llvm::zip(indices1, indices2)) {
    if (!valsAreEquivalent(idx1, idx2, loopsIVsMap))
      return false;
  }
  return true;
}

/// Check if the loadOp reads from the same memory location (same buffer,
/// same indices and same properties) as written by the storeOp.
static bool
loadsFromSameMemoryLocationWrittenBy(Operation *loadOp, Operation *storeOp,
                                     const IRMapping &firstToSecondPloopIVsMap,
                                     OpBuilder &b) {
  if (!loadOp || !storeOp)
    return false;
  // Support only these memory-reading ops for now
  if (!isa<memref::LoadOp, vector::TransferReadOp, vector::LoadOp>(loadOp))
    return false;
  bool accessSameMemory =
      llvm::TypeSwitch<Operation *, bool>(loadOp)
          .Case([&](memref::LoadOp memLoadOp) {
            if (auto memStoreOp = dyn_cast<memref::StoreOp>(storeOp))
              return opsAccessSameIndices(memLoadOp, memStoreOp,
                                          firstToSecondPloopIVsMap, b);
            if (auto vecWriteOp = dyn_cast<vector::TransferWriteOp>(storeOp))
              return loadMatchesVectorWrite(memLoadOp, vecWriteOp,
                                            firstToSecondPloopIVsMap);
            if (auto vecStoreOp = dyn_cast<vector::StoreOp>(storeOp))
              return loadMatchesVectorStore(memLoadOp, vecStoreOp,
                                            firstToSecondPloopIVsMap);
            return false;
          })
          .Case([&](vector::TransferReadOp vecReadOp) {
            auto vecWriteOp = dyn_cast<vector::TransferWriteOp>(storeOp);
            if (!vecWriteOp)
              return false;
            return opsAccessSameIndices(vecReadOp, vecWriteOp,
                                        firstToSecondPloopIVsMap, b) &&
                   (vecReadOp.getMask() == vecWriteOp.getMask()) &&
                   (vecReadOp.getInBounds() == vecWriteOp.getInBounds());
          })
          .Case([&](vector::LoadOp vecLoadOp) {
            auto vecStoreOp = dyn_cast<vector::StoreOp>(storeOp);
            if (!vecStoreOp)
              return false;
            return opsAccessSameIndices(vecLoadOp, vecStoreOp,
                                        firstToSecondPloopIVsMap, b) &&
                   (vecLoadOp.getAlignment() == vecStoreOp.getAlignment());
          })
          .Default([](Operation *) { return false; });
  return accessSameMemory;
}

static Value getStoreOpTargetBuffer(Operation *op) {
  return llvm::TypeSwitch<Operation *, Value>(op)
      .Case([&](memref::StoreOp storeOp) { return storeOp.getMemRef(); })
      .Case([&](vector::TransferWriteOp writeOp) { return writeOp.getBase(); })
      .Case([&](vector::StoreOp vecStoreOp) { return vecStoreOp.getBase(); })
      .Default([](Operation *) { return Value(); });
}

/// To be called when `mayAlias(val1, val2)` is true. Check if the potential
/// aliasing between the loadOp and storeOp can be resolved by analyzing their
/// access patterns.
static bool canResolveAlias(Operation *loadOp, Operation *storeOp,
                            const IRMapping &loopsIVsMap) {
  if (auto transfWriteOp = dyn_cast<vector::TransferWriteOp>(storeOp);
      transfWriteOp && isa<memref::LoadOp>(loadOp))
    return loadMatchesVectorWrite(cast<memref::LoadOp>(loadOp), transfWriteOp,
                                  loopsIVsMap);
  if (auto vecStoreOp = dyn_cast<vector::StoreOp>(storeOp);
      vecStoreOp && isa<memref::LoadOp>(loadOp))
    return loadMatchesVectorStore(cast<memref::LoadOp>(loadOp), vecStoreOp,
                                  loopsIVsMap);
  return false;
}

/// Check that the parallel loops have no mixed access to the same buffers.
/// Return `true` if the second parallel loop does not read or write the buffers
/// written by the first loop using different indices.
static bool haveNoDataDependenciesExceptSameIndex(
    ParallelOp firstPloop, ParallelOp secondPloop,
    const IRMapping &firstToSecondPloopIndices,
    llvm::function_ref<bool(Value, Value)> mayAlias, OpBuilder &b) {
  // Map buffers to their store/write ops in the firstPloop
  DenseMap<Value, SmallVector<Operation *>> bufferStoresInFirstPloop;
  // Record all the memory buffers used in store/write ops found in firstPloop
  llvm::SmallSetVector<Value, 4> buffersWrittenInFirstPloop;

  auto collectStoreOpsInWalk = [&](Operation *op) {
    auto memOpInterf = dyn_cast_if_present<MemoryEffectOpInterface>(op);
    // Ignore ops that don't write to memory
    if (!memOpInterf || (!memOpInterf.hasEffect<MemoryEffects::Write>() &&
                         !memOpInterf.hasEffect<MemoryEffects::Free>()))
      return WalkResult::advance();

    // Only these memory-writing ops are supported for now:
    // memref.store, vector.transfer_write, vector.store
    Value storeOpBase = getStoreOpTargetBuffer(op);
    if (!storeOpBase)
      return WalkResult::interrupt();

    // Expect the base operand to be a Memref
    MemrefValue storeOpBaseMemref = dyn_cast<MemrefValue>(storeOpBase);
    if (!storeOpBaseMemref)
      return WalkResult::interrupt();
    // Get the original memref buffer, skipping full view-like ops
    Value buffer = memref::skipFullyAliasingOperations(storeOpBaseMemref);
    bufferStoresInFirstPloop[buffer].push_back(op);
    buffersWrittenInFirstPloop.insert(buffer);
    return WalkResult::advance();
  };

  // Walk the first parallel loop to collect all store/write ops and their
  // target buffers
  if (firstPloop.getBody()->walk(collectStoreOpsInWalk).wasInterrupted())
    return false;

  // Check that this load/read op encountered while walking the second parallel
  // loop does not have incompatible data dependencies with the store/write ops
  // collected from the first parallel loop: the loops can be fused only if in
  // the 2nd loop there are no loads/stores from/to the buffers written in the
  // 1st loop, except when on the same exact memory location (same indices) as
  // written in the 1st loop.
  auto checkLoadInWalkHasNoIncompatibleDataDeps = [&](Operation *loadOp) {
    auto memOpInterf = dyn_cast_if_present<MemoryEffectOpInterface>(loadOp);
    // To be conservative, we should stop on ops that don't advertise their
    // memory effects. However, many ops don't implement MemoryEffectOpInterface
    // yet, so for now we just skip them.
    // TODO: once more ops add MemoryEffectOpInterface, interrupt the walk here.
    if (!memOpInterf &&
        !loadOp->hasTrait<mlir::OpTrait::HasRecursiveMemoryEffects>())
      return WalkResult::advance();
    // Ignore ops that don't read from memory, and wrapping ops that have nested
    // memory effects (e.g. loops, conditionals) as they will be analyzed when
    // visiting their nested ops.
    if ((!memOpInterf &&
         loadOp->hasTrait<mlir::OpTrait::HasRecursiveMemoryEffects>()) ||
        (memOpInterf && !memOpInterf.hasEffect<MemoryEffects::Read>()))
      return WalkResult::advance();
    // Support only these memory-reading ops for now
    if (!isa<memref::LoadOp, vector::TransferReadOp, vector::LoadOp>(loadOp) ||
        !isa<MemrefValue>(loadOp->getOperand(0)))
      return WalkResult::interrupt();

    MemrefValue loadOpBase = cast<MemrefValue>(loadOp->getOperand(0));
    MemrefValue loadedOrigBuf = memref::skipFullyAliasingOperations(loadOpBase);

    for (Value storedMem : buffersWrittenInFirstPloop)
      if ((storedMem != loadedOrigBuf) && mayAlias(storedMem, loadedOrigBuf) &&
          !llvm::all_of(bufferStoresInFirstPloop[storedMem],
                        [&](Operation *storeOp) {
                          return canResolveAlias(loadOp, storeOp,
                                                 firstToSecondPloopIndices);
                        })) {
        return WalkResult::interrupt();
      }

    auto writeOpsIt = bufferStoresInFirstPloop.find(loadedOrigBuf);
    if (writeOpsIt == bufferStoresInFirstPloop.end())
      return WalkResult::advance();
    // Store/write ops to this buffer in the firstPloop
    SmallVector<mlir::Operation *> &writeOps = writeOpsIt->second;

    // If the first loop has no writes to this buffer, continue
    if (writeOps.empty())
      return WalkResult::advance();

    Operation *writeOp = writeOps.front();

    // In the first parallel loop, multiple writes to the same memref are
    // allowed only on the same memory location
    if (!llvm::all_of(writeOps, [&](Operation *otherWriteOp) {
          return opsWriteSameMemLocation(writeOp, otherWriteOp);
        })) {
      return WalkResult::interrupt();
    }

    // Check that the load in secondPloop reads from the same memory location as
    // written by the corresponding store in firstPloop
    if (!loadsFromSameMemoryLocationWrittenBy(loadOp, writeOp,
                                              firstToSecondPloopIndices, b)) {
      return WalkResult::interrupt();
    }

    return WalkResult::advance();
  };

  // Walk the second parallel loop to check load/read ops against the stores
  // collected from the first parallel loop.
  return !secondPloop.getBody()
              ->walk(checkLoadInWalkHasNoIncompatibleDataDeps)
              .wasInterrupted();
}

/// Check that in each loop there are no read ops on the buffers written
/// by the other loop, except when reading from the same exact memory location
/// (same indices) as written in the other loop.
static bool
noIncompatibleDataDependencies(ParallelOp firstPloop, ParallelOp secondPloop,
                               const IRMapping &firstToSecondPloopIndices,
                               llvm::function_ref<bool(Value, Value)> mayAlias,
                               OpBuilder &b) {
  if (!haveNoDataDependenciesExceptSameIndex(
          firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias, b))
    return false;

  IRMapping secondToFirstPloopIndices;
  secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),
                                firstPloop.getBody()->getArguments());
  return haveNoDataDependenciesExceptSameIndex(
      secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias, b);
}

/// Check if fusion of the two parallel loops is legal:
/// i.e. no nested parallel loops, equal iteration spaces,
/// and no incompatible data dependencies between the loops.
static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
                          const IRMapping &firstToSecondPloopIndices,
                          llvm::function_ref<bool(Value, Value)> mayAlias,
                          OpBuilder &b) {
  if (hasNestedParallelOp(firstPloop) || hasNestedParallelOp(secondPloop) ||
      !equalIterationSpaces(firstPloop, secondPloop) ||
      !noIncompatibleDataDependencies(firstPloop, secondPloop,
                                      firstToSecondPloopIndices, mayAlias, b))
    return false;

  // We are fusing first loop into second, make sure there are no users of the
  // first loop results between loops.
  DominanceInfo dom;
  for (Operation *user : firstPloop->getUsers()) {
    if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
      return false;
  }
  return true;
}

// Returns new parallel loop where two loops matching indices param are
// interchanged
static std::optional<ParallelOp>
interchangeLoops(OpBuilder &builder, ParallelOp &loop,
                 const ArrayRef<int64_t> &indices) {
  assert(loop.getNumLoops() == indices.size());
  if (loop.getNumLoops() < 2)
    return std::nullopt;

  // Replace the parallel loop with the same parallel loop.
  builder.setInsertionPoint(loop);
  SmallVector<Value> newLB =
      applyPermutation(SmallVector<Value>(loop.getLowerBound()), indices);
  SmallVector<Value> newUB =
      applyPermutation(SmallVector<Value>(loop.getUpperBound()), indices);
  SmallVector<Value> newStep =
      applyPermutation(SmallVector<Value>(loop.getStep()), indices);
  auto newOp =
      ParallelOp::create(builder, loop.getLoc(), newLB, newUB, newStep,
                         loop.getInitVals(), nullptr, loop.getUnsignedCmp());
  auto ivs = loop.getInductionVars();
  SmallVector<Value> newIvs = applyPermutation(
      newOp.getInductionVars(), invertPermutationVector(indices));
  IRMapping mapping;
  for (auto [iv, riv] : llvm::zip(ivs, newIvs)) {
    mapping.map(iv, riv);
  }

  // Copy parallel loop body
  auto b = OpBuilder::atBlockBegin(newOp.getBody());
  for (auto &o : loop.getNumReductions()
                     ? loop.getBodyRegion().front()
                     : loop.getBodyRegion().front().without_terminator()) {
    b.clone(o, mapping);
  }
  return newOp;
}

struct LoopIV {
  Value lBound, uBound, step;
  bool operator!=(LoopIV const &other) const { return !(*this == other); }
  bool operator==(LoopIV const &other) const {
    return lBound == other.lBound && uBound == other.uBound &&
           step == other.step;
  }
};

template <>
struct llvm::DenseMapInfo<LoopIV> {
  static inline bool isEqual(const LoopIV &lhs, const LoopIV &rhs) {
    return (lhs == rhs);
  }

  static inline unsigned getHashValue(const LoopIV &val) {
    return llvm::hash_combine(
        DenseMapInfo<mlir::Value>::getHashValue(val.lBound),
        DenseMapInfo<mlir::Value>::getHashValue(val.uBound),
        DenseMapInfo<mlir::Value>::getHashValue(val.step));
  }
};

// Returns vector of candidate permutation indices vectors,
// can be empty. Caps the number of extra candidate permutations
// explored to avoid combinatorial explosion. This makes the search
// intentionally incomplete.
static SmallVector<SmallVector<int64_t>>
computeCandidateInterchangePermutations(ParallelOp &firstPloop,
                                        ParallelOp &secondPloop,
                                        int permBudget = 120) {
  // Check preconditions
  if (firstPloop.getNumLoops() < 2 ||
      firstPloop.getNumLoops() != secondPloop.getNumLoops())
    return {};

  SmallVector<LoopIV> firstIVs(firstPloop.getNumLoops());
  SmallVector<LoopIV> secondIVs(secondPloop.getNumLoops());
  llvm::SmallSetVector<LoopIV, 6> unique;
  for (unsigned index : llvm::seq(firstPloop.getNumLoops())) {
    firstIVs[index].lBound = firstPloop.getLowerBound()[index];
    firstIVs[index].uBound = firstPloop.getUpperBound()[index];
    firstIVs[index].step = firstPloop.getStep()[index];
    secondIVs[index].lBound = secondPloop.getLowerBound()[index];
    secondIVs[index].uBound = secondPloop.getUpperBound()[index];
    secondIVs[index].step = secondPloop.getStep()[index];
    unique.insert(firstIVs[index]);
  }

  SmallVector<bool> diffIVs(firstPloop.getNumLoops());
  llvm::transform(
      llvm::zip(firstIVs, secondIVs), diffIVs.begin(),
      [](auto const &pair) { return std::get<0>(pair) != std::get<1>(pair); });

  SmallVector<int64_t> indices;
  for (auto [idx, val] : enumerate(diffIVs))
    if (val)
      indices.push_back(idx);

  // Not a permutation shortcut
  if (indices.size() == 1)
    return {};

  // Initialize with identity permutations
  SmallVector<int64_t> basic(firstIVs.size());
  std::iota(basic.begin(), basic.end(), 0);

  if (indices.empty() && unique.size() == firstIVs.size())
    return {};

  if (indices.size() > 1) {
    // Determine whether the iteration space of the first loop is a permutation
    // of the second and collect remaps.
    SmallVector<int64_t> remaps;
    for (auto fIdx : indices) {
      for (auto sIdx : indices) {
        // can be remapped
        if (fIdx != sIdx && firstIVs[fIdx] == secondIVs[sIdx] &&
            remaps.end() == std::find(remaps.begin(), remaps.end(), sIdx)) {
          remaps.push_back(sIdx);
          break;
        }
      }
    }

    // Not a permutation
    if (indices.size() != remaps.size())
      return {};

    // compose permutation indices
    for (auto [from, to] : zip(indices, remaps)) {
      basic[from] = to;
    }

    LDBG() << "Collected basic permutations: "
           << llvm::interleaved_array(basic);

    // All axes are unique, no further permutatons needed
    if (unique.size() == firstIVs.size()) {
      return {basic};
    }
  }

  //
  // Permute equal axes
  assert(unique.size() != firstIVs.size() &&
         "Expected at least two equal axes");

  // Collect equal axes to groups
  SmallVector<SmallVector<int64_t>> extraResults{basic};
  SmallVector<SmallVector<int64_t>> groups;
  for (auto iv : unique) {
    SmallVector<int64_t> group;
    for (unsigned index : llvm::seq(firstIVs.size())) {
      if (firstIVs[index] == iv)
        group.push_back(index);
    }
    if (group.size() > 1)
      groups.push_back(std::move(group));
  }

  // Permute axes groups
  SmallVector<SmallVector<int64_t>> rmpdGroups(groups);
  bool repeat = true;
  while (repeat && permBudget) {
    repeat = false;
    for (auto const &[group, groupRemaps] : zip(groups, rmpdGroups)) {
      repeat |= std::next_permutation(groupRemaps.begin(), groupRemaps.end());
      if (repeat)
        break;
    }

    if (repeat) {
      SmallVector<int64_t> extra(basic);
      for (auto const &[group, groupRemaps] : zip(groups, rmpdGroups)) {
        for (auto [from, to] : zip(group, groupRemaps))
          extra[from] = basic[to];
      }
      if (basic != extra) {
        LDBG() << "Collected extra permutations: "
               << llvm::interleaved_array(extra);

        extraResults.push_back(std::move(extra));
        permBudget--;
      }
    }
  }

  return extraResults;
}

/// Prepend operations of firstPloop's body into secondPloop's body.
/// Update secondPloop with new loop.
static void applyLoopFusion(ParallelOp &firstPloop, ParallelOp &secondPloop,
                            OpBuilder &builder) {
  Block *block1 = firstPloop.getBody();
  Block *block2 = secondPloop.getBody();
  ValueRange inits1 = firstPloop.getInitVals();
  ValueRange inits2 = secondPloop.getInitVals();

  SmallVector<Value> newInitVars(inits1.begin(), inits1.end());
  newInitVars.append(inits2.begin(), inits2.end());

  IRRewriter b(builder);
  b.setInsertionPoint(secondPloop);
  auto newSecondPloop =
      ParallelOp::create(b, secondPloop.getLoc(), secondPloop.getLowerBound(),
                         secondPloop.getUpperBound(), secondPloop.getStep(),
                         newInitVars, nullptr, secondPloop.getUnsignedCmp());

  Block *newBlock = newSecondPloop.getBody();
  auto term1 = cast<ReduceOp>(block1->getTerminator());
  auto term2 = cast<ReduceOp>(block2->getTerminator());

  b.inlineBlockBefore(block2, newBlock, newBlock->begin(),
                      newBlock->getArguments());
  b.inlineBlockBefore(block1, newBlock, newBlock->begin(),
                      newBlock->getArguments());

  ValueRange results = newSecondPloop.getResults();
  if (!results.empty()) {
    b.setInsertionPointToEnd(newBlock);

    ValueRange reduceArgs1 = term1.getOperands();
    ValueRange reduceArgs2 = term2.getOperands();
    SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());
    newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());

    auto newReduceOp = scf::ReduceOp::create(b, term2.getLoc(), newReduceArgs);

    for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(
             term1.getReductions(), term2.getReductions()))) {
      Block &oldRedBlock = reg.front();
      Block &newRedBlock = newReduceOp.getReductions()[i].front();
      b.inlineBlockBefore(&oldRedBlock, &newRedBlock, newRedBlock.begin(),
                          newRedBlock.getArguments());
    }

    firstPloop.replaceAllUsesWith(results.take_front(inits1.size()));
    secondPloop.replaceAllUsesWith(results.take_back(inits2.size()));
  }
  term1->erase();
  term2->erase();
  firstPloop.erase();
  secondPloop.erase();
  secondPloop = newSecondPloop;
}

/// Check fusion pre-conditions and call fusion if it is possible
static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
                        OpBuilder builder,
                        llvm::function_ref<bool(Value, Value)> mayAlias) {
  Block *block1 = firstPloop.getBody();
  Block *block2 = secondPloop.getBody();
  IRMapping firstToSecondPloopIndices;
  firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());

  if (isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
                    mayAlias, builder)) {
    applyLoopFusion(firstPloop, secondPloop, builder);
    return;
  }

  // If iteration space of the second parallel loop is a permutation of the
  // first one then interchange iteration space of the second parallel loop
  // and re-asses possibility of fusion.
  for (auto &perms :
       computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
    OpBuilder::InsertionGuard guard(builder);
    LDBG() << "Applied permutation: " << llvm::interleaved_array(perms);

    auto newLoop = interchangeLoops(builder, secondPloop, perms);
    firstToSecondPloopIndices.clear();
    firstToSecondPloopIndices.map(block1->getArguments(),
                                  newLoop->getBody()->getArguments());
    if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
                       mayAlias, builder)) {
      LDBG() << "Rejected: " << newLoop;

      newLoop->erase();
      continue;
    }

    secondPloop.replaceAllUsesWith(newLoop->getResults());
    secondPloop->erase();
    secondPloop = *newLoop;
    applyLoopFusion(firstPloop, secondPloop, builder);
    break;
  }
}

void mlir::scf::naivelyFuseParallelOps(
    Region &region, llvm::function_ref<bool(Value, Value)> mayAlias) {
  OpBuilder b(region);
  // Consider every single block and attempt to fuse adjacent loops.
  SmallVector<SmallVector<ParallelOp>, 1> ploopChains;
  for (auto &block : region) {
    ploopChains.clear();
    ploopChains.push_back({});

    // Not using `walk()` to traverse only top-level parallel loops and also
    // make sure that there are no side-effecting ops between the parallel
    // loops.
    bool noSideEffects = true;
    for (auto &op : block) {
      if (auto ploop = dyn_cast<ParallelOp>(op)) {
        if (noSideEffects) {
          ploopChains.back().push_back(ploop);
        } else {
          ploopChains.push_back({ploop});
          noSideEffects = true;
        }
        continue;
      }
      // TODO: Handle region side effects properly.
      noSideEffects &= isMemoryEffectFree(&op) && op.getNumRegions() == 0;
    }
    for (MutableArrayRef<ParallelOp> ploops : ploopChains) {
      for (int i = 0, e = ploops.size(); i + 1 < e; ++i)
        fuseIfLegal(ploops[i], ploops[i + 1], b, mayAlias);
    }
  }
}

namespace {
struct ParallelLoopFusion
    : public impl::SCFParallelLoopFusionBase<ParallelLoopFusion> {
  void runOnOperation() override {
    auto &aa = getAnalysis<AliasAnalysis>();

    auto mayAlias = [&](Value val1, Value val2) -> bool {
      // If the memref is defined in one of the parallel loops body, careful
      // alias analysis is needed.
      // TODO: check if this is still needed as a separate check.
      auto val1Def = val1.getDefiningOp();
      auto val2Def = val2.getDefiningOp();
      auto val1Loop =
          val1Def ? val1Def->getParentOfType<ParallelOp>() : nullptr;
      auto val2Loop =
          val2Def ? val2Def->getParentOfType<ParallelOp>() : nullptr;
      if (val1Loop != val2Loop)
        return true;

      return !aa.alias(val1, val2).isNo();
    };

    getOperation()->walk([&](Operation *child) {
      for (Region &region : child->getRegions())
        naivelyFuseParallelOps(region, mayAlias);
    });
  }
};
} // namespace

std::unique_ptr<Pass> mlir::createParallelLoopFusionPass() {
  return std::make_unique<ParallelLoopFusion>();
}