TilingInterfaceImpl.cpp

Go to the documentation of this file.

//===- TilingInterfaceImpl.cpp - Implementation of TilingInterface -------===//
//
// 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/Linalg/Transforms/TilingInterfaceImpl.h"
 
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/Support/Debug.h"
#include <optional>
 
#define DEBUG_TYPE "linalg-tiling-interface-impl"
 
using namespace mlir;
using namespace mlir::linalg;
 
//===----------------------------------------------------------------------===//
// Utility methods for implementation of Tiling Interface for Linalg ops
//===----------------------------------------------------------------------===//
 
/// Return the SSA values that represent the data point accessed using a given
/// `indexingMap` for a given point in the iteration space represented by `ivs`.
static SmallVector<Value> getIndicesForAccess(OpBuilder &b, Location loc,
                                              AffineMap indexingMap,
                                              ValueRange ivs) {
  SmallVector<Value> indices;
  indices.reserve(indexingMap.getNumResults());
  for (auto result : indexingMap.getResults()) {
    AffineMap m = AffineMap::get(indexingMap.getNumDims(),
                                 indexingMap.getNumSymbols(), result);
    Value v = affine::AffineApplyOp::create(b, loc, m, ivs);
    indices.push_back(v);
  }
  return indices;
}
 
/// Method to inline the payload of a `linalgOp` given the iteration space
/// point and values for the arguments of the payload.
static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp,
                                   ValueRange ivs, ValueRange argValues) {
  Block *body = linalgOp.getBlock();
  IRMapping map;
  map.map(body->getArguments(), argValues);
  for (auto &op : body->without_terminator()) {
    if (auto indexOp = dyn_cast<IndexOp>(&op)) {
      map.map(indexOp.getResult(), ivs[indexOp.getDim()]);
      continue;
    }
    b.clone(op, map);
  }
 
  Operation *terminator = body->getTerminator();
  Location loc = terminator->getLoc();
  for (const auto &operand : llvm::enumerate(terminator->getOperands())) {
    Value toStore = map.lookupOrDefault(operand.value());
    OpOperand *storeInto = linalgOp.getDpsInitOperand(operand.index());
    auto indices = getIndicesForAccess(
        b, loc, linalgOp.getMatchingIndexingMap(storeInto), ivs);
    memref::StoreOp::create(b, loc, toStore,
                            linalgOp.getDpsInitOperand(operand.index())->get(),
                            indices);
  }
  return success();
}
 
/// Verify that tiling can be applied in presence of semi-affine maps.
static LogicalResult
validateTilingSemiAffineMaps(LinalgOp linalgOp, ArrayRef<OpFoldResult> sizes) {
  // Precompute each dimension's constant tile-size upper bound once.
  // A failed entry marks a dynamic tile with no static bound.
  SmallVector<FailureOr<int64_t>> tileSizeBounds =
      llvm::map_to_vector(sizes, [](OpFoldResult size) {
        return ValueBoundsConstraintSet::computeConstantBound(
            presburger::BoundType::UB, size,
            /*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
      });
  SmallVector<int64_t> loopRanges = linalgOp.getStaticLoopRanges();
 
  // Dynamic tiles or dynamic loop ranges are conservatively treated as tiled.
  SmallVector<bool> tiledDims(loopRanges.size(), false);
  for (auto [pos, tileSize] : llvm::enumerate(tileSizeBounds)) {
    if (failed(tileSize)) {
      tiledDims[pos] = true;
      continue;
    }
    if (*tileSize == 0)
      continue;
    tiledDims[pos] =
        ShapedType::isDynamic(loopRanges[pos]) || *tileSize < loopRanges[pos];
  }
 
  for (AffineMap map : linalgOp.getIndexingMapsArray()) {
    for (AffineExpr result : map.getResults()) {
      WalkResult status = result.walk([&](AffineExpr expr) -> WalkResult {
        auto binExpr = dyn_cast<AffineBinaryOpExpr>(expr);
        if (!binExpr)
          return WalkResult::advance();
        AffineExprKind kind = binExpr.getKind();
        if (kind != AffineExprKind::Mod && kind != AffineExprKind::FloorDiv &&
            kind != AffineExprKind::CeilDiv)
          return WalkResult::advance();
 
        // Skip if the semi-affine expression does not involve any tiled
        // dimension: an untiled dimension keeps its full extent in every tile,
        // so re-applying the map on the slice is exact.
        bool involvesTiledDim = expr.walk([&](AffineExpr e) -> WalkResult {
                                      auto dim = dyn_cast<AffineDimExpr>(e);
                                      if (dim && tiledDims[dim.getPosition()])
                                        return WalkResult::interrupt();
                                      return WalkResult::advance();
                                    })
                                    .wasInterrupted();
        if (!involvesTiledDim)
          return WalkResult::advance();
 
        // Allow only `d OP C` map where `d` is a dimension and `C` is a
        // constant. A compound LHS (e.g. `(d0 + d1)`, `(d0 * 2)`, a nested
        // semi-affine expression) or a non-constant step is not provably safe,
        // so reject it.
        auto dimExpr = dyn_cast<AffineDimExpr>(binExpr.getLHS());
        auto stepExpr = dyn_cast<AffineConstantExpr>(binExpr.getRHS());
        if (!dimExpr || !stepExpr || stepExpr.getValue() <= 0) {
          linalgOp.emitOpError()
              << "tiling is not supported for the semi-affine indexing map: "
                 "only a single iteration dimension divided by a positive "
                 "constant step can be tiled over a tiled dimension";
          return WalkResult::interrupt();
        }
 
        // Tiles are spaced by the full tile size, so tile origins are its
        // multiples (0, tileSize, 2*tileSize, ...).
        // A tile's indices are `origin + d'`, with `origin` the tile's start
        // and `0 <= d' < tileSize`. A trailing partial tile is a full tile
        // truncated at the same origin, spanning a subset of the same `d'`, so
        // full-tile validity implies partial-tile validity and validating the
        // upper-bound tile size suffices.
        unsigned dimPos = dimExpr.getPosition();
        FailureOr<int64_t> tileSize = tileSizeBounds[dimPos];
 
        // Dynamic tile sizes are assumed to be valid.
        // Unit tile is always valid.
        if (failed(tileSize) || *tileSize == 1)
          return WalkResult::advance();
 
        // Tiled op reuses the same map on a slice whose base offset is
        // `m(origin) - m(0)`, so it is correct only when
        // `m(origin + d') == (m(origin) - m(0)) + m(d')` for every `d'`.
        // Slice origins are tile-size multiples, so this reduces to a relation
        // between the tile size and the step `C`:
        //  - `floordiv`/`mod` are locally affine within a step window (floordiv
        //    is constant, mod is linear), so they compose when the origin is
        //    step-aligned (`C | tileSize`) or the whole tile fits in one window
        //    (`tileSize | C`);
        //  - `ceildiv` jumps at `k * C + 1` instead of `k * C`, so a
        //    non-step-aligned origin already straddles the jump. It composes
        //    only from a step-aligned origin, i.e. `C | tileSize`.
        int64_t step = stepExpr.getValue();
        bool isCeil = kind == AffineExprKind::CeilDiv;
        bool safe = *tileSize % step == 0 || (!isCeil && step % *tileSize == 0);
        if (!safe) {
          linalgOp.emitOpError()
              << "tiling is not supported for the semi-affine indexing map: "
                 "tile size "
              << *tileSize << " for dimension d" << dimPos
              << (isCeil ? " must be a multiple of the step "
                         : " must divide or be divisible by the step ")
              << step;
          return WalkResult::interrupt();
        }
        return WalkResult::advance();
      });
      if (status.wasInterrupted())
        return failure();
    }
  }
  return success();
}
 
//===----------------------------------------------------------------------===//
// External Model for implementing `TilingInterface` for `LinalgOp`s.
//===----------------------------------------------------------------------===//
 
namespace {
/// External model implementation of TilingInterface for LinalgOps. An external
/// model implementation is used for now till the use of `TilingInterface` is
/// on-par with the current Linalg tiling + fusion patterns. Once it is
/// maybe possible to move this into the op-definition (though there are
/// advantages to leaving it as an external model)
template <typename LinalgOpTy>
struct LinalgOpTilingInterface
    : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
                                            LinalgOpTy> {
  using Base =
      TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
                                     LinalgOpTy>;
  // Inherit the defaulted hint-bearing overloads; these ops do not require the
  // hint (no inner tiles).
  using Base::generateResultTileValue;
  using Base::getIterationDomainTileFromOperandTiles;
  using Base::getTiledImplementation;
  using Base::getTiledImplementationFromOperandTiles;
 
  /// Return the loop iterator type.
  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
    LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
    return concreteOp.getIteratorTypesArray();
  }
 
  /// Return the iteration domain range.
  SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
    OpBuilder::InsertionGuard g(b);
    b.setInsertionPoint(op);
    Location loc = op->getLoc();
    LinalgOp linalgOp = cast<LinalgOp>(op);
    SmallVector<OpFoldResult> allShapesSizes =
        linalgOp.createFlatListOfOperandDims(b, loc);
    AffineMap map = linalgOp.getShapesToLoopsMap();
 
    return llvm::map_to_vector(map.getResults(), [&](AffineExpr loopExpr) {
      OpFoldResult ofr = affine::makeComposedFoldedAffineApply(b, loc, loopExpr,
                                                               allShapesSizes);
      return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
    });
  }
 
  /// Instantiate the tiled implementation of the operation.
  FailureOr<TilingResult>
  getTiledImplementation(Operation *op, OpBuilder &b,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes) const {
    // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
    // specified could lead to out of bounds accesses.
    Location loc = op->getLoc();
    LinalgOp linalgOp = cast<LinalgOp>(op);
    // In case of a semi-affine expression, generalized tracking of tiles would
    // require a per-tile-position shift that cannot be expressed by the
    // symbol-free indexing maps.
    // Thus, tiling is allowed only when the semi-affine maps can be proven safe
    // for the current tiling configuration. Otherwise, tiling can end up
    // producing incorrect results.
    if (failed(validateTilingSemiAffineMaps(linalgOp, sizes)))
      return failure();
    SmallVector<Value> valuesToTile = linalgOp->getOperands();
    SmallVector<Value> tiledOperands = makeTiledShapes(
        b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
    SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
        llvm::make_filter_range(
            tiledOperands,
            [](Value v) -> bool {
              return isa_and_nonnull<tensor::ExtractSliceOp, memref::SubViewOp>(
                  v.getDefiningOp());
            }),
        [](Value v) -> Operation * { return v.getDefiningOp(); });
 
    SmallVector<Type> resultTensorTypes =
        getTensorOutputTypes(linalgOp, tiledOperands);
 
    Operation *tiledOp = clone(b, linalgOp, resultTensorTypes, tiledOperands);
    offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
 
    return TilingResult{
        {tiledOp}, SmallVector<Value>(tiledOp->getResults()), generatedSlices};
  }
 
  /// Utility to fetch the offsets and sizes when applied as per the indexing
  /// map of the linalg op. This helps in fusing the linalg op as a consumer of
  /// a given slice op.
  static LogicalResult
  getMappedOffsetAndSize(LinalgOp linalgOp, OpBuilder &b,
                         ArrayRef<AffineMap> indexingMaps,
                         ArrayRef<SmallVector<OpFoldResult>> allOffsets,
                         ArrayRef<SmallVector<OpFoldResult>> allSizes,
                         SmallVectorImpl<OpFoldResult> &mappedOffsetsVec,
                         SmallVectorImpl<OpFoldResult> &mappedSizesVec) {
    DenseMap<unsigned, OpFoldResult> mappedOffsets, mappedSizes;
 
    for (auto [indexingMap, offsets, sizes] :
         llvm::zip_equal(indexingMaps, allOffsets, allSizes)) {
      for (auto [resultExpr, offset, size] :
           llvm::zip_equal(indexingMap.getResults(), offsets, sizes)) {
        auto dimExpr = dyn_cast<AffineDimExpr>(resultExpr);
        if (!dimExpr)
          return failure();
        unsigned position = dimExpr.getPosition();
        auto it = mappedOffsets.find(position);
        if (it != mappedOffsets.end()) {
          OpFoldResult seenOffset = it->second;
          OpFoldResult seenSize = mappedSizes.lookup(position);
          if (seenOffset != offset || seenSize != size) {
            LLVM_DEBUG({
              llvm::dbgs() << "inconsistent iteration space mapping from "
                              "offsets/sizes of operands/results";
            });
            return failure();
          }
        } else {
          mappedOffsets[position] = offset;
          mappedSizes[position] = size;
        }
      }
    }
 
    // Aggregate from the given operand offsets and sizes, or default to
    // iteration space values.
    SmallVector<Range> iterationDomain =
        cast<TilingInterface>(linalgOp.getOperation()).getIterationDomain(b);
    mappedOffsetsVec.resize(iterationDomain.size());
    mappedSizesVec.resize(iterationDomain.size());
    for (auto [index, domain] : llvm::enumerate(iterationDomain)) {
      auto it = mappedOffsets.find(index);
      if (it != mappedOffsets.end()) {
        mappedOffsetsVec[index] = it->second;
        mappedSizesVec[index] = mappedSizes.lookup(index);
        continue;
      }
      mappedOffsetsVec[index] = domain.offset;
      mappedSizesVec[index] = domain.size;
    }
    return success();
  }
 
  /// Method to return the position of the result tile computed by the tiled
  /// operation.
  LogicalResult getIterationDomainTileFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes,
      SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
      SmallVectorImpl<OpFoldResult> &iterDomainSizes) const {
    auto linalgOp = cast<LinalgOp>(op);
 
    SmallVector<AffineMap> indexingMaps =
        llvm::map_to_vector(operandNumbers, [&](unsigned operandNumber) {
          OpOperand &opOperand = linalgOp->getOpOperand(operandNumber);
          return linalgOp.getMatchingIndexingMap(&opOperand);
        });
    if (failed(getMappedOffsetAndSize(linalgOp, b, indexingMaps, allOffsets,
                                      allSizes, iterDomainOffsets,
                                      iterDomainSizes))) {
      return failure();
    }
    return success();
  }
 
  /// Return the details of the output tile generated by the tiled
  /// implementation.
  LogicalResult
  getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                        ArrayRef<OpFoldResult> offsets,
                        ArrayRef<OpFoldResult> sizes,
                        SmallVector<OpFoldResult> &resultOffsets,
                        SmallVector<OpFoldResult> &resultSizes) const {
    Location loc = op->getLoc();
    LinalgOp linalgOp = cast<LinalgOp>(op);
 
    AffineExpr d0;
    bindDims(b.getContext(), d0);
    SmallVector<OpFoldResult> subShapeSizes =
        llvm::map_to_vector(sizes, [&](OpFoldResult ofr) {
          return affine::makeComposedFoldedAffineApply(b, loc, d0 - 1, ofr);
        });
 
    OpOperand *outOperand = linalgOp.getDpsInitOperand(resultNumber);
    SliceParameters sliceParams = computeSliceParameters(
        b, loc, outOperand->get(), sizes,
        linalgOp.getMatchingIndexingMap(outOperand), offsets,
        /*ubs*/ {}, subShapeSizes, true);
    resultOffsets = sliceParams.offsets;
    resultSizes = sliceParams.sizes;
    return success();
  }
 
  LogicalResult getIterationDomainTileFromResultTile(
      Operation *op, OpBuilder &b, unsigned resultNumber,
      ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
      SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
      SmallVectorImpl<OpFoldResult> &iterDomainSizes) const {
    auto linalgOp = cast<LinalgOp>(op);
 
    // Check that the indexing map used for the output is a projected
    // permutation. This could be relaxed with a more general approach that can
    // map the offsets and sizes from the result to iteration space tiles
    // (filling in full extent for dimensions not used to access the result).
    AffineMap indexingMap =
        linalgOp.getIndexingMapMatchingResult(op->getResult(resultNumber));
    if (!indexingMap.isProjectedPermutation()) {
      return op->emitOpError(
          "unhandled tiled implementation generation when result is not "
          "accessed using a permuted projection");
    }
 
    SmallVector<OpFoldResult> allOffsets = llvm::to_vector(offsets);
    SmallVector<OpFoldResult> allSizes = llvm::to_vector(sizes);
    auto status =
        getMappedOffsetAndSize(linalgOp, b, indexingMap, {allOffsets},
                               {allSizes}, iterDomainOffsets, iterDomainSizes);
    (void)status;
    assert(succeeded(status) && "unexpected error in offset calculation");
    return success();
  }
 
  FailureOr<TilingResult>
  generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
    SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
    if (failed(getIterationDomainTileFromResultTile(
            op, b, resultNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
      return failure();
    }
    auto tilingInterfaceOp = cast<TilingInterface>(op);
    FailureOr<TilingResult> tilingResult =
        tilingInterfaceOp.getTiledImplementation(b, mappedOffsets, mappedSizes);
 
    if (failed(tilingResult))
      return failure();
 
    if (tilingResult->tiledOps.size() != 1)
      return op->emitOpError("failed to generate tiled implementation");
 
    return TilingResult{
        tilingResult->tiledOps,
        SmallVector<Value>{tilingResult->tiledValues[resultNumber]},
        tilingResult->generatedSlices};
  }
 
  /// Method to generate the tiled implementation of an operation from the tile
  /// of the operand.
  FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
    SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
    if (failed(getIterationDomainTileFromOperandTiles(
            op, b, operandNumbers, allOffsets, allSizes, mappedOffsets,
            mappedSizes))) {
      return failure();
    }
    return getTiledImplementation(op, b, mappedOffsets, mappedSizes);
  }
 
  LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
                                             Location loc,
                                             ValueRange ivs) const {
    auto linalgOp = cast<LinalgOp>(op);
    if (!linalgOp.hasPureBufferSemantics())
      return op->emitOpError("expected operation to have buffer semantics");
 
    SmallVector<Value> indexedValues;
    indexedValues.reserve(linalgOp->getNumOperands());
    Location linalgOpLoc = op->getLoc();
    /// Load the data corresponding to the block arguments that
    /// represent input operands.
    for (OpOperand &operand : linalgOp->getOpOperands()) {
      if (!linalgOp.payloadUsesValueFromOperand(&operand)) {
        indexedValues.push_back(nullptr);
        continue;
      }
      if (linalgOp.isScalar(&operand)) {
        indexedValues.push_back(operand.get());
        continue;
      }
      SmallVector<Value> indices = getIndicesForAccess(
          builder, linalgOpLoc, linalgOp.getMatchingIndexingMap(&operand), ivs);
      Value load =
          memref::LoadOp::create(builder, linalgOpLoc, operand.get(), indices);
      indexedValues.push_back(load);
    }
 
    /// Inline the op payload and store the result.
    return inlinePayload(builder, linalgOp, ivs, indexedValues);
  }
 
  bool isOpFusableWithConsumerSlice(Operation *op, unsigned resultNumber,
                                    ArrayRef<OpFoldResult> offsets,
                                    ArrayRef<OpFoldResult> sizes) const {
    // The verifier gives all the necessary requirements for consumer fusion.
    return true;
  }
 
  bool isOpFusableWithProducerSlices(
      Operation *op, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
 
    auto linalgOp = cast<LinalgOp>(op);
    SmallVector<AffineMap> indexingMaps =
        llvm::map_to_vector(operandNumbers, [&](unsigned operandNumber) {
          OpOperand &opOperand = linalgOp->getOpOperand(operandNumber);
          return linalgOp.getMatchingIndexingMap(&opOperand);
        });
    // Check that offsets/sizes are consistent across all operands.
    OpBuilder b(op);
    SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
    return succeeded(getMappedOffsetAndSize(linalgOp, b, indexingMaps,
                                            allOffsets, allSizes, mappedOffsets,
                                            mappedSizes));
  }
};
 
//===----------------------------------------------------------------------===//
// External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
//===----------------------------------------------------------------------===//
 
/// In a given set vector, get the position of a particular element.
std::optional<int> getPositionIn(const llvm::SetVector<unsigned> &reductionDims,
                                 unsigned value) {
  for (auto [index, reductionDim] : llvm::enumerate(reductionDims)) {
    if (reductionDim == value) {
      return index;
    }
  }
  return std::nullopt;
}
 
/// Return an AffineMaps to use for the `outs` operands of the linalg op
/// generated for partial results. The new AffineMap is the AffineMap of the
/// untiled op with reduction dimensions appended at end in order in which they
/// were specified during tiling.
static SmallVector<AffineMap>
getPartialResultAffineMaps(LinalgOp linalgOp,
                           const SetVector<unsigned> &reductionDims) {
  auto partialReductionMaps = llvm::map_to_vector(
      linalgOp.getDpsInitsMutable(), [&](OpOperand &opOperand) {
        AffineMap map = linalgOp.getMatchingIndexingMap(&opOperand);
        for (auto redPos : reductionDims) {
          map =
              map.insertResult(getAffineDimExpr(redPos, linalgOp.getContext()),
                               map.getNumResults());
        }
        return map;
      });
  return partialReductionMaps;
}
 
struct InitSliceInfo {
  SmallVector<int64_t> resultShape;
  SmallVector<OpFoldResult> offsets;
  SmallVector<OpFoldResult> sizes;
  SmallVector<OpFoldResult> strides;
};
 
/// Return the result shape, offsets, sizes and strides of the slice of the
/// `initValue` to use as the destination of the partial reduction op generated
/// with outer reduction strategy.
static InitSliceInfo getInitSliceInfoForOuterReduction(
    MLIRContext *context, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, const SetVector<unsigned> &reductionDims,
    ArrayRef<OpFoldResult> splitReductionIvs, AffineMap partialReductionMap,
    ArrayRef<OpFoldResult> initOperandShape) {
  int64_t initRank = partialReductionMap.getNumResults();
  SmallVector<OpFoldResult> initOffsets, initSizes;
  Attribute zero = IntegerAttr::get(IndexType::get(context), 0);
  Attribute one = IntegerAttr::get(IndexType::get(context), 1);
  SmallVector<OpFoldResult> initStrides(initRank, one);
  for (auto [resultIdx, dimExpr] :
       llvm::enumerate(partialReductionMap.getResults())) {
    if (isa<AffineConstantExpr>(dimExpr)) {
      // A constant index in the output map accesses a fixed position; keep
      // the full output dimension to match the original output operand shape.
      initOffsets.push_back(zero);
      initSizes.push_back(initOperandShape[resultIdx]);
      continue;
    }
    unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
    if (reductionDims.contains(dim)) {
      initOffsets.push_back(zero);
    } else {
      initOffsets.push_back(offsets[dim]);
    }
    initSizes.push_back(sizes[dim]);
  }
  SmallVector<int64_t> resultShape;
  std::tie(resultShape, std::ignore) = decomposeMixedValues(initSizes);
  return {resultShape, initOffsets, initSizes, initStrides};
}
 
/// Return the result shape, offsets, sizes and strides of the slice of the
/// `initValue` to use as destination of the partial reduction op generated with
/// outer parallel strategy.
static InitSliceInfo getInitSliceInfoForOuterParallel(
    MLIRContext *context, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, const SetVector<unsigned> &reductionDims,
    ArrayRef<OpFoldResult> splitReductionIvs, AffineMap partialReductionMap,
    ArrayRef<OpFoldResult> initOperandShape) {
  int64_t initRank = partialReductionMap.getNumResults();
  SmallVector<OpFoldResult> initOffsets, initSizes;
  Attribute zero = IntegerAttr::get(IndexType::get(context), 0);
  Attribute one = IntegerAttr::get(IndexType::get(context), 1);
  SmallVector<OpFoldResult> initStrides(initRank, one);
  SmallVector<OpFoldResult> resultShape;
  for (auto [resultIdx, dimExpr] :
       llvm::enumerate(partialReductionMap.getResults())) {
    if (isa<AffineConstantExpr>(dimExpr)) {
      // A constant index accesses a fixed position; keep the full output
      // dimension to match the original output operand shape.
      initOffsets.push_back(zero);
      initSizes.push_back(initOperandShape[resultIdx]);
      resultShape.push_back(initOperandShape[resultIdx]);
      continue;
    }
    unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
    if (std::optional<unsigned> dimPos = getPositionIn(reductionDims, dim)) {
      initOffsets.push_back(splitReductionIvs[dimPos.value()]);
      initSizes.push_back(one);
    } else {
      initOffsets.push_back(offsets[dim]);
      initSizes.push_back(sizes[dim]);
      resultShape.push_back(sizes[dim]);
    }
  }
  SmallVector<int64_t> staticShapes;
  std::tie(staticShapes, std::ignore) = decomposeMixedValues(resultShape);
  return {staticShapes, initOffsets, initSizes, initStrides};
}
 
/// Return the result shape, offsets, sizes and strides of the slice of the
/// `initValue` to use as destination of the partial reduction op.
static InitSliceInfo getInitSliceInfo(MLIRContext *context,
                                      ReductionTilingStrategy strategy,
                                      ArrayRef<OpFoldResult> offsets,
                                      ArrayRef<OpFoldResult> sizes,
                                      const SetVector<unsigned> &reductionDims,
                                      ArrayRef<OpFoldResult> splitReductionIvs,
                                      AffineMap partialReductionMap,
                                      ArrayRef<OpFoldResult> initOperandShape) {
  if (strategy == ReductionTilingStrategy::PartialReductionOuterReduction) {
    return getInitSliceInfoForOuterReduction(
        context, offsets, sizes, reductionDims, splitReductionIvs,
        partialReductionMap, initOperandShape);
  }
  assert(strategy == ReductionTilingStrategy::PartialReductionOuterParallel &&
         "unexpected ReductionTilingStrategy");
  return getInitSliceInfoForOuterParallel(
      context, offsets, sizes, reductionDims, splitReductionIvs,
      partialReductionMap, initOperandShape);
}
 
/// External model implementation of PartialReductionInterface for
/// LinalgOps.
template <typename LinalgOpTy>
struct LinalgOpPartialReductionInterface
    : public PartialReductionOpInterface::ExternalModel<
          LinalgOpPartialReductionInterface<LinalgOpTy>, LinalgOpTy> {
  FailureOr<SmallVector<Value>> generateInitialTensorForPartialReduction(
      Operation *op, OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
      const SetVector<unsigned> &reductionDims) const {
    auto linalgOp = cast<LinalgOp>(op);
 
    OpBuilder::InsertionGuard guard(b);
    if (linalgOp.hasPureBufferSemantics())
      return op->emitOpError("expected operation to have tensor semantics");
 
    SmallVector<AffineMap> partialResultMaps =
        getPartialResultAffineMaps(linalgOp, reductionDims);
 
    SmallVector<Value> inits;
    for (auto [initIdx, result, partialMap] :
         llvm::enumerate(linalgOp->getResults(), partialResultMaps)) {
      SmallVector<Operation *, 4> combinerOps;
      if (!matchReduction(linalgOp.getRegionOutputArgs(), initIdx,
                          combinerOps) ||
          combinerOps.size() != 1)
        return op->emitOpError("Failed to anaysis the reduction operation.");
 
      Operation *reductionOp = combinerOps[0];
      std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
      if (!identity.has_value())
        return op->emitOpError(
            "Failed to get an identity value for the reduction operation.");
 
      // Append the new partial result dimensions.
      SmallVector<OpFoldResult> partialResultShape;
      Value initValue = linalgOp.getDpsInits()[initIdx];
      SmallVector<OpFoldResult> initShape =
          tensor::getMixedSizes(b, loc, initValue);
      for (auto [resultIdx, dimExpr] :
           llvm::enumerate(partialMap.getResults())) {
        if (isa<AffineConstantExpr>(dimExpr)) {
          // A constant index in the output map accesses a fixed position; use
          // the actual output dimension size (not a hardcoded 1).
          partialResultShape.push_back(initShape[resultIdx]);
          continue;
        }
        auto dim = cast<AffineDimExpr>(dimExpr);
        partialResultShape.push_back(sizes[dim.getPosition()]);
      }
 
      Type elType = getElementTypeOrSelf(result.getType());
      Value emptyTensor =
          tensor::EmptyOp::create(b, loc, partialResultShape, elType);
      Value constantOp = arith::ConstantOp::create(b, loc, *identity);
      auto identityTensor =
          linalg::FillOp::create(b, loc, constantOp, emptyTensor);
      inits.push_back(identityTensor.getResult(0));
    }
 
    return inits;
  }
 
  FailureOr<TilingResult>
  tileToPartialReduction(Operation *op, OpBuilder &b, Location loc,
                         ReductionTilingStrategy tilingStrategy,
                         ValueRange init, ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes,
                         const SetVector<unsigned> &reductionDims,
                         ArrayRef<OpFoldResult> splitReductionIvs) const {
    OpBuilder::InsertionGuard guard(b);
    auto linalgOp = cast<LinalgOp>(op);
 
    SmallVector<AffineMap> partialReductionMaps =
        getPartialResultAffineMaps(linalgOp, reductionDims);
 
    // Step 1. Extend init maps to have reduction dimension dims, since we
    // are converting them to parallel dimensions.
    SmallVector<AffineMap> newInitMaps;
    if (tilingStrategy ==
        ReductionTilingStrategy::PartialReductionOuterReduction) {
      newInitMaps = llvm::to_vector(partialReductionMaps);
    } else {
      newInitMaps = llvm::map_to_vector(
          linalgOp.getDpsInitsMutable(), [&](OpOperand &opOperand) {
            return linalgOp.getMatchingIndexingMap(&opOperand);
          });
    }
 
    // Step 2a: Extract a slice of the input operands.
    SmallVector<Value> tiledInputs = makeTiledShapes(
        b, loc, linalgOp, linalgOp.getDpsInputs(), offsets, sizes, {}, true);
    SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
        llvm::make_filter_range(
            tiledInputs, [](Value v) -> bool { return v.getDefiningOp(); }),
        [](Value v) -> Operation * { return v.getDefiningOp(); });
 
    // Step 2b: Extract a slice of the init operands.
    SmallVector<Value, 1> tiledInits;
    for (auto [partialReductionMap, valueToTile, initOperandValue] :
         llvm::zip_equal(partialReductionMaps, init, linalgOp.getDpsInits())) {
      // Compute the actual shape of the original init operand for handling
      // constant expressions in the partial reduction map.
      SmallVector<OpFoldResult> initOperandShape =
          tensor::getMixedSizes(b, loc, initOperandValue);
      InitSliceInfo sliceInfo = getInitSliceInfo(
          b.getContext(), tilingStrategy, offsets, sizes, reductionDims,
          splitReductionIvs, partialReductionMap, initOperandShape);
      auto valueToTileType = cast<RankedTensorType>(valueToTile.getType());
      RankedTensorType sliceResultType = RankedTensorType::get(
          sliceInfo.resultShape, valueToTileType.getElementType(),
          valueToTileType.getEncoding());
      auto sliceOp = tensor::ExtractSliceOp::create(
          b, loc, sliceResultType, valueToTile, sliceInfo.offsets,
          sliceInfo.sizes, sliceInfo.strides);
      tiledInits.push_back(sliceOp.getResult());
      generatedSlices.push_back(sliceOp);
    }
 
    // Update the indexing maps.
    SmallVector<AffineMap> newMaps = linalgOp.getIndexingMapsArray();
    for (auto [initOperand, newInitMap] :
         llvm::zip_equal(linalgOp.getDpsInitsMutable(), newInitMaps)) {
      int mapIdx = linalgOp.getIndexingMapIndex(&initOperand);
      newMaps[mapIdx] = newInitMap;
    }
 
    // Step 3. Change the reduction dim iterator types.
    SmallVector<utils::IteratorType> newIteratorTypes =
        linalgOp.getIteratorTypesArray();
    if (tilingStrategy ==
        ReductionTilingStrategy::PartialReductionOuterReduction) {
      for (int dim : reductionDims)
        newIteratorTypes[dim] = utils::IteratorType::parallel;
    }
 
    // Step 4. Create the new generic op.
    Operation *partialReductionOp;
    auto resultTypes = ValueRange(tiledInits).getTypes();
    if (tilingStrategy ==
        ReductionTilingStrategy::PartialReductionOuterReduction) {
      auto genericOp = GenericOp::create(b, loc, resultTypes, tiledInputs,
                                         tiledInits, newMaps, newIteratorTypes);
      IRMapping mapping;
      op->getRegion(0).cloneInto(&genericOp.getRegion(),
                                 genericOp.getRegion().begin(), mapping);
      offsetIndices(b, genericOp, offsets);
      partialReductionOp = genericOp.getOperation();
    } else {
      SmallVector<Value> operands = std::move(tiledInputs);
      llvm::append_range(operands, tiledInits);
      partialReductionOp = mlir::clone(b, op, resultTypes, operands);
      offsetIndices(b, cast<LinalgOp>(partialReductionOp), offsets);
    }
    return TilingResult{
        {partialReductionOp},
        llvm::map_to_vector(partialReductionOp->getResults(),
                            [](OpResult r) -> Value { return r; }),
        generatedSlices};
  }
 
  FailureOr<MergeResult>
  mergeReductions(Operation *op, OpBuilder &b, Location loc,
                  ValueRange partialReduce,
                  const SetVector<unsigned> &reductionDims) const {
    auto linalgOp = cast<LinalgOp>(op);
    SmallVector<AffineMap> partialReductionMaps =
        getPartialResultAffineMaps(linalgOp, reductionDims);
 
    // Permute the reduction dims as permuted by the partial result map.
    SmallVector<Operation *> mergeOperations;
    SmallVector<Value> replacements;
    for (auto [idx, init, partialResult, partialMap] : llvm::enumerate(
             linalgOp.getDpsInits(), partialReduce, partialReductionMaps)) {
      unsigned initIdx = idx;
      // linalg.reduce's iteration space is the tiled result's iteration space
      // (and not the tiled operation's iteration space). To account for this,
      // permute the reduction dimensions based on the partial result map of the
      // tiled result.
      SmallVector<int64_t> partialReductionDims;
      for (auto [resultNum, dimExpr] :
           llvm::enumerate(partialMap.getResults())) {
        if (isa<AffineConstantExpr>(dimExpr))
          continue; // Constant dims are never reduction dims.
        unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
        if (llvm::is_contained(reductionDims, dim)) {
          partialReductionDims.push_back(resultNum);
        }
      }
 
      auto reduction = linalg::ReduceOp::create(
          b, loc, partialResult, init, partialReductionDims,
          [&linalgOp, &initIdx](OpBuilder &b, Location loc, ValueRange inputs) {
            // Get the combiner op.
            SmallVector<Operation *, 4> combinerOps;
            matchReduction(linalgOp.getRegionOutputArgs(), initIdx,
                           combinerOps);
            Operation *clonedReductionOp = b.clone(*combinerOps[0]);
            // Combine the input at idx and output at numInits + idx.
            clonedReductionOp->setOperand(0, inputs[0]);
            clonedReductionOp->setOperand(1, inputs[1]);
            linalg::YieldOp::create(b, loc, clonedReductionOp->getResult(0));
          });
 
      mergeOperations.push_back(reduction);
      replacements.push_back(reduction->getResult(0));
    }
 
    return MergeResult{mergeOperations, replacements};
  }
 
  LogicalResult getPartialResultTilePosition(
      Operation *op, OpBuilder &b, unsigned resultNumber,
      ReductionTilingStrategy tilingStrategy, ArrayRef<OpFoldResult> offsets,
      ArrayRef<OpFoldResult> sizes, const SetVector<unsigned> &reductionDims,
      ArrayRef<OpFoldResult> splitReductionIvs,
      SmallVector<OpFoldResult> &resultOffsets,
      SmallVector<OpFoldResult> &resultSizes) const {
    auto linalgOp = cast<LinalgOp>(op);
    SmallVector<AffineMap> partialReductionMaps =
        getPartialResultAffineMaps(linalgOp, reductionDims);
    // Compute the actual shape of the init operand for handling constant
    // expressions in the partial reduction map.
    Value initOperandValue = linalgOp.getDpsInits()[resultNumber];
    Location loc = op->getLoc();
    SmallVector<OpFoldResult> initOperandShape =
        tensor::getMixedSizes(b, loc, initOperandValue);
    InitSliceInfo sliceInfo =
        getInitSliceInfo(b.getContext(), tilingStrategy, offsets, sizes,
                         reductionDims, splitReductionIvs,
                         partialReductionMaps[resultNumber], initOperandShape);
    std::swap(resultOffsets, sliceInfo.offsets);
    std::swap(resultSizes, sliceInfo.sizes);
 
    return success();
  }
};
 
template <typename OpTy>
static SmallVector<Range> getPackUnPackIterationDomain(OpTy op,
                                                       OpBuilder &builder) {
  static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                "applies to only pack or unpack operations");
  OpBuilder::InsertionGuard g(builder);
  int64_t rank = (std::is_same<OpTy, PackOp>::value) ? op.getSourceRank()
                                                     : op.getDestRank();
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  ReifiedRankedShapedTypeDims resultShape;
  (void)op.reifyResultShapes(builder, resultShape);
  SmallVector<Range> loopBounds(rank);
  for (auto dim : llvm::seq<int64_t>(0, rank)) {
    loopBounds[dim].offset = zero;
    loopBounds[dim].stride = one;
    loopBounds[dim].size = resultShape[0][dim];
  }
  return loopBounds;
}
 
static void applyPermToRange(SmallVector<OpFoldResult> &offsets,
                             SmallVector<OpFoldResult> &sizes,
                             ArrayRef<int64_t> permutation) {
  if (permutation.empty())
    return;
  applyPermutationToVector<OpFoldResult>(offsets, permutation);
  applyPermutationToVector<OpFoldResult>(sizes, permutation);
}
 
/// Compute the permutation vector to interchange `elements` such that the
/// elements at positions in `dimsPos` are moved to the positions `[0, ...,
/// dimsPos.size())` in order.
static SmallVector<int64_t>
computeInterchangeFromDimPos(ArrayRef<int64_t> dimsPos, int64_t rank) {
  SmallVector<int64_t> interchangeVector;
  interchangeVector.reserve(dimsPos.size());
  // First map dims and their position. For example, dims_pos = [2, 0] will map
  // to:
  // [
  //  [ key: 2, value: 0]
  //  [ key: 0, value: 1]
  // ]
  // where key is the idx in dims_pos while value its position in dims_pos.
  DenseMap<int64_t, int64_t> dimsAndPosMapping;
  for (int64_t dimsIdx = 0, end = dimsPos.size(); dimsIdx < end; dimsIdx++)
    dimsAndPosMapping[dimsPos[dimsIdx]] = dimsIdx;
 
  // Scan the position in order and insert the value in the map
  // to compute the interchange vector.
  for (int64_t dimsIdx = 0; dimsIdx < rank; dimsIdx++) {
    if (dimsAndPosMapping.count(dimsIdx))
      interchangeVector.push_back(dimsAndPosMapping[dimsIdx]);
  }
  return interchangeVector;
}
 
/// Permute the elements of `vec` starting at position `offset` according to
/// `interchangeVector`. The permutation maps position `i` in the permuted range
/// to position `interchangeVector[i]` in the original range. Elements before
/// `offset` are unchanged.
///
/// Example: interchange([a, b, c, d, e], [2, 0, 1], offset=2)
///          returns [a, b, e, c, d] (permutes the suffix [c, d, e])
///
/// Note: This is similar to `applyPermutationToVector` but supports an offset
/// for permuting a suffix of the vector. It is only used for pack/unpack scalar
/// implementation where we need to permute inner tile dimensions which are
/// stored at the end of the index vector.
template <typename T>
static SmallVector<T> interchange(ArrayRef<T> elements,
                                  ArrayRef<int64_t> interchangeVector,
                                  int offset = 0) {
  SmallVector<T> vec = llvm::to_vector(elements);
  for (auto [idx, val] : llvm::enumerate(interchangeVector))
    vec[idx + offset] = elements[val + offset];
  return vec;
}
 
/// Generate the body of the innermost loop of the scalar implementation
/// of `pack` operation.
static void generatePackOpScalarImplementationBody(PackOp packOp,
                                                   OpBuilder &builder,
                                                   Location loc,
                                                   ValueRange ivs) {
  // Note: `ivs` are already in the correct order, possibly interchanged based
  // on `dims_pos`. However, connecting the loops with the access patterns is
  // difficult - What is the relation between the position of the tile loop and
  // the point loop? However, if we interchange `ivs` once more to go to the
  // canonical blocking format: ABCabc, this connection becomes trivial: Each
  // point loop is pointLoopsOffset + inputRank away from the tiled loop.
  ArrayRef<int64_t> dimsToInnerBlock = packOp.getInnerDimsPos();
  ArrayRef<int64_t> dimsToOuterBlock = packOp.getOuterDimsPerm();
 
  SmallVector<Value> interchangedIvs = ivs;
  SmallVector<int64_t> interchangeVector =
      computeInterchangeFromDimPos(dimsToInnerBlock, packOp.getSourceRank());
  interchangedIvs = interchange<Value>(interchangedIvs, interchangeVector,
                                       /*offset=*/packOp.getSourceRank());
  if (!dimsToOuterBlock.empty()) {
    interchangeVector =
        computeInterchangeFromDimPos(dimsToOuterBlock, packOp.getSourceRank());
    interchangedIvs =
        interchange<Value>(interchangedIvs, interchangeVector, /*offset=*/0);
  }
  DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
      packOp.getDimAndTileMapping();
  SmallVector<OpFoldResult> sourceIndices;
  size_t pointLoopsOffset = 0;
  int64_t sourceRank = packOp.getSourceRank();
  for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
    if (dimAndTileMapping.contains(dim)) {
      AffineExpr i, j, tile;
      bindDims(builder.getContext(), i, j);
      bindSymbols(builder.getContext(), tile);
      OpFoldResult sourceIndex = affine::makeComposedFoldedAffineApply(
          builder, loc, i * tile + j,
          ArrayRef<OpFoldResult>{
              interchangedIvs[dim],
              interchangedIvs[pointLoopsOffset + packOp.getSourceRank()],
              dimAndTileMapping[dim]});
      sourceIndices.push_back(sourceIndex);
      ++pointLoopsOffset;
    } else {
      sourceIndices.push_back(interchangedIvs[dim]);
    }
  }
 
  auto createLoad = [&]() -> Value {
    return memref::LoadOp::create(
        builder, loc, packOp.getSource(),
        getValueOrCreateConstantIndexOp(builder, loc, sourceIndices));
  };
  Value scalar;
  if (auto paddingValue = packOp.getPaddingValue()) {
    ArithBuilder arithBuilder(builder, loc);
    Value isInBounds;
    for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
      Value idx =
          getValueOrCreateConstantIndexOp(builder, loc, sourceIndices[dim]);
      Value cond = arithBuilder.slt(
          idx, createOrFoldDimOp(builder, loc, packOp.getSource(), dim));
      isInBounds = dim == 0 ? cond : arithBuilder._and(isInBounds, cond);
    }
    scalar = scf::IfOp::create(
                 builder, loc, isInBounds, /*thenBuilder=*/
                 [&](OpBuilder &b, Location l) {
                   scf::YieldOp::create(b, l, createLoad());
                 },
                 /*elseBuilder=*/
                 [&](OpBuilder &b, Location l) {
                   scf::YieldOp::create(b, l, paddingValue);
                 })
                 .getResult(0);
  } else {
    scalar = createLoad();
  }
 
  memref::StoreOp::create(builder, loc, scalar, packOp.getDest(), ivs);
}
 
struct PackOpTiling
    : public TilingInterface::ExternalModel<PackOpTiling, linalg::PackOp> {
  using Base = TilingInterface::ExternalModel<PackOpTiling, linalg::PackOp>;
  using Base::getTiledImplementation;
 
  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
    // Note that here we only consider untiled dimensions and outer tiled data
    // dimensions, the inner tiled data dimensions are materialized when
    // building the body of the operation.
    auto packOp = cast<PackOp>(op);
    SmallVector<utils::IteratorType> iteratorTypes(
        packOp.getSourceRank(), utils::IteratorType::parallel);
    return iteratorTypes;
  }
 
  SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
    return getPackUnPackIterationDomain<PackOp>(cast<PackOp>(op), b);
  }
 
  FailureOr<TilingResult>
  getTiledImplementation(Operation *op, OpBuilder &b,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes) const {
    auto packOp = cast<PackOp>(op);
    // TODO: Support Memref PackOp. Temporarily return failure.
    if (!packOp.hasPureTensorSemantics())
      return failure();
 
    Location loc = packOp.getLoc();
 
    // The tiling is applied on interchanged dimensions. We have to undo the
    // interchange to map sizes and offsets to the original input.
    int64_t inputRank = packOp.getSourceRank();
    SmallVector<OpFoldResult> origOffsets(offsets);
    SmallVector<OpFoldResult> origSizes(sizes);
    applyPermToRange(origOffsets, origSizes,
                     invertPermutationVector(packOp.getOuterDimsPerm()));
 
    DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
        packOp.getDimAndTileMapping();
    SmallVector<OpFoldResult> srcDimValues =
        tensor::getMixedSizes(b, loc, packOp.getSource());
    SmallVector<OpFoldResult> inputIndices, inputSizes;
    for (auto dim : llvm::seq<int64_t>(0, inputRank)) {
      using AV = affine::AffineValueExpr;
      affine::AffineBuilder ab(b, loc);
      AffineExpr dim0, dim1, sym;
      bindDims(b.getContext(), dim0, dim1);
      bindSymbols(b.getContext(), sym);
      if (dimAndTileMapping.count(dim)) {
        // If the data dimension is tiled, the i-th index is the product of
        // offset_i and tile_i, and the i-th size is the product of sizes_i and
        // tile_i.
        auto avOffset = AV(dim0).bind(origOffsets[dim]);
        auto avSize = AV(dim0).bind(origSizes[dim]);
        auto avTileSize = AV(sym).bind(dimAndTileMapping[dim]);
        inputIndices.push_back(ab.mul(avOffset, avTileSize));
        inputSizes.push_back(ab.mul(avSize, avTileSize));
      } else {
        inputIndices.push_back(origOffsets[dim]);
        inputSizes.push_back(origSizes[dim]);
      }
 
      // Limit the size of the input operand for incomplete tiles.
      if (packOp.getPaddingValue()) {
        OpFoldResult dimSize = srcDimValues[dim];
        auto avDimSize = AV(dim0).bind(dimSize);
        auto avInputIdx = AV(dim1).bind(inputIndices.back());
        inputSizes.back() =
            ab.min({inputSizes.back(), ab.sub(avDimSize, avInputIdx)});
      }
    }
 
    auto oneAttr = b.getI64IntegerAttr(1);
    SmallVector<OpFoldResult> strides(inputRank, oneAttr);
 
    SmallVector<Value> tiledOperands;
    auto sourceSlice = tensor::ExtractSliceOp::create(
        b, loc, packOp.getSource(), inputIndices, inputSizes, strides);
    tiledOperands.push_back(sourceSlice);
 
    SmallVector<OpFoldResult> outputOffsets, outputSizes;
    if (failed(getResultTilePosition(op, b, 0, offsets, sizes, outputOffsets,
                                     outputSizes)))
      return {};
 
    strides.append(packOp.getDestRank() - inputRank, oneAttr);
    auto outSlice = tensor::ExtractSliceOp::create(
        b, loc, packOp.getDest(), outputOffsets, outputSizes, strides);
    tiledOperands.push_back(outSlice);
 
    if (auto val = packOp.getPaddingValue())
      tiledOperands.push_back(val);
    for (auto tile : packOp.getInnerTiles())
      tiledOperands.push_back(tile);
 
    Operation *tiledPackOp = PackOp::create(
        b, loc, TypeRange{outSlice.getType()}, tiledOperands, op->getAttrs());
 
    return TilingResult{
        {tiledPackOp},
        SmallVector<Value>(tiledPackOp->getResults()),
        llvm::to_vector(ArrayRef<Operation *>{sourceSlice, outSlice})};
  }
 
  LogicalResult
  getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                        ArrayRef<OpFoldResult> offsets,
                        ArrayRef<OpFoldResult> sizes,
                        SmallVector<OpFoldResult> &resultOffsets,
                        SmallVector<OpFoldResult> &resultSizes) const {
    // The iteration domain is over outer dimensions of packed layout. In this
    // context, the outer dimensions of `resultOffsets` are `offsets`. The
    // inner dimensions of `resultOffsets` are zeros because tiling is not
    // applied to them.
    auto packOp = cast<PackOp>(op);
    int64_t inputRank = packOp.getSourceRank();
    int64_t outputRank = packOp.getDestRank();
    auto zeroAttr = b.getI64IntegerAttr(0);
    resultOffsets.assign(offsets.begin(), offsets.end());
    resultOffsets.append(outputRank - inputRank, zeroAttr);
 
    ReifiedRankedShapedTypeDims outputShape;
    (void)reifyResultShapes(b, packOp, outputShape);
    resultSizes.assign(sizes.begin(), sizes.end());
    for (auto dataTileDim : llvm::seq<unsigned>(inputRank, outputRank))
      resultSizes.push_back(outputShape[0][dataTileDim]);
 
    return success();
  }
 
  FailureOr<TilingResult>
  generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
    return generateResultTileValue(op, b, resultNumber, offsets, sizes,
                                   /*innerTileAlignments=*/{});
  }
 
  FailureOr<TilingResult> generateResultTileValue(
      Operation *op, OpBuilder &b, unsigned resultNumber,
      ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
      ArrayRef<InnerTileAlignment> innerTileAlignments) const {
    auto packOp = cast<PackOp>(op);
    int64_t numTiles = packOp.getInnerDimsPos().size();
 
    // linalg.pack op is fusible (as a producer) only if full inner tiles are
    // iterated or inner dims are not tiled. Otherwise, it will generate a
    // sequence of non-trivial ops (for partial tiles).
    for (auto offset : offsets.take_back(numTiles))
      if (!isZeroInteger(offset))
        return failure();
 
    // Each requested inner-dim size must cover a full inner tile. A caller may
    // instead assert this via an `Equal` alignment hint. The hint is indexed by
    // source dim, matching the consumer-fusion path.
    ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
    SmallVector<OpFoldResult> mixedTiles = packOp.getMixedTiles();
    ArrayRef<OpFoldResult> innerSizes = sizes.take_back(numTiles);
    for (auto [i, pos] : llvm::enumerate(innerDimsPos)) {
      InnerTileAlignment alignment =
          pos < static_cast<int64_t>(innerTileAlignments.size())
              ? innerTileAlignments[pos]
              : InnerTileAlignment::Unknown;
      if (alignment != InnerTileAlignment::Equal &&
          !isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
        return failure();
    }
 
    FailureOr<TilingResult> tilingResult = getTiledImplementation(
        op, b, offsets.drop_back(numTiles), sizes.drop_back(numTiles));
    if (failed(tilingResult))
      return failure();
    return tilingResult.value();
  }
 
  LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
                                             Location loc,
                                             ValueRange ivs) const {
    auto packOp = cast<PackOp>(op);
    assert(packOp.hasPureBufferSemantics() &&
           "expected operation to have buffer semantics");
    OpBuilder::InsertionGuard g(builder);
    // The `ivs` already represent the position into the output for the non
    // data-tile dimensions.
    SmallVector<Value> ivVec(ivs);
 
    // Get output shape - for memrefs, get dimensions from dest directly.
    SmallVector<OpFoldResult> outputShape;
    Value dest = packOp.getDest();
    for (auto dim : llvm::seq<int64_t>(0, packOp.getDestRank()))
      outputShape.push_back(createOrFoldDimOp(builder, loc, dest, dim));
 
    // Generate the loops that iterate over the data tile.
    Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
    Value one = arith::ConstantIndexOp::create(builder, loc, 1);
 
    // All loops except the innermost are simple loops that just iterate
    // over the tile dimensions.
    for (auto dataTileDim : llvm::seq<unsigned>(packOp.getSourceRank(),
                                                packOp.getDestRank() - 1)) {
      Value ub = getValueOrCreateConstantIndexOp(builder, loc,
                                                 outputShape[dataTileDim]);
      scf::ForOp loop = scf::ForOp::create(builder, loc, zero, ub, one);
      builder.setInsertionPointToStart(loop.getBody());
      ivVec.push_back(loop.getInductionVar());
    }
    // The body of the innermost loops does the actual data movement.
    scf::ForOp::create(
        builder, loc, zero,
        getValueOrCreateConstantIndexOp(builder, loc, outputShape.back()), one,
        ValueRange{},
        [&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
            ValueRange regionIterArgs) {
          ivVec.push_back(iv);
          generatePackOpScalarImplementationBody(packOp, bodyBuilder, bodyLoc,
                                                 ivVec);
          scf::YieldOp::create(bodyBuilder, bodyLoc);
        });
    return success();
  }
 
  LogicalResult getIterationDomainTileFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes,
      SmallVectorImpl<OpFoldResult> &resultOffsets,
      SmallVectorImpl<OpFoldResult> &resultSizes) const {
    return getIterationDomainTileFromOperandTiles(
        op, b, operandNumbers, allOffsets, allSizes, resultOffsets, resultSizes,
        /*innerTileAlignments=*/{});
  }
 
  /// Method to return the position of iteration domain tile computed by the
  /// tiled operation. In current `linalg.pack` context, the `resultOffsets` and
  /// `resultSizes` only cover outer dimensions.
  LogicalResult getIterationDomainTileFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes,
      SmallVectorImpl<OpFoldResult> &resultOffsets,
      SmallVectorImpl<OpFoldResult> &resultSizes,
      ArrayRef<InnerTileAlignment> innerTileAlignments) const {
    if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
      LLVM_DEBUG(
          { llvm::dbgs() << "unsupported operands for consumer fusion"; });
      return failure();
    }
 
    ArrayRef<OpFoldResult> offsets(allOffsets[0]);
    ArrayRef<OpFoldResult> sizes(allSizes[0]);
    auto packOp = cast<PackOp>(op);
    Location loc = packOp.getLoc();
    SmallVector<OpFoldResult> outerDimOffsets, outerDimSizes;
    DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
        packOp.getDimAndTileMapping();
    SmallVector<int64_t> outerShapeWithoutTranspose(
        packOp.getDestType().getShape().take_front(packOp.getSourceRank()));
    if (!packOp.getOuterDimsPerm().empty()) {
      applyPermutationToVector(
          outerShapeWithoutTranspose,
          invertPermutationVector(packOp.getOuterDimsPerm()));
    }
    for (auto dim : llvm::seq<int64_t>(packOp.getSourceRank())) {
      if (dimAndTileMapping.count(dim)) {
        FailureOr<int64_t> cstTileSize =
            ValueBoundsConstraintSet::computeConstantBound(
                presburger::BoundType::UB, sizes[dim],
                /*stopCondition=*/nullptr,
                ValueBoundsOptions{/*closedUB=*/true});
        std::optional<int64_t> cstInnerSize =
            getConstantIntValue(dimAndTileMapping[dim]);
 
        // A caller-supplied alignment hint (see InnerTileAlignment) asserts
        // that this packed dimension is tiled and how its loop tile size
        // relates to the pack op inner tile size.
        InnerTileAlignment innerTileAlignment =
            dim < static_cast<int64_t>(innerTileAlignments.size())
                ? innerTileAlignments[dim]
                : InnerTileAlignment::Unknown;
 
        // If a dimension is not tiled, it is always valid to fuse the pack op,
        // even if the op has padding semantics. Because it always generates a
        // full slice along the dimension. The tile sizes are for unpacked
        // domain, i.e., `srcDimSize`, so `tileSize < srcDimSize` means that the
        // dimension is tiled.
        // TODO: It could be untiled if the `srcDimSize` is dynamic. It is a
        // hard check to determine if a dimension is tiled or not.
        // A non-`Unknown` hint also means the caller asserts the dimension is
        // tiled: `cstTileSize` is an upper bound, so a scalable/`min`-shaped
        // tile (whose bound equals `srcDimSize`) would otherwise be mistaken
        // for untiled and bypass the hint below.
        int64_t srcDimSize = packOp.getSourceType().getDimSize(dim);
        int64_t destDimSize = outerShapeWithoutTranspose[dim];
        bool isTiled = innerTileAlignment != InnerTileAlignment::Unknown ||
                       failed(cstTileSize) ||
                       ShapedType::isDynamic(srcDimSize) ||
                       cstTileSize.value() < srcDimSize;
        if (!isTiled) {
          outerDimOffsets.push_back(offsets[dim]);
          if (ShapedType::isStatic(destDimSize)) {
            outerDimSizes.push_back(b.getIndexAttr(destDimSize));
          } else {
            outerDimSizes.push_back(
                b.createOrFold<tensor::DimOp>(loc, packOp.getDest(), dim));
          }
          continue;
        }
 
        // Currently fusing `packOp` as consumer only expects perfect tiling
        // scenario because even if without padding semantic, the `packOp` may
        // also yield incomplete tiles. E.g. tensor<30xf32> -> tensor<5x6xf32>,
        // where the `tileSize` from operand of `packOp` is 5, which is not
        // exactly divided by `innerTile`(=6) of `packOp`. As the result:
        // 1. the first slice is extracted from (0) to (4) and inserted into
        // (0,0)~(0,4) at first row.
        // 2. the second slice is extracted from (5) to (9) and SHOULD BE
        // respectively inserted into two rows with different length, including
        // first row: (0,5) and second row (1,0)~(1,3). It is hard to coordinate
        // them, thus adding below constraint to bypass them temporarily. In
        // another word, we can only support tiling with consumer if the tile
        // size for the producer is a multiple of the inner tile size for the
        // packed dimensions at this moment.
 
        // The caller may assert how this packed dimension's loop tile size
        // relates to the inner tile size via `innerTileAlignments` (see
        // InnerTileAlignment). The hint is the source of truth and is honored
        // when present. When both sizes are also statically known we assert the
        // hint agrees with them (a contradicting hint is a caller bug). When
        // the hint is `Unknown`, fall back to requiring a statically-provable
        // multiple.
        bool assumeInnerTileSizesMatchTiles =
            innerTileAlignment == InnerTileAlignment::Equal;
        bool staticallyDecidable =
            !failed(cstTileSize) && cstInnerSize.has_value();
        if (innerTileAlignment == InnerTileAlignment::Unknown) {
          if (!staticallyDecidable || *cstTileSize % *cstInnerSize != 0)
            return failure();
        } else if (staticallyDecidable) {
          assert(*cstTileSize % *cstInnerSize == 0 &&
                 "InnerTileAlignment hint contradicts statically known tile "
                 "sizes");
          assert((innerTileAlignment != InnerTileAlignment::Equal ||
                  *cstTileSize == *cstInnerSize) &&
                 "InnerTileAlignment::Equal contradicts statically known tile "
                 "sizes");
        }
 
        using AV = affine::AffineValueExpr;
        affine::AffineBuilder ab(b, loc);
        AffineExpr dim0, sym;
        bindDims(b.getContext(), dim0);
        bindSymbols(b.getContext(), sym);
        auto avOffset = AV(dim0).bind(offsets[dim]);
        auto avSize = AV(dim0).bind(sizes[dim]);
        auto avTileSize = AV(sym).bind(dimAndTileMapping[dim]);
        outerDimOffsets.push_back(ab.floor(avOffset, avTileSize));
        // If the tile size equals the inner tile size, the outer dims are
        // always 1.
        outerDimSizes.push_back(assumeInnerTileSizesMatchTiles
                                    ? b.getIndexAttr(1)
                                    : ab.ceil(avSize, avTileSize));
      } else {
        outerDimOffsets.push_back(offsets[dim]);
        outerDimSizes.push_back(sizes[dim]);
      }
    }
    applyPermToRange(outerDimOffsets, outerDimSizes, packOp.getOuterDimsPerm());
    resultOffsets = outerDimOffsets;
    resultSizes = outerDimSizes;
    return success();
  }
 
  FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
    return getTiledImplementationFromOperandTiles(op, b, operandNumbers,
                                                  allOffsets, allSizes,
                                                  /*innerTileAlignments=*/{});
  }
 
  /// Method to return the tiled implementation of linalg.pack as a consumer.
  FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes,
      ArrayRef<InnerTileAlignment> innerTileAlignments) const {
    if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
      LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
      return failure();
    }
 
    ArrayRef<OpFoldResult> offsets(allOffsets[0]);
    ArrayRef<OpFoldResult> sizes(allSizes[0]);
 
    auto packOp = cast<PackOp>(op);
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!packOp.hasPureTensorSemantics())
      return failure();
 
    Location loc = packOp.getLoc();
 
    int64_t inputRank = packOp.getSourceRank();
    auto oneAttr = b.getI64IntegerAttr(1);
    SmallVector<OpFoldResult> strides(inputRank, oneAttr);
 
    SmallVector<Value> tiledOperands;
    auto sourceSlice = tensor::ExtractSliceOp::create(
        b, loc, packOp.getSource(), offsets, sizes, strides);
    tiledOperands.push_back(sourceSlice);
 
    SmallVector<OpFoldResult> outerDimOffsets, outerDimSizes;
    if (failed(getIterationDomainTileFromOperandTiles(
            op, b, operandNumbers, allOffsets, allSizes, outerDimOffsets,
            outerDimSizes, innerTileAlignments)))
      return failure();
 
    SmallVector<OpFoldResult> outputOffsets, outputSizes;
    if (failed(getResultTilePosition(op, b, 0, outerDimOffsets, outerDimSizes,
                                     outputOffsets, outputSizes)))
      return failure();
 
    strides.append(packOp.getDestRank() - inputRank, oneAttr);
    auto outSlice = tensor::ExtractSliceOp::create(
        b, loc, packOp.getDest(), outputOffsets, outputSizes, strides);
    tiledOperands.push_back(outSlice);
 
    if (auto val = packOp.getPaddingValue())
      tiledOperands.push_back(val);
    for (auto tile : packOp.getInnerTiles())
      tiledOperands.push_back(tile);
 
    Operation *tiledPackOp = PackOp::create(
        b, loc, TypeRange{outSlice.getType()}, tiledOperands, op->getAttrs());
 
    return TilingResult{
        {tiledPackOp},
        SmallVector<Value>(tiledPackOp->getResults()),
        llvm::to_vector(ArrayRef<Operation *>{sourceSlice, outSlice})};
  }
};
 
struct UnpackTileDimInfo {
  bool isAlignedToInnerTileSize;
  OpFoldResult sourceOffset;
  OpFoldResult sourceSize;
  OpFoldResult resultOffset;
  OpFoldResult destExpandedSize;
};
 
/// Returns the needed information for tiling unpack op on `tileDim` with given
/// `tileOffset` and `tileSize`. For more details, see the comment of the
/// `getTiledImplementation`.
static UnpackTileDimInfo
getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp, int64_t tileDim,
                     OpFoldResult tileOffset, OpFoldResult tileSize,
                     InnerTileAlignment innerTileAlignment) {
  UnpackTileDimInfo info;
  Attribute zeroAttr = b.getIndexAttr(0);
  Attribute oneAttr = b.getIndexAttr(1);
  DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
      unpackOp.getDimAndTileMapping();
  // The dimension is not one of packed data dimension.
  if (!dimAndTileMapping.count(tileDim)) {
    info.isAlignedToInnerTileSize = true;
    info.sourceOffset = tileOffset;
    info.sourceSize = tileSize;
    info.resultOffset = zeroAttr;
    info.destExpandedSize = tileSize;
    return info;
  }
 
  Location loc = unpackOp.getLoc();
  using AV = affine::AffineValueExpr;
  affine::AffineBuilder ab(b, loc);
  AffineExpr dim0, dim1, sym0;
  bindDims(b.getContext(), dim0, dim1);
  bindSymbols(b.getContext(), sym0);
 
  OpFoldResult innerTileSize = dimAndTileMapping[tileDim];
 
  info.isAlignedToInnerTileSize = false;
  FailureOr<int64_t> cstSize = ValueBoundsConstraintSet::computeConstantBound(
      presburger::BoundType::UB, tileSize,
      /*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
  std::optional<int64_t> cstInnerSize = getConstantIntValue(innerTileSize);
  // The caller may assert how this dimension's loop tile size relates to the
  // op's inner tile size via `innerTileAlignment` (see InnerTileAlignment). The
  // hint is the source of truth and is honored when present: `Equal`/`Multiple`
  // both mean the tile is aligned to (a multiple of) the inner tile, and
  // `Equal` additionally collapses the source slice to a single inner tile.
  // When both sizes are also statically known we assert the hint agrees with
  // them (a contradicting hint is a caller bug). When `Unknown`, fall back to
  // the static upper-bound path below.
  bool assumeInnerTileSizesMatchTiles =
      innerTileAlignment == InnerTileAlignment::Equal;
  bool staticallyDecidable = !failed(cstSize) && cstInnerSize.has_value();
  if (innerTileAlignment != InnerTileAlignment::Unknown) {
    info.isAlignedToInnerTileSize = true;
    if (staticallyDecidable) {
      assert(*cstSize % *cstInnerSize == 0 &&
             "InnerTileAlignment hint contradicts statically known tile sizes");
      assert((innerTileAlignment != InnerTileAlignment::Equal ||
              *cstSize == *cstInnerSize) &&
             "InnerTileAlignment::Equal contradicts statically known tile "
             "sizes");
    }
  }
  if (info.isAlignedToInnerTileSize || (!failed(cstSize) && cstInnerSize)) {
    if (!info.isAlignedToInnerTileSize && *cstSize % *cstInnerSize == 0)
      info.isAlignedToInnerTileSize = true;
 
    // If the tiling size equals to the inner tiling size, the outer dims are
    // always 1.
    if (assumeInnerTileSizesMatchTiles ||
        (cstInnerSize && !failed(cstSize) && *cstInnerSize == *cstSize)) {
      auto lhs = AV(dim0).bind(tileOffset);
      auto rhs = AV(dim1).bind(innerTileSize);
      info.sourceOffset = ab.floor(lhs, rhs);
      info.sourceSize = oneAttr;
      info.resultOffset = zeroAttr;
      info.destExpandedSize = tileSize;
      return info;
    }
  }
 
  if (info.isAlignedToInnerTileSize) {
    info.sourceOffset =
        ab.floor(AV(dim0).bind(tileOffset), AV(dim1).bind(innerTileSize));
    info.resultOffset = zeroAttr;
    info.destExpandedSize = tileSize;
 
    // The ceilDiv is needed here because there could be incomplete tile even
    // it is perfect tiling cases. E.g.,
    //   %0 = unpack tensor<33x2xf32> into tensor<64xf32>
    // If the tiling size is 32, there will be 3 tiles. Two of them have
    // size=32; one of them have size=2. The size is represented using
    // affine_min op; we need ceilDiv.
    info.sourceSize =
        ab.ceil(AV(dim0).bind(tileSize), AV(dim1).bind(innerTileSize));
    return info;
  }
 
  affine::DivModValue firstCoord = affine::getDivMod(
      b, loc, getValueOrCreateConstantIndexOp(b, loc, tileOffset),
      getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
  OpFoldResult tileExclusiveBound =
      ab.add(AV(dim0).bind(tileOffset), AV(dim1).bind(tileSize));
  affine::DivModValue lastCoord = affine::getDivMod(
      b, loc,
      getValueOrCreateConstantIndexOp(
          b, loc,
          ab.sub(AV(dim0).bind(tileExclusiveBound), AV(dim1).bind(oneAttr))),
      getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
 
  OpFoldResult lengthMinusOne = ab.sub(AV(dim0).bind(lastCoord.quotient),
                                       AV(dim1).bind(firstCoord.quotient));
  info.sourceSize =
      ab.add(AV(dim0).bind(lengthMinusOne), AV(dim1).bind(oneAttr));
  info.sourceOffset = firstCoord.quotient;
  info.resultOffset = firstCoord.remainder;
  // Do not create an Affine ops for expanded size because the affine op is too
  // complicated which would trigger an issue in affine ops simplification.
  info.destExpandedSize = b.createOrFold<arith::MulIOp>(
      loc, getValueOrCreateConstantIndexOp(b, loc, info.sourceSize),
      getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
  return info;
}
 
struct UnPackOpTiling
    : public TilingInterface::ExternalModel<UnPackOpTiling, linalg::UnPackOp> {
  using Base = TilingInterface::ExternalModel<UnPackOpTiling, linalg::UnPackOp>;
  using Base::getIterationDomainTileFromOperandTiles;
 
  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
    auto unpackOp = cast<UnPackOp>(op);
    SmallVector<utils::IteratorType> iteratorTypes(
        unpackOp.getDestRank(), utils::IteratorType::parallel);
    return iteratorTypes;
  }
 
  SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
    return getPackUnPackIterationDomain<UnPackOp>(cast<UnPackOp>(op), b);
  }
 
  /// There are two cases in tiling unpack ops. If the tiling size is aligned to
  /// the inner tile size, the corresponding tiles of source are all complete.
  /// Otherwise, there are in-complete tiles. We will need to expand the slice
  /// of source for getting complete tiles. The tiled unpack op unpacks more
  /// data from source, so We'll need an extract_slice op to shift and truncate
  /// the output.
  /// Take Nn_to_N as an example. Say that N=32, n=8, and tiling_size=15. The
  /// coordinates of second tile (i.e., result[15..31]) are
  /// [(1, 7), (2, 0,), (2, 1) ... (3, 6), (3, 7)]. The first row and the last
  /// row are incomplete tiles. To represent the unpack op, we have to complete
  /// the rows. I.e., the input coordinates would start with (1, 0); end with
  /// (3, 7). In this context, the tiled unpack produces a (3 * n) elements
  /// because there are 3 rows in total. Follow by a tensor.extract_slice op, we
  /// can get the actual result.
  FailureOr<TilingResult>
  getTiledImplementation(Operation *op, OpBuilder &b,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes) const {
    return getTiledImplementation(op, b, offsets, sizes,
                                  /*innerTileAlignments=*/{});
  }
 
  FailureOr<TilingResult> getTiledImplementation(
      Operation *op, OpBuilder &b, ArrayRef<OpFoldResult> offsets,
      ArrayRef<OpFoldResult> sizes,
      ArrayRef<InnerTileAlignment> innerTileAlignments) const {
    auto unpackOp = cast<UnPackOp>(op);
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unpackOp.hasPureTensorSemantics())
      return failure();
 
    int64_t srcRank = unpackOp.getSourceRank();
    int64_t destRank = unpackOp.getDestRank();
    int64_t numInnerTiles = srcRank - destRank;
    Location loc = unpackOp.getLoc();
 
    // The perfect tiling case indicates that the tiling sizes are multiple of
    // inner_tile_size. In this context, no extra data is needed when
    // representing the tiled unpack op.
    bool isPerfectTilingCase = true;
    Attribute oneAttr = b.getIndexAttr(1);
    SmallVector<OpFoldResult> sliceSrcStrides(destRank, oneAttr);
    SmallVector<OpFoldResult> sliceSrcIndices, sliceSrcSizes;
    SmallVector<OpFoldResult> destExpandedSizes, resultOffsetsFromDest;
    for (auto dim : llvm::seq<int64_t>(0, destRank)) {
      UnpackTileDimInfo info = getUnpackTileDimInfo(
          b, unpackOp, dim, offsets[dim], sizes[dim],
          dim < static_cast<int64_t>(innerTileAlignments.size())
              ? innerTileAlignments[dim]
              : InnerTileAlignment::Unknown);
      if (!info.isAlignedToInnerTileSize)
        isPerfectTilingCase = false;
      sliceSrcIndices.push_back(info.sourceOffset);
      sliceSrcSizes.push_back(info.sourceSize);
      destExpandedSizes.push_back(info.destExpandedSize);
      resultOffsetsFromDest.push_back(info.resultOffset);
    }
 
    // The tiling is applied on destination dimensions. We have to apply the
    // interchange on source dimensions if outer_dims_perm is set.
    applyPermToRange(sliceSrcIndices, sliceSrcSizes,
                     unpackOp.getOuterDimsPerm());
    Attribute zeroAttr = b.getIndexAttr(0);
    sliceSrcIndices.append(numInnerTiles, zeroAttr);
    sliceSrcSizes.append(unpackOp.getMixedTiles());
    sliceSrcStrides.append(numInnerTiles, oneAttr);
    SmallVector<Operation *> generatedSlices;
    tensor::ExtractSliceOp sliceSource = tensor::ExtractSliceOp::create(
        b, loc, unpackOp.getSource(), sliceSrcIndices, sliceSrcSizes,
        sliceSrcStrides);
    generatedSlices.push_back(sliceSource);
 
    SmallVector<OpFoldResult> destStrides(destRank, oneAttr);
    Value sliceDest;
    if (isPerfectTilingCase) {
      auto destSliceOp = tensor::ExtractSliceOp::create(
          b, loc, unpackOp.getDest(), offsets, sizes, destStrides);
      sliceDest = destSliceOp;
      generatedSlices.push_back(destSliceOp);
    } else {
      sliceDest = tensor::EmptyOp::create(
          b, loc, destExpandedSizes, unpackOp.getDestType().getElementType());
    }
 
    SmallVector<Value> tiledOperands = {sliceSource.getResult(), sliceDest};
    for (auto tile : unpackOp.getInnerTiles())
      tiledOperands.push_back(tile);
 
    Operation *tiledUnpackOp = UnPackOp::create(
        b, loc, TypeRange{sliceDest.getType()}, tiledOperands, op->getAttrs());
 
    if (isPerfectTilingCase)
      return TilingResult{{tiledUnpackOp},
                          SmallVector<Value>(tiledUnpackOp->getResults()),
                          generatedSlices};
 
    auto extractSlice = tensor::ExtractSliceOp::create(
        b, loc, tiledUnpackOp->getResult(0), resultOffsetsFromDest, sizes,
        destStrides);
    return TilingResult{
        {tiledUnpackOp}, {extractSlice.getResult()}, generatedSlices};
  }
 
  LogicalResult
  getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                        ArrayRef<OpFoldResult> offsets,
                        ArrayRef<OpFoldResult> sizes,
                        SmallVector<OpFoldResult> &resultOffsets,
                        SmallVector<OpFoldResult> &resultSizes) const {
    resultOffsets = llvm::to_vector(offsets);
    resultSizes = llvm::to_vector(sizes);
    return success();
  }
 
  FailureOr<TilingResult>
  generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes) const {
    return generateResultTileValue(op, b, resultNumber, offsets, sizes,
                                   /*innerTileAlignments=*/{});
  }
 
  FailureOr<TilingResult> generateResultTileValue(
      Operation *op, OpBuilder &b, unsigned resultNumber,
      ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
      ArrayRef<InnerTileAlignment> innerTileAlignments) const {
    FailureOr<TilingResult> tilingResult =
        getTiledImplementation(op, b, offsets, sizes, innerTileAlignments);
    if (failed(tilingResult))
      return failure();
    return tilingResult.value();
  }
 
  LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
                                             Location loc,
                                             ValueRange ivs) const {
    auto unpackOp = cast<UnPackOp>(op);
    assert(unpackOp.hasPureBufferSemantics() &&
           "expected operation to have buffer semantics");
    assert(ivs.size() == unpackOp.getDestRank() &&
           "number of ivs must match the rank of the output tensor");
    OpBuilder::InsertionGuard g(builder);
 
    DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
        unpackOp.getDimAndTileMapping();
    // Untiled loops and tile loops induction variables.
    SmallVector<Value> inputIvs;
    // Point loops induction variables.
    SmallVector<Value> inputIvsPointLoops;
    inputIvs.reserve(unpackOp.getDestRank());
    inputIvsPointLoops.reserve(dimAndTileMapping.size());
    for (auto dim : llvm::seq<int64_t>(0, unpackOp.getDestRank())) {
      if (dimAndTileMapping.count(dim)) {
        affine::DivModValue divMod =
            affine::getDivMod(builder, loc, ivs[dim],
                              getValueOrCreateConstantIndexOp(
                                  builder, loc, dimAndTileMapping[dim]));
        inputIvsPointLoops.push_back(divMod.remainder);
        inputIvs.push_back(divMod.quotient);
      } else {
        inputIvs.push_back(ivs[dim]);
      }
    }
 
    // TODO: (lorenzo) simplify the logic a bit. There is `ivs`,
    // `inputIvsPointLoops` and `inputIvs`.
    assert(inputIvsPointLoops.size() + inputIvs.size() ==
               unpackOp.getSourceRank() &&
           "expect same number of induction variables equals to input rank");
    // Interchange the point loops induction variables based on `inner_dim_pos`.
    ArrayRef<int64_t> innerDims = unpackOp.getInnerDimsPos();
    SmallVector<int64_t> interchangeVector =
        computeInterchangeFromDimPos(innerDims, unpackOp.getDestRank());
    SmallVector<Value> interchangedInputIvsPointLoops = inputIvsPointLoops;
    interchangedInputIvsPointLoops = interchange<Value>(
        interchangedInputIvsPointLoops, interchangeVector, /*offset=*/0);
    // Interchange the tiled loops induction variables based on
    // `outer_dims_perm`.
    ArrayRef<int64_t> outerDims = unpackOp.getOuterDimsPerm();
    if (!outerDims.empty())
      inputIvs = interchange<Value>(inputIvs, outerDims, /*offset=*/0);
 
    llvm::append_range(inputIvs, interchangedInputIvsPointLoops);
    Value scalar =
        memref::LoadOp::create(builder, loc, unpackOp.getSource(), inputIvs);
    memref::StoreOp::create(builder, loc, scalar, unpackOp.getDest(), ivs);
    return success();
  }
 
  /// Method to return the position of iteration domain tile computed by the
  /// tiled operation.
  LogicalResult getIterationDomainTileFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes,
      SmallVectorImpl<OpFoldResult> &resultOffsets,
      SmallVectorImpl<OpFoldResult> &resultSizes) const {
    if (operandNumbers.size() != 1) {
      LLVM_DEBUG({ llvm::dbgs() << "unable to handle multiple operands"; });
      return failure();
    }
    auto unPackOp = cast<UnPackOp>(op);
    unsigned operandNumber = operandNumbers[0];
    ArrayRef<OpFoldResult> offsets(allOffsets[0]);
    ArrayRef<OpFoldResult> sizes(allSizes[0]);
 
    // If the operand tile is the dest, then no adjustment is needed.
    if (operandNumber == unPackOp.getDestMutable().getOperandNumber()) {
      resultOffsets = llvm::to_vector(offsets);
      resultSizes = llvm::to_vector(sizes);
      return success();
    }
    Location loc = unPackOp.getLoc();
 
    int64_t numTiles = unPackOp.getInnerDimsPos().size();
    auto destOffsets = offsets.drop_back(numTiles);
    auto destSizes = sizes.drop_back(numTiles);
    // The tiling is applied on interchanged dimensions. We have to undo the
    // interchange to map sizes and offsets to the original input.
    int64_t outputRank = unPackOp.getDestRank();
    ReifiedRankedShapedTypeDims reifiedReturnShapes;
    if (failed(reifyResultShapes(b, unPackOp, reifiedReturnShapes)))
      return failure();
    SmallVector<OpFoldResult> outputMixedSizes = reifiedReturnShapes.front();
    SmallVector<OpFoldResult> origOffsets(destOffsets);
    SmallVector<OpFoldResult> origSizes(destSizes);
    applyPermToRange(origOffsets, origSizes,
                     invertPermutationVector(unPackOp.getOuterDimsPerm()));
 
    DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
        unPackOp.getDimAndTileMapping();
 
    for (auto dim : llvm::seq<int64_t>(0, outputRank)) {
      using AV = affine::AffineValueExpr;
      affine::AffineBuilder ab(b, loc);
      AffineExpr dim0, dim1, sym0;
      bindDims(b.getContext(), dim0, dim1);
      bindSymbols(b.getContext(), sym0);
      if (dimAndTileMapping.count(dim)) {
        // If the data dimension is tiled, the i-th index is the product of
        // offset_i and tile_i, and the i-th size is the product of sizes_i and
        // tile_i. The sizes must be clamped to the sizes of the unpack result.
        auto avOffset = AV(dim0).bind(origOffsets[dim]);
        auto avSize = AV(dim0).bind(origSizes[dim]);
        auto avTileSize = AV(sym0).bind(dimAndTileMapping[dim]);
        auto avResultSize = AV(dim0).bind(outputMixedSizes[dim]);
        resultOffsets.push_back(ab.mul(avOffset, avTileSize));
        auto avResultOffset = AV(dim1).bind(resultOffsets.back());
        resultSizes.push_back(ab.min({ab.mul(avSize, avTileSize),
                                      ab.sub(avResultSize, avResultOffset)}));
      } else {
        resultOffsets.push_back(origOffsets[dim]);
        resultSizes.push_back(origSizes[dim]);
      }
    }
    return success();
  }
 
  FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
    return getTiledImplementationFromOperandTiles(op, b, operandNumbers,
                                                  allOffsets, allSizes,
                                                  /*innerTileAlignments=*/{});
  }
 
  /// Method to return the tiled implementation of linalg.unpack as a consumer.
  FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
      ArrayRef<SmallVector<OpFoldResult>> allSizes,
      ArrayRef<InnerTileAlignment> innerTileAlignments) const {
    if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
      LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
      return failure();
    }
    auto unPackOp = cast<UnPackOp>(op);
    // TODO: Support Memref UnPackOp. Temporarily return failure.
    if (!unPackOp.hasPureTensorSemantics())
      return failure();
 
    ArrayRef<OpFoldResult> offsets(allOffsets[0]);
    ArrayRef<OpFoldResult> sizes(allSizes[0]);
 
    // linalg.unpack op is fusible (as a consumer) only if the inner dims are
    // not tiled, i.e. each inner-dim loop tile size equals the inner tile size.
    // The caller may assert this per inner dim via InnerTileAlignment::Equal;
    // otherwise we require a statically-provable equality.
    int64_t numTiles = unPackOp.getInnerDimsPos().size();
    ArrayRef<int64_t> innerDimsPos = unPackOp.getInnerDimsPos();
    SmallVector<OpFoldResult> mixedTiles = unPackOp.getMixedTiles();
    ArrayRef<OpFoldResult> innerSizes = sizes.take_back(numTiles);
    for (int64_t i = 0; i < numTiles; ++i) {
      // `innerTileAlignments` is indexed by the unpack iteration domain (the
      // dest dims); the i-th inner tile lives on dest dim `innerDimsPos[i]`.
      int64_t destDim = innerDimsPos[i];
      bool hintedEqual =
          destDim < static_cast<int64_t>(innerTileAlignments.size()) &&
          innerTileAlignments[destDim] == InnerTileAlignment::Equal;
      // The hint is the source of truth: honor a caller `Equal` assertion. When
      // both sizes are also statically known, assert the hint agrees with them
      // (a contradicting hint is a caller bug) rather than silently ignoring
      // it. Without an `Equal` hint, require a statically-provable equality
      // (the inner dim must not be tiled).
      if (hintedEqual) {
        assert((!getConstantIntValue(mixedTiles[i]) ||
                !getConstantIntValue(innerSizes[i]) ||
                isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i])) &&
               "InnerTileAlignment::Equal contradicts statically known tile "
               "sizes");
        continue;
      }
      if (isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
        continue;
      return failure();
    }
 
    Location loc = unPackOp.getLoc();
 
    // Fetch offset/size for creating the slice of the dest operand of
    // unpack op.
    SmallVector<OpFoldResult> outputOffsets, outputSizes;
    if (failed(getIterationDomainTileFromOperandTiles(
            op, b, operandNumbers, allOffsets, allSizes, outputOffsets,
            outputSizes)))
      return failure();
 
    auto oneAttr = b.getI64IntegerAttr(1);
    int64_t outputRank = unPackOp.getDestRank();
    SmallVector<OpFoldResult> strides(outputRank, oneAttr);
 
    SmallVector<Value> tiledOperands;
    // Create slice of the dest operand.
    auto extractDestSlice = tensor::ExtractSliceOp::create(
        b, loc, unPackOp.getDest(), outputOffsets, outputSizes, strides);
    tiledOperands.push_back(extractDestSlice);
 
    strides.append(unPackOp.getSourceRank() - outputRank, oneAttr);
    // Create slice of the source operand.
    auto extractSourceSlice = tensor::ExtractSliceOp::create(
        b, loc, unPackOp.getSource(), offsets, sizes, strides);
    tiledOperands.insert(tiledOperands.begin(), extractSourceSlice);
    for (auto tile : unPackOp.getInnerTiles())
      tiledOperands.push_back(tile);
 
    // Create tiled unpack op.
    Operation *tiledUnPackOp =
        UnPackOp::create(b, loc, TypeRange{extractDestSlice.getType()},
                         tiledOperands, op->getAttrs());
 
    return TilingResult{{tiledUnPackOp},
                        SmallVector<Value>(tiledUnPackOp->getResults()),
                        llvm::to_vector(ArrayRef<Operation *>{
                            extractSourceSlice, extractDestSlice})};
  }
};
 
} // namespace
 
template <typename OpType>
static void registerOne(MLIRContext *ctx) {
  OpType::template attachInterface<LinalgOpTilingInterface<OpType>>(*ctx);
  OpType::template attachInterface<LinalgOpPartialReductionInterface<OpType>>(
      *ctx);
}
 
/// Variadic helper function.
template <typename... OpTypes>
static void registerAll(MLIRContext *ctx) {
  (registerOne<OpTypes>(ctx), ...);
}
 
#define GET_OP_LIST
 
void mlir::linalg::registerTilingInterfaceExternalModels(
    DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
    registerOne<linalg::GenericOp>(ctx);
    linalg::PackOp::attachInterface<PackOpTiling>(*ctx);
    linalg::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
    registerAll<
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
        >(ctx);
  });
}
 
void mlir::linalg::registerTilingInterfaceExternalModelsForPackUnPackOps(
    DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *dialect) {
    linalg::PackOp::attachInterface<PackOpTiling>(*ctx);
    linalg::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
  });
}