TheRock/compiler/amd-llvm/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

//===- DropUnitDims.cpp - Pass to drop use of unit-extent for broadcasting ===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements patterns/pass to remove usage of unit-extent dimensions
// to specify broadcasting in favor of more canonical representation of the
// computation
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Linalg/Passes.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
#include "llvm/Support/Debug.h"

namespace mlir {
#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS
#include "mlir/Dialect/Linalg/Passes.h.inc"
} // namespace mlir

#define DEBUG_TYPE "linalg-drop-unit-dims"

using namespace mlir;
using namespace mlir::linalg;

namespace {
/// Pattern to move init operands to ins when all the loops are parallel and
/// blockArgument corresponding to init is used in the region. This is a fix-up
/// when unit reduction dimensions are all folded away. In this context, it
/// becomes a elementwise generic op. E.g., it converts
///
///  %0 = tensor.empty() : tensor<1x1xf32>
///  %1 = linalg.fill
///    ins(%cst : f32)
///    outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>
///  %2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,
///                                        affine_map<(d0) -> (0, d0)>],
///                       iterator_types = ["parallel"]}
///    ins(%arg0 : tensor<1x?x1x1xf32>)
///    outs(%1 : tensor<1x1xf32>) {
///  ^bb0(%in: f32, %out: f32):
///    %3 = arith.addf %in, %out : f32
///    linalg.yield %3 : f32
///  } -> tensor<1x1xf32>
///
///  into
///
///  %0 = tensor.empty() : tensor<1x1xf32>
///  %1 = linalg.fill
///    ins(%cst : f32)
///    outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>
///  %2 = tensor.empty() : tensor<1x1xf32>
///  %3 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,
///                                        affine_map<(d0) -> (0, d0)>,
///                                        affine_map<(d0) -> (0, d0)>],
///                       iterator_types = ["parallel"]}
///   ins(%arg0, %1 : tensor<1x?x1x1xf32>, tensor<1x1xf32>)
///   outs(%2 : tensor<1x1xf32>) {
///  ^bb0(%in: f32, %in_0: f32, %out: f32):
///    %4 = arith.addf %in, %in_0 : f32
///    linalg.yield %4 : f32
///  } -> tensor<1x1xf32>
struct MoveInitOperandsToInput : public OpRewritePattern<GenericOp> {
  using OpRewritePattern<GenericOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(GenericOp genericOp,
                                PatternRewriter &rewriter) const override {
    if (!genericOp.hasPureTensorSemantics())
      return failure();
    if (genericOp.getNumParallelLoops() != genericOp.getNumLoops())
      return failure();

    auto outputOperands = genericOp.getDpsInitsMutable();
    SetVector<OpOperand *> candidates;
    for (OpOperand &op : outputOperands) {
      if (genericOp.getMatchingBlockArgument(&op).use_empty())
        continue;
      candidates.insert(&op);
    }

    if (candidates.empty())
      return failure();

    // Compute the modified indexing maps.
    int64_t origNumInput = genericOp.getNumDpsInputs();
    SmallVector<Value> newInputOperands = genericOp.getDpsInputs();
    SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
    SmallVector<AffineMap> newIndexingMaps;
    newIndexingMaps.append(indexingMaps.begin(),
                           std::next(indexingMaps.begin(), origNumInput));
    for (OpOperand *op : candidates) {
      newInputOperands.push_back(op->get());
      newIndexingMaps.push_back(genericOp.getMatchingIndexingMap(op));
    }
    newIndexingMaps.append(std::next(indexingMaps.begin(), origNumInput),
                           indexingMaps.end());

    Location loc = genericOp.getLoc();
    SmallVector<Value> newOutputOperands =
        llvm::to_vector(genericOp.getDpsInits());
    for (OpOperand *op : candidates) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointAfterValue(op->get());
      auto elemType = cast<ShapedType>(op->get().getType()).getElementType();
      auto empty = tensor::EmptyOp::create(
          rewriter, loc, tensor::getMixedSizes(rewriter, loc, op->get()),
          elemType);

      unsigned start = genericOp.getDpsInits().getBeginOperandIndex();
      newOutputOperands[op->getOperandNumber() - start] = empty.getResult();
    }

    auto newOp = GenericOp::create(
        rewriter, loc, genericOp.getResultTypes(), newInputOperands,
        newOutputOperands, newIndexingMaps, genericOp.getIteratorTypesArray(),
        /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));

    OpBuilder::InsertionGuard guard(rewriter);
    Region &region = newOp.getRegion();
    Block *block = rewriter.createBlock(&region);
    IRMapping mapper;
    for (auto bbarg : genericOp.getRegionInputArgs())
      mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));

    for (OpOperand *op : candidates) {
      BlockArgument bbarg = genericOp.getMatchingBlockArgument(op);
      mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
    }

    for (OpOperand &op : outputOperands) {
      BlockArgument bbarg = genericOp.getMatchingBlockArgument(&op);
      if (candidates.count(&op))
        block->addArgument(bbarg.getType(), loc);
      else
        mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
    }

    for (auto &op : genericOp.getBody()->getOperations()) {
      rewriter.clone(op, mapper);
    }
    rewriter.replaceOp(genericOp, newOp.getResults());

    return success();
  }
};
} // namespace

//===---------------------------------------------------------------------===//
// Drop loops that are unit-extents within Linalg operations.
//===---------------------------------------------------------------------===//

/// Implements a pass that canonicalizes the uses of unit-extent dimensions for
/// broadcasting. For example,
///
/// ```mlir
/// #accesses = [
///   affine_map<(d0, d1) -> (0, d1)>,
///   affine_map<(d0, d1) -> (d0, 0)>,
///   affine_map<(d0, d1) -> (d0, d1)>
/// ]
///
/// #trait = {
///   indexing_maps = #accesses,
///   iterator_types = ["parallel", "parallel"],
///   library_call = "some_external_fn"
/// }
///
/// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
/// tensor<5x5xf32>
/// {
///   %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :
///        tensor<5xf32> into tensor<1x5xf32>
///   %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :
///        tensor<5xf32> into tensor<5x1xf32>
///   %2 = linalg.generic #trait %0, %1 {
///        ^bb0(%arg2: f32, %arg3: f32):
///          %3 = arith.addf %arg2, %arg3 : f32
///          linalg.yield %3 : f32
///        } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>
///   return %2 : tensor<5x5xf32>
/// }
///
/// would canonicalize to
///
/// ```mlir
/// #accesses = [
///   affine_map<(d0, d1) -> (d1)>,
///   affine_map<(d0, d1) -> (d0)>,
///   affine_map<(d0, d1) -> (d0, d1)>
/// ]
///
/// #trait = {
///   indexing_maps = #accesses,
///   iterator_types = ["parallel", "parallel"],
///   library_call = "some_external_fn"
/// }
///
/// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
/// tensor<5x5xf32>
/// {
///   %0 = linalg.generic #trait %arg0, %arg1 {
///        ^bb0(%arg2: f32, %arg3: f32):
///          %3 = arith.addf %arg2, %arg3 : f32
///          linalg.yield %3 : f32
///        } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32>
///   return %0 : tensor<5x5xf32>
/// }

/// Update the index accesses of linalg operations having index semantics.
static void
replaceUnitDimIndexOps(GenericOp genericOp,
                       const llvm::SmallDenseSet<unsigned> &unitDims,
                       RewriterBase &rewriter) {
  for (IndexOp indexOp :
       llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPoint(indexOp);
    if (unitDims.count(indexOp.getDim()) != 0) {
      rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);
    } else {
      // Update the dimension of the index operation if needed.
      unsigned droppedDims = llvm::count_if(
          unitDims, [&](unsigned dim) { return dim < indexOp.getDim(); });
      if (droppedDims != 0)
        rewriter.replaceOpWithNewOp<IndexOp>(indexOp,
                                             indexOp.getDim() - droppedDims);
    }
  }
}

FailureOr<Value>
ControlDropUnitDims::expandValue(RewriterBase &rewriter, Location loc,
                                 Value result, Value origDest,
                                 ArrayRef<ReassociationIndices> reassociation,
                                 const ControlDropUnitDims &control) {
  // There are no results for memref outputs.
  auto origResultType = cast<RankedTensorType>(origDest.getType());
  if (origResultType.getEncoding() != nullptr) {
    // Do not expand tensors with encoding.
    return failure();
  }
  if (control.rankReductionStrategy ==
      ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {
    unsigned rank = origResultType.getRank();
    SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
    SmallVector<OpFoldResult> sizes =
        tensor::getMixedSizes(rewriter, loc, origDest);
    SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
    return rewriter.createOrFold<tensor::InsertSliceOp>(
        loc, result, origDest, offsets, sizes, strides);
  }

  assert(control.rankReductionStrategy ==
             ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&
         "unknown rank reduction strategy");
  return tensor::ExpandShapeOp::create(rewriter, loc, origResultType, result,
                                       reassociation)
      .getResult();
}

FailureOr<Value>
ControlDropUnitDims::collapseValue(RewriterBase &rewriter, Location loc,
                                   Value operand, ArrayRef<int64_t> targetShape,
                                   ArrayRef<ReassociationIndices> reassociation,
                                   const ControlDropUnitDims &control) {
  if (auto memrefType = dyn_cast<MemRefType>(operand.getType())) {
    if (!memrefType.getLayout().isIdentity()) {
      // Do not collapse memrefs with a non-identity layout.
      return failure();
    }
    if (control.rankReductionStrategy ==
        ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {
      FailureOr<Value> rankReducingExtract =
          memref::SubViewOp::rankReduceIfNeeded(rewriter, loc, operand,
                                                targetShape);
      assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
      return *rankReducingExtract;
    }

    assert(
        control.rankReductionStrategy ==
            ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&
        "unknown rank reduction strategy");
    MemRefLayoutAttrInterface layout;
    auto targetType = MemRefType::get(targetShape, memrefType.getElementType(),
                                      layout, memrefType.getMemorySpace());
    return memref::CollapseShapeOp::create(rewriter, loc, targetType, operand,
                                           reassociation)
        .getResult();
  }
  if (auto tensorType = dyn_cast<RankedTensorType>(operand.getType())) {
    if (tensorType.getEncoding() != nullptr) {
      // Do not collapse tensors with an encoding.
      return failure();
    }
    if (control.rankReductionStrategy ==
        ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {
      FailureOr<Value> rankReducingExtract =
          tensor::ExtractSliceOp::rankReduceIfNeeded(rewriter, loc, operand,
                                                     targetShape);
      assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
      return *rankReducingExtract;
    }

    assert(
        control.rankReductionStrategy ==
            ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&
        "unknown rank reduction strategy");
    auto targetType =
        RankedTensorType::get(targetShape, tensorType.getElementType());
    return tensor::CollapseShapeOp::create(rewriter, loc, targetType, operand,
                                           reassociation)
        .getResult();
  }
  llvm_unreachable("unsupported operand type");
}

/// Compute the modified metadata for an operands of operation
/// whose unit dims are being dropped. Return the new indexing map
/// to use, the shape of the operand in the replacement op
/// and the `reassocation` to use to go from original operand shape
/// to modified operand shape.
struct UnitExtentReplacementInfo {
  AffineMap indexMap;
  SmallVector<ReassociationIndices> reassociation;
  SmallVector<int64_t> targetShape;
};
static UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata(
    MLIRContext *context, IndexingMapOpInterface op, OpOperand *opOperand,
    llvm::SmallDenseMap<unsigned, unsigned> &oldDimsToNewDimsMap,
    ArrayRef<AffineExpr> dimReplacements) {
  UnitExtentReplacementInfo info;
  ReassociationIndices reassociationGroup;
  SmallVector<AffineExpr> newIndexExprs;
  AffineMap indexingMap = op.getMatchingIndexingMap(opOperand);
  SmallVector<int64_t> operandShape = op.getStaticOperandShape(opOperand);
  ArrayRef<AffineExpr> exprs = indexingMap.getResults();

  auto isUnitDim = [&](unsigned dim) {
    if (auto dimExpr = dyn_cast<AffineDimExpr>(exprs[dim])) {
      unsigned oldPosition = dimExpr.getPosition();
      return !oldDimsToNewDimsMap.count(oldPosition) &&
             (operandShape[dim] == 1);
    }
    // Handle the other case where the shape is 1, and is accessed using a
    // constant 0.
    if (operandShape[dim] == 1) {
      // Use the new expression after replacing dimensions that will be dropped
      // here to handle cases where an affine expression with multiple
      // dimensions (e.g., `d0 + d2`) can be simplified to 0 after dropping all
      // dimensions used in the expression (`d0` and `d2` in this example).
      AffineExpr newExpr = exprs[dim].replaceDims(dimReplacements);
      auto constAffineExpr = dyn_cast<AffineConstantExpr>(newExpr);
      return constAffineExpr && constAffineExpr.getValue() == 0;
    }
    return false;
  };

  unsigned dim = 0;
  while (dim < operandShape.size() && isUnitDim(dim))
    reassociationGroup.push_back(dim++);
  while (dim < operandShape.size()) {
    assert(!isUnitDim(dim) && "expected non unit-extent");
    reassociationGroup.push_back(dim);
    AffineExpr newExpr = exprs[dim].replaceDims(dimReplacements);
    newIndexExprs.push_back(newExpr);
    info.targetShape.push_back(operandShape[dim]);
    ++dim;
    // Fold all following dimensions that are unit-extent.
    while (dim < operandShape.size() && isUnitDim(dim)) {
      reassociationGroup.push_back(dim++);
    }
    info.reassociation.push_back(reassociationGroup);
    reassociationGroup.clear();
  }
  info.indexMap =
      AffineMap::get(oldDimsToNewDimsMap.size(), indexingMap.getNumSymbols(),
                     newIndexExprs, context);
  return info;
}

FailureOr<DropUnitDimsResult>
linalg::dropUnitDims(RewriterBase &rewriter, IndexingMapOpInterface op,
                     const DroppedUnitDimsBuilder &droppedUnitDimsBuilder,
                     const ControlDropUnitDims &options) {
  auto dpsOp = dyn_cast<DestinationStyleOpInterface>(op.getOperation());
  if (!dpsOp) {
    return rewriter.notifyMatchFailure(
        op, "op should implement DestinationStyleOpInterface");
  }

  SmallVector<AffineMap> indexingMaps = op.getIndexingMapsArray();
  if (indexingMaps.empty())
    return failure();

  // 1. Check if any of the iteration dimensions are unit-trip count. They will
  //    end up being unit-trip count if they are used to index into a unit-dim
  //    tensor/memref.
  AffineMap invertedMap =
      inversePermutation(concatAffineMaps(indexingMaps, rewriter.getContext()));
  if (!invertedMap) {
    return rewriter.notifyMatchFailure(op,
                                       "invalid indexing maps for operation");
  }

  SmallVector<int64_t> allShapesSizes;
  for (OpOperand &opOperand : op->getOpOperands())
    llvm::append_range(allShapesSizes, op.getStaticOperandShape(&opOperand));

  // 1a. Get the allowed list of dimensions to drop from the `options`.
  SmallVector<unsigned> allowedUnitDims = options.controlFn(op);
  if (allowedUnitDims.empty()) {
    return rewriter.notifyMatchFailure(
        op, "control function returns no allowed unit dims to prune");
  }
  llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),
                                               allowedUnitDims.end());
  llvm::SmallDenseSet<unsigned> unitDims;
  for (const auto &expr : enumerate(invertedMap.getResults())) {
    if (AffineDimExpr dimExpr = dyn_cast<AffineDimExpr>(expr.value())) {
      if (allShapesSizes[dimExpr.getPosition()] == 1 &&
          unitDimsFilter.count(expr.index()))
        unitDims.insert(expr.index());
    }
  }

  // 2. Compute the new loops of the modified op by dropping the one-trip
  //    count loops.
  llvm::SmallDenseMap<unsigned, unsigned> oldDimToNewDimMap;
  SmallVector<AffineExpr> dimReplacements;
  unsigned newDims = 0;
  for (auto index : llvm::seq<int64_t>(op.getStaticLoopRanges().size())) {
    if (unitDims.count(index)) {
      dimReplacements.push_back(
          getAffineConstantExpr(0, rewriter.getContext()));
    } else {
      oldDimToNewDimMap[index] = newDims;
      dimReplacements.push_back(
          getAffineDimExpr(newDims, rewriter.getContext()));
      newDims++;
    }
  }

  // 3. For each of the operands, find the
  //    - modified affine map to use.
  //    - shape of the operands after the unit-dims are dropped.
  //    - the reassociation indices used to convert from the original
  //      operand type to modified operand (needed only when using reshapes
  //      for rank reduction strategy)
  // Note that the indexing maps might need changing even if there are no
  // unit dimensions that are dropped to handle cases where `0` is used to
  // access a unit-extent tensor. Consider moving this out of this specific
  // transformation as a stand-alone transformation. Kept here right now due
  // to legacy.
  SmallVector<AffineMap> newIndexingMaps;
  SmallVector<SmallVector<ReassociationIndices>> reassociations;
  SmallVector<SmallVector<int64_t>> targetShapes;
  SmallVector<bool> collapsed;
  for (OpOperand &opOperand : op->getOpOperands()) {
    auto indexingMap = op.getMatchingIndexingMap(&opOperand);
    auto replacementInfo =
        dropUnitExtentFromOperandMetadata(rewriter.getContext(), op, &opOperand,
                                          oldDimToNewDimMap, dimReplacements);
    reassociations.push_back(replacementInfo.reassociation);
    newIndexingMaps.push_back(replacementInfo.indexMap);
    targetShapes.push_back(replacementInfo.targetShape);
    collapsed.push_back(!(replacementInfo.indexMap.getNumResults() ==
                          indexingMap.getNumResults()));
  }

  // Abort if the indexing maps of the result operation are not invertible
  // (i.e. not legal) or if no dimension was reduced.
  if (newIndexingMaps == indexingMaps ||
      !inversePermutation(
          concatAffineMaps(newIndexingMaps, rewriter.getContext())))
    return failure();

  Location loc = op.getLoc();
  // 4. For each of the operands, collapse the operand to convert
  //    from original shape to shape in the modified operation if needed,
  //    either through use of reshapes or rank-reducing slices as
  //    specified in `options`.
  //    Abort if one of the operands cannot be collapsed.
  SmallVector<Value> newOperands;
  for (OpOperand &opOperand : op->getOpOperands()) {
    int64_t idx = opOperand.getOperandNumber();
    if (!collapsed[idx]) {
      newOperands.push_back(opOperand.get());
      continue;
    }
    FailureOr<Value> collapsed =
        options.collapseFn(rewriter, loc, opOperand.get(), targetShapes[idx],
                           reassociations[idx], options);
    if (failed(collapsed)) {
      // Abort if the operand could not be collapsed.
      return failure();
    }
    newOperands.push_back(collapsed.value());
  }

  IndexingMapOpInterface replacementOp = droppedUnitDimsBuilder(
      loc, rewriter, op, newOperands, newIndexingMaps, unitDims);

  // 6. If any result type changes, insert a reshape/slice to convert from the
  //    original type to the new type.
  //    Abort the transformation if the result cannot be expanded back to its
  //    original shape.
  SmallVector<Value> resultReplacements;
  for (auto [index, result] : llvm::enumerate(replacementOp->getResults())) {
    unsigned opOperandIndex = index + dpsOp.getNumDpsInputs();
    Value origDest = dpsOp.getDpsInitOperand(index)->get();
    if (!collapsed[opOperandIndex]) {
      resultReplacements.push_back(result);
      continue;
    }
    FailureOr<Value> expanded =
        options.expandFn(rewriter, loc, result, origDest,
                         reassociations[opOperandIndex], options);
    if (failed(expanded)) {
      // Abort if expansion is not successful.
      return failure();
    }
    resultReplacements.push_back(expanded.value());
  }

  return DropUnitDimsResult{replacementOp, resultReplacements};
}

FailureOr<DropUnitDimsResult>
linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
                     const ControlDropUnitDims &options) {

  DroppedUnitDimsBuilder build =
      [](Location loc, OpBuilder &b, IndexingMapOpInterface op,
         ArrayRef<Value> newOperands, ArrayRef<AffineMap> newIndexingMaps,
         const llvm::SmallDenseSet<unsigned> &droppedDims)
      -> IndexingMapOpInterface {
    auto genericOp = cast<GenericOp>(op);
    // Compute the iterator types of the modified op by dropping the one-trip
    // count loops.
    SmallVector<utils::IteratorType> newIteratorTypes;
    for (auto [index, attr] :
         llvm::enumerate(genericOp.getIteratorTypesArray())) {
      if (!droppedDims.count(index))
        newIteratorTypes.push_back(attr);
    }

    // Create the `linalg.generic` operation with the new operands,
    //    indexing maps, iterator types and result types.
    ArrayRef<Value> newInputs =
        ArrayRef<Value>(newOperands).take_front(genericOp.getNumDpsInputs());
    ArrayRef<Value> newOutputs =
        ArrayRef<Value>(newOperands).take_back(genericOp.getNumDpsInits());
    SmallVector<Type> resultTypes;
    resultTypes.reserve(genericOp.getNumResults());
    for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))
      resultTypes.push_back(newOutputs[i].getType());
    GenericOp replacementOp =
        GenericOp::create(b, loc, resultTypes, newInputs, newOutputs,
                          newIndexingMaps, newIteratorTypes);
    b.cloneRegionBefore(genericOp.getRegion(), replacementOp.getRegion(),
                        replacementOp.getRegion().begin());
    // 5a. Replace `linalg.index` operations that refer to the dropped unit
    //     dimensions.
    IRRewriter rewriter(b);
    replaceUnitDimIndexOps(replacementOp, droppedDims, rewriter);

    return replacementOp;
  };

  return dropUnitDims(rewriter, genericOp, build, options);
}

namespace {
struct DropUnitDims : public OpRewritePattern<GenericOp> {
  DropUnitDims(MLIRContext *context, ControlDropUnitDims options = {},
               PatternBenefit benefit = 1)
      : OpRewritePattern(context, benefit), options(std::move(options)) {}

  LogicalResult matchAndRewrite(GenericOp genericOp,
                                PatternRewriter &rewriter) const override {
    FailureOr<DropUnitDimsResult> result =
        dropUnitDims(rewriter, genericOp, options);
    if (failed(result)) {
      return failure();
    }
    rewriter.replaceOp(genericOp, result->replacements);
    return success();
  }

private:
  ControlDropUnitDims options;
};
} // namespace

//===---------------------------------------------------------------------===//
// Drop dimensions that are unit-extents within tensor operations.
//===---------------------------------------------------------------------===//

namespace {
struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
  DropPadUnitDims(MLIRContext *context, ControlDropUnitDims options = {},
                  PatternBenefit benefit = 1)
      : OpRewritePattern(context, benefit), options(std::move(options)) {}

  LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                PatternRewriter &rewriter) const override {
    // 1a. Get the allowed list of dimensions to drop from the `options`.
    SmallVector<unsigned> allowedUnitDims = options.controlFn(padOp);
    if (allowedUnitDims.empty()) {
      return rewriter.notifyMatchFailure(
          padOp, "control function returns no allowed unit dims to prune");
    }

    if (padOp.getSourceType().getEncoding()) {
      return rewriter.notifyMatchFailure(
          padOp, "cannot collapse dims of tensor with encoding");
    }

    // Fail for non-constant padding values. The body of the pad could
    // depend on the padding indices and/or properties of the padded
    // tensor so for now we fail.
    // TODO: Support non-constant padding values.
    Value paddingVal = padOp.getConstantPaddingValue();
    if (!paddingVal) {
      return rewriter.notifyMatchFailure(
          padOp, "unimplemented: non-constant padding value");
    }

    ArrayRef<int64_t> sourceShape = padOp.getSourceType().getShape();
    ArrayRef<int64_t> resultShape = padOp.getResultType().getShape();
    int64_t padRank = sourceShape.size();

    auto isStaticZero = [](OpFoldResult f) {
      return getConstantIntValue(f) == 0;
    };

    llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),
                                                 allowedUnitDims.end());
    llvm::SmallDenseSet<unsigned> unitDims;
    SmallVector<int64_t> newShape;
    SmallVector<int64_t> newResultShape;
    SmallVector<OpFoldResult> newLowPad;
    SmallVector<OpFoldResult> newHighPad;
    for (const auto [dim, size, outSize, low, high] : zip_equal(
             llvm::seq(static_cast<int64_t>(0), padRank), sourceShape,
             resultShape, padOp.getMixedLowPad(), padOp.getMixedHighPad())) {
      if (unitDimsFilter.contains(dim) && size == 1 && isStaticZero(low) &&
          isStaticZero(high)) {
        unitDims.insert(dim);
      } else {
        newShape.push_back(size);
        newResultShape.push_back(outSize);
        newLowPad.push_back(low);
        newHighPad.push_back(high);
      }
    }

    if (unitDims.empty()) {
      return rewriter.notifyMatchFailure(padOp, "no unit dims to collapse");
    }

    ReassociationIndices reassociationGroup;
    SmallVector<ReassociationIndices> reassociationMap;
    int64_t dim = 0;
    while (dim < padRank && unitDims.contains(dim))
      reassociationGroup.push_back(dim++);
    while (dim < padRank) {
      assert(!unitDims.contains(dim) && "expected non unit-extent");
      reassociationGroup.push_back(dim);
      dim++;
      // Fold all following dimensions that are unit-extent.
      while (dim < padRank && unitDims.contains(dim))
        reassociationGroup.push_back(dim++);
      reassociationMap.push_back(reassociationGroup);
      reassociationGroup.clear();
    }

    FailureOr<Value> collapsedSource =
        options.collapseFn(rewriter, padOp.getLoc(), padOp.getSource(),
                           newShape, reassociationMap, options);
    if (failed(collapsedSource)) {
      return rewriter.notifyMatchFailure(padOp, "Failed to collapse source");
    }

    auto newResultType = RankedTensorType::get(
        newResultShape, padOp.getResultType().getElementType());
    auto newPadOp = tensor::PadOp::create(
        rewriter, padOp.getLoc(), /*result=*/newResultType,
        collapsedSource.value(), newLowPad, newHighPad, paddingVal,
        padOp.getNofold());

    Value dest = padOp.getResult();
    if (options.rankReductionStrategy ==
        ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {
      SmallVector<OpFoldResult> expandedSizes;
      int64_t numUnitDims = 0;
      for (auto dim : llvm::seq(static_cast<int64_t>(0), padRank)) {
        if (unitDims.contains(dim)) {
          expandedSizes.push_back(rewriter.getIndexAttr(1));
          numUnitDims++;
          continue;
        }
        expandedSizes.push_back(tensor::getMixedSize(
            rewriter, padOp.getLoc(), newPadOp, dim - numUnitDims));
      }
      dest = tensor::EmptyOp::create(rewriter, padOp.getLoc(), expandedSizes,
                                     padOp.getResultType().getElementType());
    }

    FailureOr<Value> expandedValue =
        options.expandFn(rewriter, padOp.getLoc(), newPadOp.getResult(), dest,
                         reassociationMap, options);
    if (failed(expandedValue)) {
      return rewriter.notifyMatchFailure(padOp, "Failed to expand result");
    }
    rewriter.replaceOp(padOp, expandedValue.value());
    return success();
  }

private:
  ControlDropUnitDims options;
};
} // namespace

namespace {
/// Convert `extract_slice` operations to rank-reduced versions.
struct RankReducedExtractSliceOp
    : public OpRewritePattern<tensor::ExtractSliceOp> {
  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                PatternRewriter &rewriter) const override {
    RankedTensorType resultType = sliceOp.getType();
    SmallVector<OpFoldResult> targetShape;
    for (auto size : resultType.getShape())
      targetShape.push_back(rewriter.getIndexAttr(size));
    auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);
    if (!reassociation ||
        reassociation->size() == static_cast<size_t>(resultType.getRank()))
      return failure();

    SmallVector<OpFoldResult> offsets = sliceOp.getMixedOffsets();
    SmallVector<OpFoldResult> strides = sliceOp.getMixedStrides();
    SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();
    auto rankReducedType = cast<RankedTensorType>(
        tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
            reassociation->size(), sliceOp.getSourceType(), sizes));

    Location loc = sliceOp.getLoc();
    Value newSlice = tensor::ExtractSliceOp::create(
        rewriter, loc, rankReducedType, sliceOp.getSource(), offsets, sizes,
        strides);
    rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
        sliceOp, resultType, newSlice, *reassociation);
    return success();
  }
};

/// Convert `insert_slice` operations to rank-reduced versions.
/// This patterns works with both InsertSliceOp and ParallelInsertSliceOp.
template <typename InsertOpTy>
struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {
  using OpRewritePattern<InsertOpTy>::OpRewritePattern;

  LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
                                PatternRewriter &rewriter) const override {
    RankedTensorType sourceType = insertSliceOp.getSourceType();
    SmallVector<OpFoldResult> targetShape;
    for (auto size : sourceType.getShape())
      targetShape.push_back(rewriter.getIndexAttr(size));
    auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);
    if (!reassociation ||
        reassociation->size() == static_cast<size_t>(sourceType.getRank()))
      return failure();

    Location loc = insertSliceOp.getLoc();
    tensor::CollapseShapeOp reshapedSource;
    {
      OpBuilder::InsertionGuard g(rewriter);
      // The only difference between InsertSliceOp and ParallelInsertSliceOp
      // is the insertion point is just before the ParallelCombiningOp in the
      // parallel case.
      if (std::is_same<InsertOpTy, tensor::ParallelInsertSliceOp>::value)
        rewriter.setInsertionPoint(insertSliceOp->getParentOp());
      reshapedSource = tensor::CollapseShapeOp::create(
          rewriter, loc, insertSliceOp.getSource(), *reassociation);
    }
    rewriter.replaceOpWithNewOp<InsertOpTy>(
        insertSliceOp, reshapedSource, insertSliceOp.getDest(),
        insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
        insertSliceOp.getMixedStrides());
    return success();
  }
};
} // namespace

/// Patterns that are used to canonicalize the use of unit-extent dims for
/// broadcasting.
void mlir::linalg::populateFoldUnitExtentDimsPatterns(
    RewritePatternSet &patterns, linalg::ControlDropUnitDims &options) {
  auto *context = patterns.getContext();
  patterns.add<DropUnitDims>(context, options);
  patterns.add<DropPadUnitDims>(context, options);
}

void mlir::linalg::populateFoldUnitExtentDimsCanonicalizationPatterns(
    RewritePatternSet &patterns, linalg::ControlDropUnitDims &options) {
  auto *context = patterns.getContext();
  bool reassociativeReshape =
      options.rankReductionStrategy ==
      linalg::ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape;
  if (reassociativeReshape) {
    patterns.add<RankReducedExtractSliceOp,
                 RankReducedInsertSliceOp<tensor::InsertSliceOp>,
                 RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
        context);
    tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
    tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
  }
  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
  tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
  tensor::populateFoldTensorEmptyPatterns(patterns);
  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
  memref::populateResolveShapedTypeResultDimsPatterns(patterns);
}

void mlir::linalg::populateMoveInitOperandsToInputPattern(
    RewritePatternSet &patterns) {
  patterns.add<MoveInitOperandsToInput>(patterns.getContext());
}

namespace {
/// Pass that removes unit-extent dims within generic ops.
struct LinalgFoldUnitExtentDimsPass
    : public impl::LinalgFoldUnitExtentDimsPassBase<
          LinalgFoldUnitExtentDimsPass> {
  using impl::LinalgFoldUnitExtentDimsPassBase<
      LinalgFoldUnitExtentDimsPass>::LinalgFoldUnitExtentDimsPassBase;
  void runOnOperation() override {
    Operation *op = getOperation();
    MLIRContext *context = op->getContext();
    ControlDropUnitDims options;
    if (useRankReducingSlices) {
      options.rankReductionStrategy = linalg::ControlDropUnitDims::
          RankReductionStrategy::ExtractInsertSlice;
    }

    // Apply fold unit extent dims patterns with walk-based driver.
    {
      RewritePatternSet patterns(context);
      linalg::populateFoldUnitExtentDimsPatterns(patterns, options);
      walkAndApplyPatterns(op, std::move(patterns));
    }

    // Apply canonicalization patterns with greedy driver.
    {
      RewritePatternSet patterns(context);
      populateMoveInitOperandsToInputPattern(patterns);
      linalg::populateFoldUnitExtentDimsCanonicalizationPatterns(patterns,
                                                                 options);
      (void)applyPatternsGreedily(op, std::move(patterns));
    }
  }
};

} // namespace

namespace {

/// Returns reassociation indices for collapsing/expanding a
/// tensor of rank `rank` at position `pos`.
static SmallVector<ReassociationIndices>
getReassociationForReshapeAtDim(int64_t rank, int64_t pos) {
  SmallVector<ReassociationIndices> reassociation(rank - 1, {0, 1});
  bool lastDim = pos == rank - 1;
  if (rank > 2) {
    for (int64_t i = 0; i < rank - 1; i++) {
      if (i == pos || (lastDim && i == pos - 1))
        reassociation[i] = ReassociationIndices{i, i + 1};
      else if (i < pos)
        reassociation[i] = ReassociationIndices{i};
      else
        reassociation[i] = ReassociationIndices{i + 1};
    }
  }
  return reassociation;
}

/// Returns a collapsed `val` where the collapsing occurs at dim `pos`.
/// If `pos < 0`, then don't collapse.
static Value collapseSingletonDimAt(PatternRewriter &rewriter, Value val,
                                    int64_t pos) {
  if (pos < 0)
    return val;
  auto valType = cast<ShapedType>(val.getType());
  SmallVector<int64_t> collapsedShape(valType.getShape());
  collapsedShape.erase(collapsedShape.begin() + pos);
  ControlDropUnitDims control{};
  FailureOr<Value> collapsed = control.collapseFn(
      rewriter, val.getLoc(), val, collapsedShape,
      getReassociationForReshapeAtDim(valType.getRank(), pos), control);
  assert(llvm::succeeded(collapsed) && "Collapsing the value failed");
  return collapsed.value();
}

/// Base class for all rank reduction patterns for contraction ops
/// with unit dimensions.  All patterns should convert one named op
/// to another named op.  Intended to reduce only one iteration space dim
/// at a time.
/// Reducing multiple dims will happen with recusive application of
/// pattern rewrites.
template <typename FromOpTy, typename ToOpTy>
struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
  using OpRewritePattern<FromOpTy>::OpRewritePattern;

  /// Collapse all collapsable operands.
  SmallVector<Value>
  collapseOperands(PatternRewriter &rewriter, ArrayRef<Value> operands,
                   ArrayRef<int64_t> operandCollapseDims) const {
    assert(operandCollapseDims.size() == 3 && operands.size() == 3 &&
           "expected 3 operands and dims");
    return llvm::map_to_vector(
        llvm::zip(operands, operandCollapseDims), [&](auto pair) {
          return collapseSingletonDimAt(rewriter, std::get<0>(pair),
                                        std::get<1>(pair));
        });
  }

  /// Expand result tensor.
  Value expandResult(PatternRewriter &rewriter, Value result,
                     RankedTensorType expandedType, int64_t dim) const {
    return tensor::ExpandShapeOp::create(
        rewriter, result.getLoc(), expandedType, result,
        getReassociationForReshapeAtDim(expandedType.getRank(), dim));
  }

  LogicalResult matchAndRewrite(FromOpTy contractionOp,
                                PatternRewriter &rewriter) const override {
    if (contractionOp.hasUserDefinedMaps()) {
      return rewriter.notifyMatchFailure(
          contractionOp, "ops with user-defined maps are not supported");
    }

    auto loc = contractionOp.getLoc();
    auto inputs = contractionOp.getDpsInputs();
    auto inits = contractionOp.getDpsInits();
    if (inputs.size() != 2 || inits.size() != 1)
      return rewriter.notifyMatchFailure(contractionOp,
                                         "expected 2 inputs and 1 init");
    auto lhs = inputs[0];
    auto rhs = inputs[1];
    auto init = inits[0];
    SmallVector<Value> operands{lhs, rhs, init};

    SmallVector<int64_t> operandUnitDims;
    if (failed(getOperandUnitDims(contractionOp, operandUnitDims)))
      return rewriter.notifyMatchFailure(contractionOp,
                                         "no reducable dims found");

    SmallVector<Value> collapsedOperands =
        collapseOperands(rewriter, operands, operandUnitDims);
    Value collapsedLhs = collapsedOperands[0];
    Value collapsedRhs = collapsedOperands[1];
    Value collapsedInit = collapsedOperands[2];
    SmallVector<Type, 1> collapsedResultTy;
    if (isa<RankedTensorType>(collapsedInit.getType()))
      collapsedResultTy.push_back(collapsedInit.getType());
    auto collapsedOp = ToOpTy::create(rewriter, loc, collapsedResultTy,
                                      ValueRange{collapsedLhs, collapsedRhs},
                                      ValueRange{collapsedInit});
    for (auto attr : contractionOp->getAttrs()) {
      if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName ||
          attr.getName() == "indexing_maps")
        continue;
      collapsedOp->setAttr(attr.getName(), attr.getValue());
    }

    auto results = contractionOp.getResults();
    assert(results.size() < 2 && "expected at most one result");
    if (results.empty()) {
      rewriter.replaceOp(contractionOp, collapsedOp);
    } else {
      rewriter.replaceOp(
          contractionOp,
          expandResult(rewriter, collapsedOp.getResultTensors()[0],
                       cast<RankedTensorType>(results[0].getType()),
                       operandUnitDims[2]));
    }

    return success();
  }

  /// Populate `operandUnitDims` with 3 indices indicating the unit dim
  /// for each operand that should be collapsed in this pattern.  If an
  /// operand shouldn't be collapsed, the index should be negative.
  virtual LogicalResult
  getOperandUnitDims(LinalgOp op,
                     SmallVectorImpl<int64_t> &operandUnitDims) const = 0;
};

/// Patterns for unbatching batched contraction ops
template <typename FromOpTy, typename ToOpTy>
struct RankReduceToUnBatched : RankReduceContractionOps<FromOpTy, ToOpTy> {
  using RankReduceContractionOps<FromOpTy, ToOpTy>::RankReduceContractionOps;

  /// Look for unit batch dims to collapse.
  LogicalResult
  getOperandUnitDims(LinalgOp op,
                     SmallVectorImpl<int64_t> &operandUnitDims) const override {
    FailureOr<ContractionDimensions> maybeContractionDims =
        inferContractionDims(op);
    if (failed(maybeContractionDims)) {
      LLVM_DEBUG(llvm::dbgs() << "could not infer contraction dims");
      return failure();
    }
    const ContractionDimensions &contractionDims = maybeContractionDims.value();

    if (contractionDims.batch.size() != 1)
      return failure();
    auto batchDim = contractionDims.batch[0];
    SmallVector<std::pair<Value, unsigned>, 3> bOperands;
    op.mapIterationSpaceDimToAllOperandDims(batchDim, bOperands);
    if (bOperands.size() != 3 || llvm::any_of(bOperands, [](auto pair) {
          return cast<ShapedType>(std::get<0>(pair).getType())
                     .getShape()[std::get<1>(pair)] != 1;
        })) {
      LLVM_DEBUG(llvm::dbgs() << "specified unit dims not found");
      return failure();
    }

    operandUnitDims = SmallVector<int64_t>{std::get<1>(bOperands[0]),
                                           std::get<1>(bOperands[1]),
                                           std::get<1>(bOperands[2])};
    return success();
  }
};

/// Patterns for reducing non-batch dimensions
template <typename FromOpTy, typename ToOpTy>
struct RankReduceMatmul : RankReduceContractionOps<FromOpTy, ToOpTy> {
  using RankReduceContractionOps<FromOpTy, ToOpTy>::RankReduceContractionOps;

  /// Helper for determining whether the lhs/init or rhs/init are reduced.
  static bool constexpr reduceLeft =
      (std::is_same_v<FromOpTy, BatchMatmulOp> &&
       std::is_same_v<ToOpTy, BatchVecmatOp>) ||
      (std::is_same_v<FromOpTy, MatmulOp> &&
       std::is_same_v<ToOpTy, VecmatOp>) ||
      (std::is_same_v<FromOpTy, MatvecOp> && std::is_same_v<ToOpTy, DotOp>);

  /// Look for non-batch spatial dims to collapse.
  LogicalResult
  getOperandUnitDims(LinalgOp op,
                     SmallVectorImpl<int64_t> &operandUnitDims) const override {
    FailureOr<ContractionDimensions> maybeContractionDims =
        inferContractionDims(op);
    if (failed(maybeContractionDims)) {
      LLVM_DEBUG(llvm::dbgs() << "could not infer contraction dims");
      return failure();
    }
    const ContractionDimensions &contractionDims = maybeContractionDims.value();

    if constexpr (reduceLeft) {
      auto m = contractionDims.m[0];
      SmallVector<std::pair<Value, unsigned>, 2> mOperands;
      op.mapIterationSpaceDimToAllOperandDims(m, mOperands);
      if (mOperands.size() != 2)
        return failure();
      if (llvm::all_of(mOperands, [](auto pair) {
            return cast<ShapedType>(std::get<0>(pair).getType())
                       .getShape()[std::get<1>(pair)] == 1;
          })) {
        operandUnitDims = SmallVector<int64_t>{std::get<1>(mOperands[0]), -1,
                                               std::get<1>(mOperands[1])};
        return success();
      }
    } else {
      auto n = contractionDims.n[0];
      SmallVector<std::pair<Value, unsigned>, 2> nOperands;
      op.mapIterationSpaceDimToAllOperandDims(n, nOperands);
      if (nOperands.size() != 2)
        return failure();
      if (llvm::all_of(nOperands, [](auto pair) {
            return cast<ShapedType>(std::get<0>(pair).getType())
                       .getShape()[std::get<1>(pair)] == 1;
          })) {
        operandUnitDims = SmallVector<int64_t>{-1, std::get<1>(nOperands[0]),
                                               std::get<1>(nOperands[1])};
        return success();
      }
    }
    LLVM_DEBUG(llvm::dbgs() << "specified unit dims not found");
    return failure();
  }
};

} // namespace

void mlir::linalg::populateContractionOpRankReducingPatterns(
    RewritePatternSet &patterns) {
  MLIRContext *context = patterns.getContext();
  // Unbatching patterns for unit batch size
  patterns.add<RankReduceToUnBatched<BatchMatmulOp, MatmulOp>>(context);
  patterns.add<RankReduceToUnBatched<BatchMatvecOp, MatvecOp>>(context);
  patterns.add<RankReduceToUnBatched<BatchVecmatOp, VecmatOp>>(context);

  // Non-batch rank 1 reducing patterns
  patterns.add<RankReduceMatmul<MatmulOp, VecmatOp>>(context);
  patterns.add<RankReduceMatmul<MatmulOp, MatvecOp>>(context);
  // Batch rank 1 reducing patterns
  patterns.add<RankReduceMatmul<BatchMatmulOp, BatchVecmatOp>>(context);
  patterns.add<RankReduceMatmul<BatchMatmulOp, BatchMatvecOp>>(context);

  // Non-batch rank 0 reducing patterns
  patterns.add<RankReduceMatmul<MatvecOp, DotOp>>(context);
  patterns.add<RankReduceMatmul<VecmatOp, DotOp>>(context);
}