//===- ReshapeOpsUtils.cpp - Utilities used by structured ops -------------===//
//
// 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 "mlir/Dialect/Utils/ReshapeOpsUtils.h"

#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"

#include <numeric>
#include <optional>

using namespace mlir;

std::optional<SmallVector<ReassociationIndices>>
mlir::getReassociationIndicesForReshape(ShapedType sourceType,
                                        ShapedType targetType) {
  if (sourceType.getRank() > targetType.getRank())
    return getReassociationIndicesForCollapse(sourceType.getShape(),
                                              targetType.getShape());
  if (sourceType.getRank() < targetType.getRank())
    return getReassociationIndicesForCollapse(targetType.getShape(),
                                              sourceType.getShape());
  return std::nullopt;
}

namespace {
/// A simple struct to represent ReassociationIndices as an inclusive interval.
/// It's designed to be feasibly minimal, so the call sites should manage the
/// validity of the range manually.
struct ReassociationIndexRange {
  /// FIXME: Signed type is used for consistency with ReassociationIndices.
  /// We should consider refactoring all reassociation utilities to use unsigned
  /// types.
  int64_t leftIdx = 0, rightIdx = 0;

  /// Util for manual checks of the range's validity
  LogicalResult verify() const {
    return leftIdx >= 0 && (leftIdx <= rightIdx) ? success() : failure();
  }

  /// Checks range's containment within another range. Treats the edges
  /// non-exclusively.
  bool isInRange(const ReassociationIndexRange &outerRange) const {
    return leftIdx >= outerRange.leftIdx && rightIdx <= outerRange.rightIdx;
  }

  unsigned size() const {
    assert(succeeded(verify()));
    return rightIdx - leftIdx + 1;
  }
  bool containsSingleIndex() const { return size() == 1; }

  /// Collects indices that do not overlap between this and another range.
  ReassociationIndices
  getNonOverlappingIndicesWith(ReassociationIndexRange &rhs) const {
    if (rightIdx < rhs.leftIdx) {
      // The intervals do not overlap - concatenate the indices from both.
      auto jointFullIndices = getFullIndices();
      jointFullIndices.append(rhs.getFullIndices());
      return jointFullIndices;
    }
    ReassociationIndices result;
    // Handle the chunk left of the overlapping range.
    int64_t leftStart = std::min(leftIdx, rhs.leftIdx);
    int64_t leftEnd = std::max(leftIdx, rhs.leftIdx);
    llvm::append_range(result, llvm::seq(leftStart, leftEnd));
    // Handle the chunk right of the overlapping range. Symmetrically, we should
    // skip the edge of the overlap AND include the rightmost index.
    int64_t rightStart = std::min(rightIdx, rhs.rightIdx) + 1;
    int64_t rightEnd = std::max(rightIdx, rhs.rightIdx);
    if (rightStart < rightEnd)
      llvm::append_range(result, llvm::seq_inclusive(rightStart, rightEnd));
    return result;
  }

  /// Converts the range into ReassociationIndices.
  ReassociationIndices getFullIndices() const {
    ReassociationIndices result;
    for (int64_t idx = leftIdx; idx <= rightIdx; ++idx) {
      result.push_back(idx);
    }
    return result;
  }
};
} // namespace

/// Starting from `sourceStartIdx`, searches `sourceShape` for the first
/// sequence that can be collapsed into a dynamic dimension (at least one must
/// be present in the source).
/// By default, lazily returns once the first dynamic dimension has been found.
/// Setting `matchGreedily` as `true` will also mark all subsequent
/// source dimensions for collapsing into the target.
static FailureOr<ReassociationIndexRange>
findReassociationRangeForDynamicDim(ArrayRef<int64_t> sourceShape,
                                    int64_t sourceStartIdx,
                                    bool matchGreedily = false) {
  const unsigned numSourceDims = sourceShape.size();
  ReassociationIndexRange sourceShapeAsRange{0, numSourceDims - 1};
  std::optional<ReassociationIndexRange> resultRange = std::nullopt;

  ReassociationIndexRange iterationRange{sourceStartIdx, sourceStartIdx};
  for (; iterationRange.isInRange(sourceShapeAsRange);
       iterationRange.rightIdx++) {
    int64_t sourceSize = sourceShape[iterationRange.rightIdx];
    if (sourceSize == ShapedType::kDynamic) {
      resultRange = iterationRange;
      break;
    }
  }
  if (!resultRange)
    return failure();
  if (matchGreedily)
    resultRange->rightIdx = sourceShapeAsRange.rightIdx;
  return *resultRange;
}

/// Starting from `sourceStartIdx`, searches `sourceShape` for the first
/// sequence of static dimensions such that their product matches `targetSize`.
/// By default, lazily returns once the product matches the target size. Setting
/// `matchGreedily` as `true` will append all neighboring unit dimensions
/// (dimensions of 1) to the match.
static FailureOr<ReassociationIndexRange>
findReassociationRangeForSize(ArrayRef<int64_t> sourceShape,
                              int64_t sourceStartIdx, int64_t targetSize,
                              bool matchGreedily = false) {
  const unsigned numSourceDims = sourceShape.size();
  ReassociationIndexRange sourceShapeAsRange{0, numSourceDims - 1};
  std::optional<ReassociationIndexRange> resultRange = std::nullopt;

  ReassociationIndexRange iterationRange{sourceStartIdx, sourceStartIdx};
  int64_t prodOfCollapsedDims = 1;
  while (iterationRange.isInRange(sourceShapeAsRange)) {
    int64_t sourceSize = sourceShape[iterationRange.rightIdx];
    if (sourceSize == ShapedType::kDynamic) {
      // Reassociation for a static dim cannot include a dynamic dim. Reset
      // induction variables to essentially restart the loop from the next
      // source dimension.
      prodOfCollapsedDims = 1;
      iterationRange = {iterationRange.rightIdx + 1,
                        iterationRange.rightIdx + 1};
      continue;
    }
    prodOfCollapsedDims *= sourceSize;
    // If the target size has been exceeded without matching, we need to shift
    // the range start right. From the start of the range, roll back the
    // multiplication until the target size exceeds the product again.
    while (prodOfCollapsedDims > targetSize &&
           !iterationRange.containsSingleIndex()) {
      int64_t frontSourceSize = sourceShape[iterationRange.leftIdx];
      prodOfCollapsedDims /= frontSourceSize;
      // Shrink the range rightwards
      iterationRange.leftIdx++;
    }
    // We could've reached the target size with the current dimension,
    // also as a result of the above shift to right.
    if (prodOfCollapsedDims == targetSize) {
      resultRange = iterationRange;
      break;
    }
    // Increment the iteration range
    iterationRange.rightIdx++;
  }
  if (!resultRange)
    return failure();
  if (matchGreedily) {
    // We now want to collect all unit dimensions directly after the target
    // product match. Advance the iterator to avoid OOB when the product match
    // happens at the last element.
    iterationRange.rightIdx++;
    while (iterationRange.isInRange(sourceShapeAsRange) &&
           sourceShape[iterationRange.rightIdx] == 1) {
      resultRange = iterationRange;
      iterationRange.rightIdx++;
    }
  }
  return *resultRange;
}

/// Attempts to find a valid collapsing reassociation of `sourceShape` into
/// `targetShape` through a simple traversal. If successful, an array of source
/// index ranges is returned, correspondingly to each dimension in the target
/// shape. The resulting indices shall fully cover the `sourceShape` without
/// overlaps.
///
/// The algorithm is essentially a lazy one, searching for non-greedy matches -
/// it will only yield a greedy match for the last target dimension.
/// FIXME: The algorithm can only backtrack when it needs to append an offset
/// for a static target dimension to the preceding dynamic one (this retains the
/// linear complexity). As feasible, consider adding further backtracking
/// routines to enable more reassociations, e.g.:
/// - ?x2x?x2 into ?x2
static FailureOr<SmallVector<ReassociationIndexRange>>
findReassociationRangesForCollapse(ArrayRef<int64_t> sourceShape,
                                   ArrayRef<int64_t> targetShape) {
  unsigned numSourceDims = sourceShape.size(),
           numTargetDims = targetShape.size();
  assert(numSourceDims > numTargetDims);
  ReassociationIndexRange sourceShapeAsRange{0, numSourceDims - 1};

  SmallVector<ReassociationIndexRange> reassocRanges;
  reassocRanges.reserve(numTargetDims);
  // We'll iterate in strides of 2 to enable pseudo-backtracking for simple
  // cases, e.g.:
  // - ?x2x3x5 into ?x15
  std::optional<int64_t> prevTargetSize = std::nullopt;
  for (unsigned targetDimIdx = 0, sourceDimIdx = 0;
       targetDimIdx < numTargetDims; ++targetDimIdx) {
    int64_t targetSize = targetShape[targetDimIdx];
    // Simply check if there are any subsequent target dimensions left - if not,
    // the match must be made greedily.
    bool shouldMatchGreedily = targetDimIdx == numTargetDims - 1;
    FailureOr<ReassociationIndexRange> sourceRange;
    if (targetSize == ShapedType::kDynamic) {
      sourceRange = findReassociationRangeForDynamicDim(
          sourceShape, sourceDimIdx, shouldMatchGreedily);
    } else {
      sourceRange = findReassociationRangeForSize(
          sourceShape, sourceDimIdx, targetSize, shouldMatchGreedily);
    }

    // Run sanity checks on the returned index range.
    if (failed(sourceRange) || failed(sourceRange->verify()) ||
        !sourceRange->isInRange(sourceShapeAsRange))
      return failure();
    if (sourceRange->leftIdx > sourceDimIdx) {
      // If some source dimensions had to be skipped in order to find a match,
      // they must be collapsed into the directly preceding dynamic dimension.
      if (!prevTargetSize || prevTargetSize != ShapedType::kDynamic)
        return failure();
      reassocRanges.back().rightIdx = sourceRange->leftIdx - 1;
    }

    // Store the gathered information as required for the next iteration.
    prevTargetSize = targetSize;
    sourceDimIdx = sourceRange->rightIdx + 1;
    reassocRanges.push_back(*sourceRange);
  }
  // Fail if the source shape wasn't a full match for the target shape. We only
  // need to check the last recorded index - any other gaps should have been
  // mended by the main loop.
  if (reassocRanges.back().rightIdx < sourceShapeAsRange.rightIdx)
    return failure();
  return reassocRanges;
}

/// A variant of `findReassociationRangesForCollapse(...)` that can also scan
/// the shapes right-to-left.
static FailureOr<SmallVector<ReassociationIndexRange>>
findReassociationRangesForCollapse(ArrayRef<int64_t> sourceShape,
                                   ArrayRef<int64_t> targetShape,
                                   bool iterateRightToLeft) {
  if (!iterateRightToLeft)
    return findReassociationRangesForCollapse(sourceShape, targetShape);
  // NB: To iterate right-to-left, we currently reverse the shapes and then
  // reverse the result back. The reversed shapes must not be temporary, as
  // we're passing through an ArrayRef.
  // FIXME: It would be preferable to avoid the expensive copies. At the moment,
  // this approach is chosen for readability of the main implementation.
  std::vector<int64_t> sourceToReverse = sourceShape.vec(),
                       targetToReverse = targetShape.vec();
  std::reverse(sourceToReverse.begin(), sourceToReverse.end());
  std::reverse(targetToReverse.begin(), targetToReverse.end());
  auto invertedRanges =
      findReassociationRangesForCollapse(sourceToReverse, targetToReverse);
  if (failed(invertedRanges))
    return failure();
  SmallVector<ReassociationIndexRange> &rangesToInvert = *invertedRanges;
  unsigned numSourceDims = sourceShape.size();
  // We have received the ranges for inverted shapes. Now we have to invert
  // the ranges back to correspond with the original source shape.
  for (auto &range : rangesToInvert) {
    int64_t invLeftIdx = range.leftIdx, invRightIdx = range.rightIdx;
    range.leftIdx = numSourceDims - 1 - invRightIdx;
    range.rightIdx = numSourceDims - 1 - invLeftIdx;
  }
  // Also invert the ordering of the ranges to correspond with the original
  // target shape.
  std::reverse(rangesToInvert.begin(), rangesToInvert.end());
  return rangesToInvert;
}

std::optional<SmallVector<ReassociationIndices>>
mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
                                         ArrayRef<int64_t> targetShape) {
  unsigned numSourceDims = sourceShape.size(),
           numTargetDims = targetShape.size();
  // We're supposed to search for a collapsing reassociation. If the sizes
  // match, there's no actual collapsing taking place - it's either a no-op or a
  // `tensor.reshape`-style reassociation (that would be beyond the scope of
  // this utility).
  if (numSourceDims <= numTargetDims)
    return std::nullopt;
  // Early handling for scalar target types. We should report an invalid
  // reassociation for non-unit static dimensions - no chance to collapse these
  // into a scalar.
  if (numTargetDims == 0) {
    for (unsigned sourceDimIdx = 0; sourceDimIdx < numSourceDims;
         ++sourceDimIdx) {
      int64_t sourceSize = sourceShape[sourceDimIdx];
      if (sourceSize != 1 && sourceSize != ShapedType::kDynamic)
        return std::nullopt;
    }
    return SmallVector<ReassociationIndices>{};
  }

  // Collect source ranges by iterating over the target shape left-to-right.
  FailureOr<SmallVector<ReassociationIndexRange>> maybeForwardRanges =
      findReassociationRangesForCollapse(sourceShape, targetShape);
  if (failed(maybeForwardRanges))
    return std::nullopt;
  auto &ranges = *maybeForwardRanges;
  // Now do the same in reverse. We need to get another valid reassociation
  // through some other strategy, and then compare the results in order to
  // disambiguate mixed subshapes, such as:
  // ?x?x? into ?x?, ?x2x? into ?x?, ?x2x3x6x? into ?x6x?
  // This leads us to lose some of the reassociation opportunities that can only
  // be found by iterating in a certain direction, e.g. 2x2x? into 2x? - without
  // backtracking, the algorithm will fail right-to-left. However, this is the
  // best way to preserve correctness.
  FailureOr<SmallVector<ReassociationIndexRange>> maybeReverseRanges =
      findReassociationRangesForCollapse(sourceShape, targetShape,
                                         /*iterateRightToLeft=*/true);
  if (failed(maybeReverseRanges))
    return std::nullopt;
  auto &reverseRanges = *maybeReverseRanges;

  if (ranges.size() != numTargetDims || reverseRanges.size() != numTargetDims)
    return std::nullopt;
  // Now we can check for ambiguity of each target dimension's reassociation. If
  // successful, we put the full indices into our result map for the target
  // shape.
  SmallVector<ReassociationIndices> reassociationMap(numTargetDims);
  for (unsigned targetDimIdx = 0; targetDimIdx < numTargetDims;
       ++targetDimIdx) {
    ReassociationIndexRange &range = ranges[targetDimIdx];
    ReassociationIndexRange &reverseRange = reverseRanges[targetDimIdx];
    // Get non-overlapping indices between the ranges
    ReassociationIndices nonMatchingIndices =
        range.getNonOverlappingIndicesWith(reverseRange);
    // Unit dimensions can be collapsed wherever - this is the only ambiguity
    // that we allow.
    for (int64_t sourceDimIdx : nonMatchingIndices) {
      if (sourceShape[sourceDimIdx] != 1)
        return std::nullopt;
    }
    reassociationMap[targetDimIdx] = range.getFullIndices();
  }
  return reassociationMap;
}

std::optional<SmallVector<ReassociationIndices>>
mlir::composeReassociationIndices(
    ArrayRef<ReassociationIndices> producerReassociations,
    ArrayRef<ReassociationIndices> consumerReassociations,
    MLIRContext *context) {
  SmallVector<ReassociationIndices> composedIndices;
  // Make the producer the larger sized vector. If they are of same size, the
  // resulting reshape is not a supported reshape op.
  if (producerReassociations.size() == consumerReassociations.size())
    return std::nullopt;
  if (producerReassociations.size() < consumerReassociations.size())
    std::swap(producerReassociations, consumerReassociations);

  // Handle the corner case of the result being a rank 0 shaped type. Return an
  // empty reassociation.
  if (consumerReassociations.empty())
    return composedIndices;

  size_t consumerDims = std::accumulate(
      consumerReassociations.begin(), consumerReassociations.end(), 0,
      [](size_t all, ReassociationIndicesRef indices) {
        return all + indices.size();
      });
  if (producerReassociations.size() != consumerDims)
    return std::nullopt;

  for (ReassociationIndicesRef consumerIndices : consumerReassociations) {
    ReassociationIndices reassociations;
    for (int64_t consumerIndex : consumerIndices) {
      llvm::append_range(reassociations, producerReassociations[consumerIndex]);
    }
    composedIndices.push_back(std::move(reassociations));
  }
  return composedIndices;
}

SmallVector<SmallVector<AffineExpr, 2>, 2>
mlir::convertReassociationIndicesToExprs(
    MLIRContext *context, ArrayRef<ReassociationIndices> reassociationIndices) {
  SmallVector<SmallVector<AffineExpr, 2>, 2> reassociationMaps;
  for (const auto &indices : reassociationIndices) {
    SmallVector<AffineExpr, 2> reassociationMap;
    reassociationMap.reserve(indices.size());
    for (int64_t index : indices)
      reassociationMap.push_back(mlir::getAffineDimExpr(index, context));
    reassociationMaps.push_back(std::move(reassociationMap));
  }
  return reassociationMaps;
}

template <typename AffineExprTy>
unsigned getMaxPosOfType(ArrayRef<ReassociationExprs> exprArrays) {
  unsigned pos = 0;
  for (const auto &exprs : exprArrays) {
    for (auto expr : exprs) {
      expr.walk([&pos](AffineExpr e) {
        if (auto d = dyn_cast<AffineExprTy>(e))
          pos = std::max(pos, d.getPosition());
      });
    }
  }
  return pos;
}

ArrayAttr mlir::getReassociationIndicesAttribute(
    Builder &b, ArrayRef<ReassociationIndices> reassociation) {
  SmallVector<Attribute, 4> reassociationAttr =
      llvm::to_vector<4>(llvm::map_range(
          reassociation, [&](const ReassociationIndices &indices) -> Attribute {
            return cast<Attribute>(b.getI64ArrayAttr(indices));
          }));
  return b.getArrayAttr(reassociationAttr);
}

SmallVector<ReassociationIndices, 2> mlir::convertReassociationMapsToIndices(
    ArrayRef<ReassociationExprs> reassociationExprs) {
  SmallVector<ReassociationIndices, 2> reassociationIndices;
  for (const auto &exprs : reassociationExprs) {
    ReassociationIndices indices;
    indices.reserve(exprs.size());
    for (const auto &expr : exprs)
      indices.push_back(cast<AffineDimExpr>(expr).getPosition());
    reassociationIndices.push_back(indices);
  }
  return reassociationIndices;
}

SmallVector<AffineMap, 4>
mlir::getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation) {
  unsigned maxDim = getMaxPosOfType<AffineDimExpr>(reassociation);
  assert(getMaxPosOfType<AffineSymbolExpr>(reassociation) == 0 &&
         "Expected symbol-less expressions");
  SmallVector<AffineMap, 4> maps;
  maps.reserve(reassociation.size());
  for (const auto &exprs : reassociation) {
    assert(!exprs.empty());
    maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
  }
  return maps;
}

bool mlir::isReassociationValid(ArrayRef<AffineMap> reassociation,
                                int *invalidIndex) {
  if (reassociation.empty())
    return true;
  unsigned nDims = reassociation[0].getNumDims();
  unsigned nextExpectedDim = 0;
  for (const auto &it : llvm::enumerate(reassociation)) {
    auto m = it.value();
    if (m.getNumDims() != nDims || m.getNumSymbols() != 0) {
      if (invalidIndex)
        *invalidIndex = it.index();
      return false;
    }
    for (auto e : m.getResults()) {
      auto d = dyn_cast<AffineDimExpr>(e);
      if (!d || d.getPosition() != nextExpectedDim++) {
        if (invalidIndex)
          *invalidIndex = it.index();
        return false;
      }
    }
  }
  if (nextExpectedDim != nDims) {
    if (invalidIndex)
      *invalidIndex = reassociation.size() - 1;
    return false;
  }
  return true;
}

LogicalResult mlir::reshapeLikeShapesAreCompatible(
    function_ref<LogicalResult(const Twine &)> emitError,
    ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
    ArrayRef<ReassociationIndices> reassociationMaps, bool isExpandingReshape) {
  unsigned expandedDimStart = 0;
  for (const auto &map : llvm::enumerate(reassociationMaps)) {
    bool foundDynamicShape = false;
    int64_t linearizedStaticShape = 1;

    for (const auto &dim : llvm::enumerate(
             expandedShape.slice(expandedDimStart, map.value().size()))) {
      if (ShapedType::isDynamic(dim.value()))
        foundDynamicShape = true;
      else
        linearizedStaticShape *= dim.value();
    }
    if (foundDynamicShape) {
      if (ShapedType::isStatic(collapsedShape[map.index()])) {
        return emitError(
            "expected dimension " + Twine(map.index()) +
            " of collapsed type to be dynamic since one or more of the "
            "corresponding dimensions in the expanded type is dynamic");
      }
    } else {
      if (collapsedShape[map.index()] != linearizedStaticShape) {
        return emitError("expected dimension " + Twine(map.index()) +
                         " of collapsed type to be static value of " +
                         Twine(linearizedStaticShape));
      }
    }
    expandedDimStart += map.value().size();
  }
  return success();
}

bool mlir::hasNonIdentityLayout(Type type) {
  if (auto memrefType = dyn_cast<MemRefType>(type))
    return !memrefType.getLayout().isIdentity();
  return false;
}

llvm::SmallBitVector
mlir::getSlicedDimensions(ArrayRef<OpFoldResult> sliceInputShape,
                          ArrayRef<Range> sliceParams) {
  assert(sliceParams.size() == sliceInputShape.size() &&
         "only supports non rank-reducing case");
  llvm::SmallBitVector mask(sliceInputShape.size());
  unsigned idx = 0;
  for (const auto &[offset, size, stride] : sliceParams) {
    std::optional<int64_t> offsetConst = getConstantIntValue(offset);
    std::optional<int64_t> strideConst = getConstantIntValue(stride);
    mask[idx] = !isEqualConstantIntOrValue(size, sliceInputShape[idx]) ||
                (!strideConst || *strideConst != 1) ||
                (!offsetConst || *offsetConst != 0);
    idx++;
  }
  return mask;
}

llvm::SmallBitVector mlir::getLinearizedDimensions(
    ArrayRef<ReassociationIndices> reassociationIndices) {
  llvm::SmallBitVector result(reassociationIndices.size());
  for (const auto &it : llvm::enumerate(reassociationIndices))
    result[it.index()] = it.value().size() > 1;
  return result;
}

SmallVector<Range> SliceFromCollapseHelper::getExtractSliceParams(
    MLIRContext *ctx, ArrayRef<ValueRange> multiIndices) {
  unsigned loopIdx = 0;
  auto oneAttr = IntegerAttr::get(IndexType::get(ctx), 1);
  auto zeroAttr = IntegerAttr::get(IndexType::get(ctx), 0);
  SmallVector<Range> offsetsSizesAndStrides;
  offsetsSizesAndStrides.reserve(collapseShapeInputShape.size());
  for (const auto &it : llvm::enumerate(reassociationIndices)) {
    // Case 1: Linearized dimensions that have also been sliced. These
    // are size of 1 because we are iterating over these dimensions. The
    // offsets are exactly the de-linearized multi-indices.
    if (slicedDimensions[it.index()] && linearizedDimensions[it.index()]) {
      llvm::append_range(
          offsetsSizesAndStrides,
          llvm::map_range(multiIndices[loopIdx++], [&](Value v) -> Range {
            return Range{getAsOpFoldResult(v), oneAttr, oneAttr};
          }));
      continue;
    }

    // Case 2: One or possibly multiple combined input dimensions, but we
    // have proven that these are not sliced. In this case we just take
    // the full extent of each dimension in the reassociation list.
    if (linearizedDimensions[it.index()]) {
      llvm::append_range(offsetsSizesAndStrides,
                         llvm::map_range(it.value(), [&](int64_t idx) -> Range {
                           return {zeroAttr, collapseShapeInputShape[idx],
                                   oneAttr};
                         }));
      continue;
    }

    // Case 3: A single index, but it may be sliced.
    offsetsSizesAndStrides.push_back(sliceParams[it.index()]);
  }
  return offsetsSizesAndStrides;
}

SmallVector<Range>
SliceFromCollapseHelper::getInsertSliceParams(MLIRContext *ctx,
                                              ValueRange tileIndices) {
  auto one = IntegerAttr::get(IndexType::get(ctx), 1);
  auto zero = IntegerAttr::get(IndexType::get(ctx), 0);
  SmallVector<Range> insertParams;
  insertParams.reserve(linearizedDimensions.size());
  unsigned loopIdx = 0;
  for (unsigned i = 0; i < linearizedDimensions.size(); i++) {
    if (linearizedDimensions[i] && slicedDimensions[i]) {
      insertParams.push_back(Range{tileIndices[loopIdx++], one, one});
      continue;
    }
    insertParams.push_back(Range{zero, sliceParams[i].size, one});
  }
  return insertParams;
}

/// Returns the index of the only non-unit dimension among `indices` of `shape`,
/// if such a dimension exists and `indices` has more than one element.
/// Otherwise, return std::nullopt.
static std::optional<int64_t> getUniqueNonUnitDim(ArrayRef<int64_t> indices,
                                                  ArrayRef<int64_t> shape) {
  // Return false if more than one of the dimensions in this group are not 1.
  std::optional<int64_t> dimIndex;
  if (indices.size() < 2)
    return std::nullopt;
  for (int64_t idx : indices) {
    if (shape[idx] != 1) {
      if (dimIndex != std::nullopt)
        return std::nullopt;
      dimIndex = idx;
    }
  }
  return dimIndex;
}

// For each segment in the reassociation indices, check whether we can
// simplify that segment with a rank-reducing extract slice. We can do this if
// all but (exactly) one of the corresponding source dims is 1.
static SmallVector<std::optional<int64_t>> getCollapseShapeTrivialSegments(
    RankedTensorType sourceType,
    ArrayRef<ReassociationIndices> reassociationIndices) {
  SmallVector<std::optional<int64_t>> trivialSegments;
  for (const auto &indices : reassociationIndices)
    trivialSegments.push_back(
        getUniqueNonUnitDim(indices, sourceType.getShape()));
  return trivialSegments;
}

/// Returns true if any of the segments of the reassociation indices for a
/// collapsing reshape can be simplified using a rank-reducing slice.
static FailureOr<SmallVector<std::optional<int64_t>>>
canCollapseShapeBeSimplifiedByRankReducingSlice(
    RankedTensorType sourceType,
    ArrayRef<ReassociationIndices> reassociationIndices) {
  SmallVector<std::optional<int64_t>> trivialSegments =
      getCollapseShapeTrivialSegments(sourceType, reassociationIndices);
  if (!llvm::any_of(trivialSegments, [](const std::optional<int64_t> &idx) {
        return idx.has_value();
      }))
    return failure();
  return trivialSegments;
}

FailureOr<CollapseShapeRankReducingSliceSimplificationInfo>
mlir::getSimplifyCollapseShapeWithRankReducingSliceInfo(
    RankedTensorType sourceType,
    ArrayRef<ReassociationIndices> reassociationIndices) {
  FailureOr<SmallVector<std::optional<int64_t>>> trivialSegments =
      canCollapseShapeBeSimplifiedByRankReducingSlice(sourceType,
                                                      reassociationIndices);
  if (failed(trivialSegments))
    return failure();

  // Create the expected result shape of the rank-reducing slice.
  SmallVector<int64_t> sliceShape;
  for (const auto &[nonUnitDim, indices] :
       llvm::zip(*trivialSegments, reassociationIndices)) {
    if (nonUnitDim) {
      sliceShape.push_back(sourceType.getDimSize(*nonUnitDim));
      continue;
    }
    llvm::append_range(sliceShape, llvm::map_range(indices, [&](int64_t idx) {
                         return sourceType.getDimSize(idx);
                       }));
  }
  auto sliceType =
      RankedTensorType::get(sliceShape, sourceType.getElementType());

  // If the rank-reducing slice simplified every segment, then we are done.
  if (sliceShape.size() == reassociationIndices.size())
    return CollapseShapeRankReducingSliceSimplificationInfo{sliceType,
                                                            std::nullopt};

  // Otherwise, we need to create a new collapse_shape op for the segments that
  // weren't covered by the slice. By design, the new reassociation indices has
  // the same number of groups as the old reassociation indices.
  SmallVector<ReassociationIndices> newReassociationIndices;
  SmallVector<int64_t, 2> reassociation;
  int64_t groupIdx = 0;
  for (int64_t dimIdx = 0; dimIdx < sliceType.getRank(); dimIdx++) {
    reassociation.push_back(dimIdx);
    if ((*trivialSegments)[groupIdx] ||
        reassociation.size() == reassociationIndices[groupIdx].size()) {
      newReassociationIndices.push_back(reassociation);
      reassociation.clear();
      groupIdx++;
    }
  }

  return CollapseShapeRankReducingSliceSimplificationInfo{
      sliceType, newReassociationIndices};
}

PackingMetadata mlir::computePackingMetadata(int64_t packedRank,
                                             ArrayRef<int64_t> innerDimPos) {
  PackingMetadata res;
  res.insertPositions.reserve(innerDimPos.size());
  // The pack insert position is the position + the number of previously
  // inserted positions + offset.
  // The offset controls whether the packing dimension is the first or last.
  //
  // Example
  // =======
  // Consider packing from a hypothetical ABCD layout to ABCDba whose
  // pack.inner_dims is [1, 0]. The first step consists in undoing the
  // permutation and producing AaBbCD. This is achieved purely by computing the
  // insert positions of `b` and `a` into `ABCD`, starting from [1, 0]. One
  // possibility, is to produce insert positions [2, 0], this would result in an
  // aAbBCD layout (i.e. offset 0). The other possibility, is to produce insert
  // positions [3, 1], this would result in an AaBbCD layout (i.e. offset 1).
  // The latter is what we expect from packing.
  int64_t offset = 1;
  for (int64_t pos : innerDimPos) {
    int64_t numInsertedBefore = llvm::count_if(
        innerDimPos, [&pos](int64_t pos2) { return pos > pos2; });
    res.insertPositions.push_back(pos + numInsertedBefore + offset);
  }

  DenseSet<int64_t> posSet(res.insertPositions.begin(),
                           res.insertPositions.end());
  res.reassociations.reserve(packedRank);
  for (int64_t i = 1; i <= packedRank; ++i) {
    res.outerPositions.push_back(i - 1);
    if (!posSet.contains(i)) {
      res.reassociations.push_back(ReassociationIndices{i - 1});
      continue;
    }
    res.reassociations.push_back(ReassociationIndices{i - 1, i});
    ++i;
  }
  return res;
}

OpFoldResult mlir::reshapeConstantSource(DenseElementsAttr source,
                                         TensorType result,
                                         std::optional<Attribute> cst) {
  if (source && source.isSplat() && result.hasStaticShape() &&
      (!cst.has_value() || source.getSplatValue<Attribute>() == cst.value()))
    return source.resizeSplat(result);

  return {};
}