brintos / llvm-project-archived public Read only

1//===- ElementwiseOpFusion.cpp - Implementation of linalg Fusion ---------===///2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements the linalg dialect Fusion on tensors operations pass.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Linalg/Passes.h"14 15#include "mlir/Dialect/Affine/IR/AffineOps.h"16#include "mlir/Dialect/Arith/IR/Arith.h"17#include "mlir/Dialect/Arith/Utils/Utils.h"18#include "mlir/Dialect/Linalg/IR/Linalg.h"19#include "mlir/Dialect/Linalg/Transforms/Transforms.h"20#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"21#include "mlir/Dialect/Tensor/Transforms/Transforms.h"22#include "mlir/IR/AffineExpr.h"23#include "mlir/IR/AffineMap.h"24#include "mlir/IR/Matchers.h"25#include "mlir/IR/PatternMatch.h"26#include "mlir/Support/LLVM.h"27#include "mlir/Transforms/GreedyPatternRewriteDriver.h"28#include "mlir/Transforms/RegionUtils.h"29#include <optional>30#include <utility>31 32namespace mlir {33#define GEN_PASS_DEF_LINALGELEMENTWISEOPFUSIONPASS34#include "mlir/Dialect/Linalg/Passes.h.inc"35} // namespace mlir36 37using namespace mlir;38using namespace mlir::linalg;39 40//===---------------------------------------------------------------------===//41// Methods and patterns that fuse elementwise `linalg.generic` operations.42//===---------------------------------------------------------------------===//43 44/// Append to `fusedOpIndexingMapAttrs` the indexing maps for the operands of45/// the `producer` to use in the fused operation given the indexing map of the46/// result of the producer in the consumer.47static AffineMap getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(48    OpOperand *producerOpOperand, AffineMap producerResultIndexMap,49    AffineMap fusedConsumerArgIndexMap) {50  // The indexing map in the consumer op (fusedConsumerArgIndexMap) is a map51  // from consumer loop -> consumer arg tensor index/producer result tensor52  // index. The fused loop is same as the consumer loop. For each producer arg53  // the indexing map to be computed is a map from consumer loop -> producer54  // arg tensor index.55  // producerResultIndexMap is a map from producer loop -> tensor index.56  // Compute the inverse to get map from tensor index -> producer loop.57  // The inverse is a map from producer result tensor index -> producer loop.58  AffineMap invProducerResultIndexMap =59      inversePermutation(producerResultIndexMap);60  assert(invProducerResultIndexMap &&61         "expected producer result indexing map to be invertible");62 63  LinalgOp producer = cast<LinalgOp>(producerOpOperand->getOwner());64  // argMap is a map from producer loop -> producer arg tensor index.65  AffineMap argMap = producer.getMatchingIndexingMap(producerOpOperand);66 67  // Compose argMap with invProducerResultIndexMap to get a map from68  // producer result tensor index -> producer arg tensor index.69  AffineMap t1 = argMap.compose(invProducerResultIndexMap);70 71  // Compose t1 with fusedConsumerArgIndexMap gives an indexing map from72  // consumer loop/ fused loop -> producer arg tensor index.73  return t1.compose(fusedConsumerArgIndexMap);74}75 76// Checks if the given operand can be dropped, and the remaining operands77// of the fused producer & consumer after the fusion can still compute the78// bounds of the op.79static bool isOpOperandCanBeDroppedAfterFusedLinalgs(80    GenericOp producer, GenericOp consumer,81    ArrayRef<OpOperand *> opOperandsToIgnore) {82  SmallVector<AffineMap> indexingMaps;83 84  SmallVector<GenericOp> ops = {producer, consumer};85  for (auto &op : ops) {86    for (auto &opOperand : op->getOpOperands()) {87      if (llvm::is_contained(opOperandsToIgnore, &opOperand)) {88        continue;89      }90      indexingMaps.push_back(op.getMatchingIndexingMap(&opOperand));91    }92  }93  if (indexingMaps.empty()) {94    // If there are no indexing maps, the operand can only be dropped95    // if neither op has loops.96    return producer.getNumLoops() == 0 && consumer.getNumLoops() == 0;97  }98 99  // The concatanation of the remained indexing maps must be invertible, so100  // the bounds of the op can be still computed after dropping the selected101  // operand. inversePermutation returns an empty AffineMap in case the102  // concatanated indexing maps are not invertible.103  return inversePermutation(concatAffineMaps(104             indexingMaps, producer.getContext())) != AffineMap();105}106 107/// Returns a set of indices of the producer's results which would108/// be preserved after the fusion.109/// * There is a chance that the implementation of the transformation does not110/// agree with the result of this method. This function gives a prediction based111/// on an optimized fusion.112llvm::SmallDenseSet<int> mlir::linalg::getPreservedProducerResults(113    GenericOp producer, GenericOp consumer, OpOperand *fusedOperand) {114  llvm::SmallDenseSet<int> preservedProducerResults;115  llvm::SmallVector<OpOperand *> opOperandsToIgnore;116 117  // The fusedOperand will be removed during the fusion118  opOperandsToIgnore.emplace_back(fusedOperand);119 120  for (const auto &producerResult : llvm::enumerate(producer->getResults())) {121    auto *outputOperand = producer.getDpsInitOperand(producerResult.index());122    opOperandsToIgnore.emplace_back(outputOperand);123    if (producer.payloadUsesValueFromOperand(outputOperand) ||124        !isOpOperandCanBeDroppedAfterFusedLinalgs(producer, consumer,125                                                  opOperandsToIgnore) ||126        llvm::any_of(producerResult.value().getUsers(), [&](Operation *user) {127          return user != consumer.getOperation();128        })) {129      preservedProducerResults.insert(producerResult.index());130 131      // In case the operand can't be dropped132      (void)opOperandsToIgnore.pop_back_val();133    }134  }135  return preservedProducerResults;136}137 138/// Conditions for elementwise fusion of generic operations.139bool mlir::linalg::areElementwiseOpsFusable(OpOperand *fusedOperand) {140  if (!fusedOperand)141    return false;142 143  auto producer = fusedOperand->get().getDefiningOp<GenericOp>();144  auto consumer = dyn_cast<GenericOp>(fusedOperand->getOwner());145 146  // Check producer and consumer are generic ops.147  if (!producer || !consumer)148    return false;149 150  // Consumer can have mixed semantics, just check operand itself has tensor151  // type. Producer must have full tensor semantics to avoid potential152  // aliasing between producer and consumer memrefs.153  if (!producer.hasPureTensorSemantics() ||154      !isa<RankedTensorType>(fusedOperand->get().getType()))155    return false;156 157  // Verify that158  // - the producer has all "parallel" iterator type.159  if (producer.getNumParallelLoops() != producer.getNumLoops())160    return false;161 162  // Only allow fusing the producer of an input operand for now.163  // TODO: allow fusing the producer of an output operand.164  if (!consumer.isDpsInput(fusedOperand))165    return false;166 167  // Get the consumer index map. The number of results of the consumer index168  // map must match the number of loops of the producer.169  AffineMap consumerIndexMap = consumer.getMatchingIndexingMap(fusedOperand);170  if (consumerIndexMap.getNumResults() != producer.getNumLoops())171    return false;172 173  // Finally the index_map for the result must be invertible. For now just174  // verify it is a permutation.175  auto producerResult = cast<OpResult>(fusedOperand->get());176  AffineMap producerResultIndexMap =177      producer.getIndexingMapMatchingResult(producerResult);178  if (!producerResultIndexMap.isPermutation())179    return false;180 181  // Ensure that the fusion does not remove size information required to182  // get the loop bounds. For non-reduction generics, this is trivially the183  // case due to the output operand. For reductions, we need to check that after184  // the fusion, each loop dimension has at least one input that defines it.185  if ((consumer.getNumReductionLoops())) {186    BitVector coveredDims(consumer.getNumLoops(), false);187 188    auto addToCoveredDims = [&](AffineMap map) {189      for (auto result : map.getResults())190        if (auto dimExpr = dyn_cast<AffineDimExpr>(result))191          coveredDims[dimExpr.getPosition()] = true;192    };193 194    for (auto pair :195         llvm::zip(consumer->getOperands(), consumer.getIndexingMapsArray())) {196      Value operand = std::get<0>(pair);197      if (operand == fusedOperand->get())198        continue;199      AffineMap operandMap = std::get<1>(pair);200      addToCoveredDims(operandMap);201    }202 203    for (OpOperand *operand : producer.getDpsInputOperands()) {204      AffineMap newIndexingMap =205          getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(206              operand, producerResultIndexMap, consumerIndexMap);207      addToCoveredDims(newIndexingMap);208    }209    if (!coveredDims.all())210      return false;211  }212 213  return true;214}215 216/// Generate the region of the fused tensor operation. The region of the fused217/// op must be empty.218static void generateFusedElementwiseOpRegion(219    RewriterBase &rewriter, GenericOp fusedOp,220    AffineMap consumerToProducerLoopsMap, OpOperand *fusedOperand,221    unsigned nloops, llvm::SmallDenseSet<int> &preservedProducerResults) {222  auto producer = cast<GenericOp>(fusedOperand->get().getDefiningOp());223  auto consumer = cast<GenericOp>(fusedOperand->getOwner());224  // Build the region of the fused op.225  Block &producerBlock = producer->getRegion(0).front();226  Block &consumerBlock = consumer->getRegion(0).front();227  OpBuilder::InsertionGuard guard(rewriter);228  Block *fusedBlock = rewriter.createBlock(&fusedOp.getRegion());229  IRMapping mapper;230 231  // 2. Add an index operation for every fused loop dimension and use the232  // `consumerToProducerLoopsMap` to map the producer indices.233  if (producer.hasIndexSemantics()) {234    // Add an index operation for every fused loop dimension.235    unsigned numFusedOpLoops = fusedOp.getNumLoops();236    SmallVector<Value> fusedIndices;237    fusedIndices.reserve(numFusedOpLoops);238    llvm::transform(llvm::seq<uint64_t>(0, numFusedOpLoops),239                    std::back_inserter(fusedIndices), [&](uint64_t dim) {240                      return IndexOp::create(rewriter, producer.getLoc(), dim);241                    });242    for (IndexOp indexOp :243         llvm::make_early_inc_range(producerBlock.getOps<IndexOp>())) {244      Value newIndex = affine::AffineApplyOp::create(245          rewriter, producer.getLoc(),246          consumerToProducerLoopsMap.getSubMap(indexOp.getDim()), fusedIndices);247      mapper.map(indexOp.getResult(), newIndex);248    }249  }250  // TODO: allow fusing the producer of an output operand.251  assert(consumer.isDpsInput(fusedOperand) &&252         "expected producer of input operand");253  // 3. Consumer input operands up to consumerIdx (exclusive).254  for (BlockArgument bbArg : consumerBlock.getArguments().take_front(255           fusedOperand->getOperandNumber())) // input assumption.256    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));257 258  // Replacing consumerIdx requires getting the cloned, yielded, value from259  // the (cloned) producer block. This happens in step 9.260 261  // 4. Splice in producer's input operands.262  for (BlockArgument bbArg :263       producerBlock.getArguments().take_front(producer.getNumDpsInputs()))264    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));265 266  // 5. Remaining consumer's input operands (drop past index `consumerIdx`).267  for (BlockArgument bbArg :268       consumerBlock.getArguments()269           .take_front(consumer.getNumDpsInputs())270           .drop_front(fusedOperand->getOperandNumber() + 1))271    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));272 273  // 6. All of the producer's output operands274  for (const auto &bbArg : llvm::enumerate(275           producerBlock.getArguments().take_back(producer.getNumDpsInits()))) {276    if (!preservedProducerResults.count(bbArg.index()))277      continue;278    mapper.map(bbArg.value(), fusedBlock->addArgument(bbArg.value().getType(),279                                                      bbArg.value().getLoc()));280  }281 282  // 7. All of consumer's output operands.283  for (BlockArgument bbArg :284       consumerBlock.getArguments().take_back(consumer.getNumDpsInits()))285    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));286 287  // 8. Clone all producer operations except for the yield and index operations288  // to the fused operation.289  for (auto &op : producerBlock.without_terminator()) {290    if (!isa<IndexOp>(op))291      rewriter.clone(op, mapper);292  }293  // 9. Now we can map the consumerBlock's `consumerIdx` block argument. Just294  // forward the yield operand.295  auto producerYieldOp = cast<linalg::YieldOp>(producerBlock.getTerminator());296  unsigned producerResultNumber =297      cast<OpResult>(fusedOperand->get()).getResultNumber();298  Value replacement =299      mapper.lookupOrDefault(producerYieldOp.getOperand(producerResultNumber));300 301  // Sanity checks, if replacement is not already in the mapper then it must be302  // produced outside.303  if (replacement == producerYieldOp.getOperand(producerResultNumber)) {304    if (auto bb = dyn_cast<BlockArgument>(replacement))305      assert(bb.getOwner() != &producerBlock &&306             "yielded block argument must have been mapped");307    else308      assert(!producer->isAncestor(replacement.getDefiningOp()) &&309             "yielded value must have been mapped");310  }311  mapper.map(consumerBlock.getArgument(fusedOperand->getOperandNumber()),312             replacement);313  // 10. Clone operations from the consumer to the fused op.314  for (auto &op : consumerBlock.without_terminator())315    rewriter.clone(op, mapper);316 317  // 11. Include the final yield (which is the remapped values for all the318  // yield)319  auto consumerYieldOp = cast<linalg::YieldOp>(consumerBlock.getTerminator());320  SmallVector<Value> fusedYieldValues;321  fusedYieldValues.reserve(producerYieldOp.getNumOperands() +322                           consumerYieldOp.getNumOperands());323  for (const auto &producerYieldVal :324       llvm::enumerate(producerYieldOp.getOperands())) {325    if (preservedProducerResults.count(producerYieldVal.index()))326      fusedYieldValues.push_back(327          mapper.lookupOrDefault(producerYieldVal.value()));328  }329  for (auto consumerYieldVal : consumerYieldOp.getOperands())330    fusedYieldValues.push_back(mapper.lookupOrDefault(consumerYieldVal));331  YieldOp::create(rewriter, fusedOp.getLoc(), fusedYieldValues);332 333  // Sanity checks.334  assert(fusedBlock->getNumArguments() == fusedOp.getNumOperands() &&335         "Ill-formed GenericOp region");336}337 338FailureOr<mlir::linalg::ElementwiseOpFusionResult>339mlir::linalg::fuseElementwiseOps(RewriterBase &rewriter,340                                 OpOperand *fusedOperand) {341  assert(areElementwiseOpsFusable(fusedOperand) &&342         "expected elementwise operation pre-conditions to pass");343  auto producerResult = cast<OpResult>(fusedOperand->get());344  auto producer = cast<GenericOp>(producerResult.getOwner());345  auto consumer = cast<GenericOp>(fusedOperand->getOwner());346  // TODO: allow fusing the producer of an output operand.347  assert(consumer.isDpsInput(fusedOperand) &&348         "expected producer of input operand");349  /// Find the results of the producer that have uses outside of the consumer,350  /// after the fusion.351  llvm::SmallDenseSet<int> preservedProducerResults =352      mlir::linalg::getPreservedProducerResults(producer, consumer,353                                                fusedOperand);354 355  // Compute the fused operands list and indexing maps.356  SmallVector<Value> fusedInputOperands, fusedOutputOperands;357  SmallVector<Type> fusedResultTypes;358  SmallVector<AffineMap> fusedIndexMaps;359  fusedInputOperands.reserve(producer.getNumDpsInputs() +360                             consumer.getNumDpsInputs());361  fusedOutputOperands.reserve(preservedProducerResults.size() +362                              consumer.getNumDpsInits());363  fusedResultTypes.reserve(preservedProducerResults.size() +364                           consumer.getNumDpsInits());365  fusedIndexMaps.reserve(producer->getNumOperands() +366                         consumer->getNumOperands());367  // In the following, numbering matches that of `generateFusedTensorOpRegion`.368  // 3. Consumer input operands/maps up to consumerIdx (exclusive).369  auto consumerInputs = consumer.getDpsInputOperands();370  auto *it = llvm::find_if(consumerInputs, [&](OpOperand *operand) {371    return operand == fusedOperand;372  });373  assert(it != consumerInputs.end() && "expected to find the consumer operand");374  for (OpOperand *opOperand : llvm::make_range(consumerInputs.begin(), it)) {375    fusedInputOperands.push_back(opOperand->get());376    fusedIndexMaps.push_back(consumer.getMatchingIndexingMap(opOperand));377  }378  // 4. Splice in producer's input operands/maps.379  AffineMap producerResultIndexMap =380      producer.getIndexingMapMatchingResult(producerResult);381  for (OpOperand *opOperand : producer.getDpsInputOperands()) {382    fusedInputOperands.push_back(opOperand->get());383    // Compute indexing maps for the producer args in the fused operation.384    AffineMap map = getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(385        opOperand, producerResultIndexMap,386        consumer.getMatchingIndexingMap(fusedOperand));387    fusedIndexMaps.push_back(map);388  }389  // 5. Remaining consumer's input operands/maps (drop past index390  // `consumerIdx`).391  for (OpOperand *opOperand :392       llvm::make_range(std::next(it), consumerInputs.end())) {393    fusedInputOperands.push_back(opOperand->get());394    fusedIndexMaps.push_back(consumer.getMatchingIndexingMap(opOperand));395  }396 397  // 6. Collect all of the producer outputs.398  for (const auto &opOperand : llvm::enumerate(producer.getDpsInitsMutable())) {399    if (!preservedProducerResults.count(opOperand.index()))400      continue;401 402    fusedOutputOperands.push_back(opOperand.value().get());403    AffineMap map = getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(404        &opOperand.value(), producerResultIndexMap,405        consumer.getMatchingIndexingMap(fusedOperand));406    fusedIndexMaps.push_back(map);407    fusedResultTypes.push_back(opOperand.value().get().getType());408  }409 410  // 7. All of consumer's output operands (skip operands: added by the builder).411  for (OpOperand &opOperand : consumer.getDpsInitsMutable()) {412    fusedOutputOperands.push_back(opOperand.get());413    fusedIndexMaps.push_back(consumer.getMatchingIndexingMap(&opOperand));414    Type resultType = opOperand.get().getType();415    if (!isa<MemRefType>(resultType))416      fusedResultTypes.push_back(resultType);417  }418 419  // Generate the fused op.420  auto fusedOp = GenericOp::create(421      rewriter, consumer.getLoc(), fusedResultTypes, fusedInputOperands,422      fusedOutputOperands, rewriter.getAffineMapArrayAttr(fusedIndexMaps),423      consumer.getIteratorTypes(),424      /*doc=*/nullptr,425      /*library_call=*/nullptr);426  if (!fusedOp.getShapesToLoopsMap()) {427    // Fused op has invalid indexing maps. Typically this means something is off428    // in the input, but going ahead here would result in verification errors.429    // So cleanup and abort.430    rewriter.eraseOp(fusedOp);431    return rewriter.notifyMatchFailure(432        fusedOp, "fused op failed loop bound computation check");433  }434 435  // Construct an AffineMap from consumer loops to producer loops.436  // consumer loop -> tensor index437  AffineMap consumerResultIndexMap =438      consumer.getMatchingIndexingMap(fusedOperand);439  // tensor index -> producer loop440  AffineMap invProducerResultIndexMap =441      inversePermutation(producerResultIndexMap);442  assert(invProducerResultIndexMap &&443         "expected producer result indexig map to be invertible");444  // consumer loop -> producer loop445  AffineMap consumerToProducerLoopsMap =446      invProducerResultIndexMap.compose(consumerResultIndexMap);447 448  generateFusedElementwiseOpRegion(449      rewriter, fusedOp, consumerToProducerLoopsMap, fusedOperand,450      consumer.getNumLoops(), preservedProducerResults);451  ElementwiseOpFusionResult result;452  result.fusedOp = fusedOp;453  int resultNum = 0;454  for (auto [index, producerResult] : llvm::enumerate(producer->getResults()))455    if (preservedProducerResults.count(index))456      result.replacements[producerResult] = fusedOp->getResult(resultNum++);457  for (auto consumerResult : consumer->getResults())458    result.replacements[consumerResult] = fusedOp->getResult(resultNum++);459  return result;460}461 462namespace {463/// Patterns to fuse a generic op, with the producer of its operands.464class FuseElementwiseOps : public OpRewritePattern<GenericOp> {465public:466  FuseElementwiseOps(MLIRContext *context, ControlFusionFn fun,467                     PatternBenefit benefit = 1)468      : OpRewritePattern<GenericOp>(context, benefit),469        controlFn(std::move(fun)) {}470 471  LogicalResult matchAndRewrite(GenericOp genericOp,472                                PatternRewriter &rewriter) const override {473    // Find the first operand that is defined by another generic op on tensors.474    for (OpOperand &opOperand : genericOp->getOpOperands()) {475      if (!areElementwiseOpsFusable(&opOperand))476        continue;477      if (!controlFn(&opOperand))478        continue;479 480      Operation *producer = opOperand.get().getDefiningOp();481 482      // Find the producer of the operand.483      FailureOr<ElementwiseOpFusionResult> fusionResult =484          fuseElementwiseOps(rewriter, &opOperand);485      if (failed(fusionResult))486        return rewriter.notifyMatchFailure(genericOp, "fusion failed");487 488      // Perform the fusion.489      for (auto [origVal, replacement] : fusionResult->replacements) {490        rewriter.replaceUsesWithIf(origVal, replacement, [&](OpOperand &use) {491          // Only replace consumer uses.492          return use.get().getDefiningOp() != producer;493        });494      }495      rewriter.eraseOp(genericOp);496      return success();497    }498    return failure();499  }500 501private:502  ControlFusionFn controlFn;503};504} // namespace505 506//===---------------------------------------------------------------------===//507// Methods and patterns that fuse reshape ops with elementwise operations by508// expanding the dimensionality of the elementwise operations.509//===---------------------------------------------------------------------===//510 511/// Conditions for folding a structured linalg operation with a reshape op by512/// expanding the iteration space dimensionality for tensor operations. These513/// are preconditions assumed by `foldReshapeByDimExpansion` which implements514/// the following fusion pattern.515///516///  Consider517///518///  %c = linalg.generic ins(%a, %b : memref<?x?x?xf32>, memref<?x?xf32>)519///         indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,520///                          affine_map<(d0, d1, d2) -> (d1, d2)>,521///                          affine_map<(d0, d1, d2) -> (d0, d2, d1)>]522///  %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]]523///       : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>524///525///  The reshape can be folded into the `linalgOp` if its loop dimensionality526///  is increased to match the result (operand) of the tensor.expand_shape.527///  The indexing_map of the fused tensor in the `linalgOp` and the528///  reassociation map helps compute the indexing maps of the modified op.529///  For the above example, based on the reassociation map it530///  can be concluded that531///532///  - The loop used to access the first dimension of the fused tensor is split533///    into two.534///  - The loop used to access the second dimension of the fused tensor is kept535///    as is.536///  - The loop used to access the third dimension of the fused tensor is split537///    into three.538///539///  i.e. (e0, e1, e2, e3, e4) is the domain of the indexing map of the modified540///  op, then541///542///   d0 -> e0, e1543///   d1 -> e2, e3, e4544///   d2 -> e5545///546///  substituting this, the structured op can be rewritten as547///548///  %d = linalg.generic ins(%0, %1 : )549///        indexing_maps =550///         [affine_map<(e0, e1, e2, e3, e4, e5) -> (e2, e3, e4, e0, e1, e5)>,551///          affine_map<(e0, e1, e2, e3, e4, e5) -> (e2, e3, e4, e5)>,552///          affine_map<(e0, e1, e2, e3, e4, e5) -> (e0, e1, e5, e2, e3, e4)>]553///554///  Since operands to the linalg generic are now 5D, reshapes can be introduced555///  to make it consistent556///557///  %0 = tensor.expand_shape %a [[0, 1, 2], [3, 4], [5]]558///       : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>559///  %1 = tensor.expand_shape %b [[0, 1, 2], [3]]560///       : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>561///562///  The added reshapes are again expanding patterns, so they will get fused563///  with its producers if possible.564static bool isFusableWithReshapeByDimExpansion(LinalgOp linalgOp,565                                               OpOperand *fusableOpOperand) {566  // Is fusable only if:567  // - All the indexing maps for operands and results are projected568  //   permutations.569  // - The fused tensor is not a scalar.570  SmallVector<utils::IteratorType> iteratorTypes =571      linalgOp.getIteratorTypesArray();572  AffineMap operandMap = linalgOp.getMatchingIndexingMap(fusableOpOperand);573  return linalgOp.hasPureTensorSemantics() &&574         llvm::all_of(linalgOp.getIndexingMaps().getValue(),575                      [](Attribute attr) {576                        return cast<AffineMapAttr>(attr)577                            .getValue()578                            .isProjectedPermutation();579                      }) &&580         operandMap.getNumResults() > 0;581}582 583namespace {584/// Information needed to expand a generic operation to fold the reshape with585/// it.586class ExpansionInfo {587public:588  // Computes the mapping from original dimensions of the op to the dimensions589  // of the expanded op given the `indexingMap` of the fused operand/result of590  // the generic op, the `reassocationMaps` of the reshape op and the shape of591  // the expanded op.592  LogicalResult compute(LinalgOp linalgOp, OpOperand *fusableOpOperand,593                        ArrayRef<AffineMap> reassociationMaps,594                        ArrayRef<OpFoldResult> expandedShape,595                        PatternRewriter &rewriter);596  unsigned getOrigOpNumDims() const { return reassociation.size(); }597  unsigned getExpandedOpNumDims() const { return expandedOpNumDims; }598  ReassociationIndicesRef getExpandedDims(unsigned i) const {599    return reassociation[i];600  }601  ArrayRef<OpFoldResult> getExpandedShapeOfDim(unsigned i) const {602    return expandedShapeMap[i];603  }604  ArrayRef<OpFoldResult> getOriginalShape() const { return originalLoopExtent; }605 606private:607  /// Reassociation from the dimensions in the original operation to the608  /// dimension of the expanded operation.609  SmallVector<ReassociationIndices> reassociation;610  /// Mapping from extent of loops in the original operation, to the extent of611  /// loops in the expanded operation.612  SmallVector<SmallVector<OpFoldResult>> expandedShapeMap;613  /// Extent of the loop in the original operation.614  SmallVector<OpFoldResult> originalLoopExtent;615  unsigned expandedOpNumDims;616};617} // namespace618 619LogicalResult ExpansionInfo::compute(LinalgOp linalgOp,620                                     OpOperand *fusableOpOperand,621                                     ArrayRef<AffineMap> reassociationMaps,622                                     ArrayRef<OpFoldResult> expandedShape,623                                     PatternRewriter &rewriter) {624  if (reassociationMaps.empty())625    return failure();626  AffineMap fusedIndexMap = linalgOp.getMatchingIndexingMap(fusableOpOperand);627 628  OpBuilder::InsertionGuard g(rewriter);629  rewriter.setInsertionPoint(linalgOp);630  originalLoopExtent = llvm::map_to_vector(631      linalgOp.createLoopRanges(rewriter, linalgOp->getLoc()),632      [](Range r) { return r.size; });633 634  reassociation.clear();635  expandedShapeMap.clear();636  // Compute the number of dimension in the expanded op that correspond to each637  // dimension of the original op.638  SmallVector<unsigned> numExpandedDims(fusedIndexMap.getNumDims(), 1);639  expandedShapeMap.resize(fusedIndexMap.getNumDims());640  for (const auto &resultExpr : llvm::enumerate(fusedIndexMap.getResults())) {641    unsigned pos = cast<AffineDimExpr>(resultExpr.value()).getPosition();642    AffineMap foldedDims = reassociationMaps[resultExpr.index()];643    numExpandedDims[pos] = foldedDims.getNumResults();644    ArrayRef<OpFoldResult> shape =645        expandedShape.slice(foldedDims.getDimPosition(0), numExpandedDims[pos]);646    expandedShapeMap[pos].assign(shape.begin(), shape.end());647  }648  // The remaining dimensions remain the same.649  for (unsigned i : llvm::seq<unsigned>(0, fusedIndexMap.getNumDims()))650    if (expandedShapeMap[i].empty())651      expandedShapeMap[i] = {originalLoopExtent[i]};652 653  // Compute reassociation map from the original op to the expanded op.654  unsigned sum = 0;655  reassociation.reserve(fusedIndexMap.getNumDims());656  for (const auto &numFoldedDim : llvm::enumerate(numExpandedDims)) {657    auto seq = llvm::seq<int64_t>(sum, sum + numFoldedDim.value());658    reassociation.emplace_back(seq.begin(), seq.end());659    sum += numFoldedDim.value();660  }661  expandedOpNumDims = sum;662  return success();663}664 665/// Return the indexing map to use in the expanded op for a given the666/// `indexingMap` of the original operation.667static AffineMap668getIndexingMapInExpandedOp(OpBuilder &builder, AffineMap indexingMap,669                           const ExpansionInfo &expansionInfo) {670  SmallVector<AffineExpr> newExprs;671  for (AffineExpr expr : indexingMap.getResults()) {672    unsigned pos = cast<AffineDimExpr>(expr).getPosition();673    SmallVector<AffineExpr, 4> expandedExprs = llvm::to_vector<4>(674        llvm::map_range(expansionInfo.getExpandedDims(pos), [&](int64_t v) {675          return builder.getAffineDimExpr(static_cast<unsigned>(v));676        }));677    newExprs.append(expandedExprs.begin(), expandedExprs.end());678  }679  return AffineMap::get(expansionInfo.getExpandedOpNumDims(),680                        indexingMap.getNumSymbols(), newExprs,681                        builder.getContext());682}683 684/// Return the shape and type of the operand/result to use in the expanded op685/// given the type in the original op.686static std::tuple<SmallVector<OpFoldResult>, RankedTensorType>687getExpandedShapeAndType(RankedTensorType originalType, AffineMap indexingMap,688                        const ExpansionInfo &expansionInfo) {689  SmallVector<OpFoldResult> expandedShape;690  for (AffineExpr expr : indexingMap.getResults()) {691    unsigned dim = cast<AffineDimExpr>(expr).getPosition();692    ArrayRef<OpFoldResult> dimExpansion =693        expansionInfo.getExpandedShapeOfDim(dim);694    expandedShape.append(dimExpansion.begin(), dimExpansion.end());695  }696  SmallVector<int64_t> expandedStaticShape;697  std::tie(expandedStaticShape, std::ignore) =698      decomposeMixedValues(expandedShape);699  return {expandedShape, RankedTensorType::get(expandedStaticShape,700                                               originalType.getElementType())};701}702 703/// Returns the reassociation maps to use in the `tensor.expand_shape`704/// operation to convert the operands of the original operation to operands of705/// the expanded operation. The same method is used to compute the706/// `tensor.collapse_shape` used to collapse the result of the expanded707/// op to get the value that can replace all uses of the results of the original708/// op.709static SmallVector<ReassociationIndices>710getReassociationForExpansion(AffineMap indexingMap,711                             const ExpansionInfo &expansionInfo) {712  SmallVector<ReassociationIndices> reassociation;713  unsigned numReshapeDims = 0;714  for (AffineExpr expr : indexingMap.getResults()) {715    unsigned dim = cast<AffineDimExpr>(expr).getPosition();716    auto numExpandedDims = expansionInfo.getExpandedDims(dim).size();717    SmallVector<int64_t, 2> indices = llvm::to_vector<2>(718        llvm::seq<int64_t>(numReshapeDims, numReshapeDims + numExpandedDims));719    reassociation.emplace_back(std::move(indices));720    numReshapeDims += numExpandedDims;721  }722  return reassociation;723}724 725/// Update the body of an expanded linalg operation having index semantics. The726/// indices of the original operation need to be recovered by linearizing the727/// indices of the correspoding dimensions of the expanded operation. For now it728/// is assumed that the shapes of the expanded operation needed for729/// linearization are static.730static void updateExpandedGenericOpRegion(PatternRewriter &rewriter,731                                          Location loc, Region &fusedRegion,732                                          const ExpansionInfo &expansionInfo) {733  // Replace the original indices by the linearization of the expanded indices.734  for (IndexOp indexOp :735       llvm::make_early_inc_range(fusedRegion.front().getOps<IndexOp>())) {736    ArrayRef<int64_t> expandedDims =737        expansionInfo.getExpandedDims(indexOp.getDim());738    assert(!expandedDims.empty() && "expected valid expansion info");739 740    // Skip index operations that are not affected by the expansion.741    if (expandedDims.size() == 1 &&742        expandedDims.front() == (int64_t)indexOp.getDim())743      continue;744 745    // Linearize the expanded indices of the original index dimension.746    OpBuilder::InsertionGuard guard(rewriter);747    rewriter.setInsertionPointAfter(indexOp);748    ArrayRef<OpFoldResult> expandedDimsShape =749        expansionInfo.getExpandedShapeOfDim(indexOp.getDim()).drop_front();750    SmallVector<Value> expandedIndices;751    expandedIndices.reserve(expandedDims.size() - 1);752    llvm::transform(753        expandedDims.drop_front(), std::back_inserter(expandedIndices),754        [&](int64_t dim) { return IndexOp::create(rewriter, loc, dim); });755    OpFoldResult newIndex =756        IndexOp::create(rewriter, loc, expandedDims.front()).getResult();757    for (auto [expandedShape, expandedIndex] :758         llvm::zip(expandedDimsShape, expandedIndices)) {759      AffineExpr idx, acc, shape;760      bindDims(rewriter.getContext(), idx, acc);761      bindSymbols(rewriter.getContext(), shape);762      newIndex = affine::makeComposedFoldedAffineApply(763          rewriter, indexOp.getLoc(), idx + acc * shape,764          ArrayRef<OpFoldResult>{expandedIndex, newIndex, expandedShape});765    }766    Value newIndexVal =767        getValueOrCreateConstantIndexOp(rewriter, indexOp.getLoc(), newIndex);768    rewriter.replaceOp(indexOp, newIndexVal);769  }770}771 772// Create an expanded transpose op.773// the reassociation map is already permuted hence we inverse permute and then774// flatten it. Then we inverse permute it again to get the final expanded775// transpose permutation. For example,776//777// permutation = [2, 0, 1]778// reassociation_map for expansion = [[0, 1], [2], [3, 4, 5]]779//780// inverse permutation = [1, 2, 0]781// applied to reassocation_map and then flattened becomes782// flatened permutation = [2, 3, 4, 5, 0, 1]783// final permuation is the inverse of the flattened permutation.784//785// Becomes786//787// permutation=[4, 5, 0, 1, 2, 3]788 789static Operation *createExpandedTransposeOp(PatternRewriter &rewriter,790                                            TransposeOp transposeOp,791                                            Value expandedInput, Value output,792                                            ExpansionInfo &expansionInfo) {793  SmallVector<int64_t> newPerm;794  for (int64_t perm : invertPermutationVector(transposeOp.getPermutation())) {795    auto reassoc = expansionInfo.getExpandedDims(perm);796    for (int64_t dim : reassoc) {797      newPerm.push_back(dim);798    }799  }800  return TransposeOp::create(rewriter, transposeOp.getLoc(), expandedInput,801                             output, invertPermutationVector(newPerm));802}803 804// Create an expanded generic op.805static Operation *createExpandedGenericOp(806    PatternRewriter &rewriter, LinalgOp linalgOp, TypeRange resultTypes,807    ArrayRef<Value> &expandedOpOperands, ArrayRef<Value> outputs,808    ExpansionInfo &expansionInfo, ArrayRef<AffineMap> expandedOpIndexingMaps) {809  // The iterator types of the expanded op are all parallel.810  SmallVector<utils::IteratorType> iteratorTypes(811      expansionInfo.getExpandedOpNumDims(), utils::IteratorType::parallel);812 813  for (auto [i, type] : llvm::enumerate(linalgOp.getIteratorTypesArray()))814    for (auto j : expansionInfo.getExpandedDims(i))815      iteratorTypes[j] = type;816 817  Operation *fused = GenericOp::create(rewriter, linalgOp.getLoc(), resultTypes,818                                       expandedOpOperands, outputs,819                                       expandedOpIndexingMaps, iteratorTypes);820 821  Region &fusedRegion = fused->getRegion(0);822  Region &originalRegion = linalgOp->getRegion(0);823  rewriter.cloneRegionBefore(originalRegion, fusedRegion, fusedRegion.begin());824 825  // Update the index accesses after the expansion.826  updateExpandedGenericOpRegion(rewriter, linalgOp.getLoc(), fusedRegion,827                                expansionInfo);828 829  return fused;830}831 832// Create an expanded fused op that retains the name for certain ops833// such as fill, copy and transpose and produce a generic op for834// rest of linalg ops.835static Operation *createExpandedOp(PatternRewriter &rewriter, LinalgOp linalgOp,836                                   TypeRange resultTypes,837                                   ArrayRef<Value> expandedOpOperands,838                                   ArrayRef<Value> outputs,839                                   ArrayRef<AffineMap> expandedOpIndexingMaps,840                                   ExpansionInfo &expansionInfo) {841 842  return TypeSwitch<Operation *, Operation *>(linalgOp.getOperation())843      .Case<TransposeOp>([&](TransposeOp transposeOp) {844        return createExpandedTransposeOp(rewriter, transposeOp,845                                         expandedOpOperands[0], outputs[0],846                                         expansionInfo);847      })848      .Case<FillOp, CopyOp>([&](Operation *op) {849        return clone(rewriter, linalgOp, resultTypes,850                     llvm::to_vector(llvm::concat<Value>(851                         llvm::to_vector(expandedOpOperands),852                         llvm::to_vector(outputs))));853      })854      .Default([&](Operation *op) {855        return createExpandedGenericOp(rewriter, linalgOp, resultTypes,856                                       expandedOpOperands, outputs,857                                       expansionInfo, expandedOpIndexingMaps);858      });859}860 861/// Implements the fusion of a tensor.collapse_shape or a tensor.expand_shape op862/// and a generic op as explained in `isFusableWithReshapeByExpansion`. Assumes863/// that those conditions have been satisfied.864static std::optional<SmallVector<Value>>865fuseWithReshapeByExpansion(LinalgOp linalgOp, Operation *reshapeOp,866                           OpOperand *fusableOpOperand,867                           PatternRewriter &rewriter) {868  assert(isFusableWithReshapeByDimExpansion(linalgOp, fusableOpOperand) &&869         "preconditions for fuse operation failed");870 871  Location loc = linalgOp.getLoc();872  SmallVector<OpFoldResult> expandedShape;873  SmallVector<AffineMap, 4> reassociationIndices;874  Value src;875  if (auto expandingReshapeOp = dyn_cast<tensor::ExpandShapeOp>(reshapeOp)) {876    // Try to move the dynamic dimensions in output shape before the `linalgOp`877    // to maintain SSA validity878    if (failed(moveValueDefinitions(879            rewriter, expandingReshapeOp.getOutputShape(), linalgOp)))880      return std::nullopt;881 882    expandedShape = expandingReshapeOp.getMixedOutputShape();883    reassociationIndices = expandingReshapeOp.getReassociationMaps();884    src = expandingReshapeOp.getSrc();885  } else {886    auto collapsingReshapeOp = dyn_cast<tensor::CollapseShapeOp>(reshapeOp);887    if (!collapsingReshapeOp)888      return std::nullopt;889 890    expandedShape = tensor::getMixedSizes(891        rewriter, collapsingReshapeOp->getLoc(), collapsingReshapeOp.getSrc());892    reassociationIndices = collapsingReshapeOp.getReassociationMaps();893    src = collapsingReshapeOp.getSrc();894  }895 896  ExpansionInfo expansionInfo;897  if (failed(expansionInfo.compute(linalgOp, fusableOpOperand,898                                   reassociationIndices, expandedShape,899                                   rewriter)))900    return std::nullopt;901 902  SmallVector<AffineMap, 4> expandedOpIndexingMaps = llvm::to_vector<4>(903      llvm::map_range(linalgOp.getIndexingMapsArray(), [&](AffineMap m) {904        return getIndexingMapInExpandedOp(rewriter, m, expansionInfo);905      }));906 907  // Set insertion point to the generic op.908  OpBuilder::InsertionGuard g(rewriter);909  rewriter.setInsertionPoint(linalgOp);910 911  SmallVector<Value> expandedOpOperands;912  expandedOpOperands.reserve(linalgOp.getNumDpsInputs());913  for (OpOperand *opOperand : linalgOp.getDpsInputOperands()) {914    if (opOperand == fusableOpOperand) {915      expandedOpOperands.push_back(src);916      continue;917    }918    if (auto opOperandType =919            dyn_cast<RankedTensorType>(opOperand->get().getType())) {920      AffineMap indexingMap = linalgOp.getMatchingIndexingMap(opOperand);921      SmallVector<OpFoldResult> expandedOperandShape;922      RankedTensorType expandedOperandType;923      std::tie(expandedOperandShape, expandedOperandType) =924          getExpandedShapeAndType(opOperandType, indexingMap, expansionInfo);925      if (expandedOperandType != opOperand->get().getType()) {926        // Reshape the operand to get the right type.927        SmallVector<ReassociationIndices> reassociation =928            getReassociationForExpansion(indexingMap, expansionInfo);929        if (failed(reshapeLikeShapesAreCompatible(930                [&](const Twine &msg) {931                  return rewriter.notifyMatchFailure(linalgOp, msg);932                },933                opOperandType.getShape(), expandedOperandType.getShape(),934                reassociation,935                /*isExpandingReshape=*/true)))936          return std::nullopt;937        expandedOpOperands.push_back(tensor::ExpandShapeOp::create(938            rewriter, loc, expandedOperandType, opOperand->get(), reassociation,939            expandedOperandShape));940        continue;941      }942    }943    expandedOpOperands.push_back(opOperand->get());944  }945 946  SmallVector<Value> outputs;947  for (OpOperand &opOperand : linalgOp.getDpsInitsMutable()) {948    AffineMap indexingMap = linalgOp.getMatchingIndexingMap(&opOperand);949    auto opOperandType = cast<RankedTensorType>(opOperand.get().getType());950    SmallVector<OpFoldResult> expandedOutputShape;951    RankedTensorType expandedOutputType;952    std::tie(expandedOutputShape, expandedOutputType) =953        getExpandedShapeAndType(opOperandType, indexingMap, expansionInfo);954    if (expandedOutputType != opOperand.get().getType()) {955      SmallVector<ReassociationIndices> reassociation =956          getReassociationForExpansion(indexingMap, expansionInfo);957      if (failed(reshapeLikeShapesAreCompatible(958              [&](const Twine &msg) {959                return rewriter.notifyMatchFailure(linalgOp, msg);960              },961              opOperandType.getShape(), expandedOutputType.getShape(),962              reassociation,963              /*isExpandingReshape=*/true)))964        return std::nullopt;965      outputs.push_back(tensor::ExpandShapeOp::create(966          rewriter, loc, expandedOutputType, opOperand.get(), reassociation,967          expandedOutputShape));968    } else {969      outputs.push_back(opOperand.get());970    }971  }972 973  TypeRange resultTypes = ValueRange(outputs).getTypes();974  Operation *fusedOp =975      createExpandedOp(rewriter, linalgOp, resultTypes, expandedOpOperands,976                       outputs, expandedOpIndexingMaps, expansionInfo);977  // Reshape the result values to their original shape if this is a collapsing978  // reshape folded into its consumer.979  SmallVector<Value> resultVals;980  for (OpResult opResult : linalgOp->getOpResults()) {981    int64_t resultNumber = opResult.getResultNumber();982    if (resultTypes[resultNumber] != opResult.getType()) {983      SmallVector<ReassociationIndices> reassociation =984          getReassociationForExpansion(985              linalgOp.getMatchingIndexingMap(986                  linalgOp.getDpsInitOperand(resultNumber)),987              expansionInfo);988      resultVals.push_back(tensor::CollapseShapeOp::create(989          rewriter, linalgOp.getLoc(), opResult.getType(),990          fusedOp->getResult(resultNumber), reassociation));991    } else {992      resultVals.push_back(fusedOp->getResult(resultNumber));993    }994  }995  // Assuming a single result.996  return resultVals;997}998 999namespace {1000 1001/// Pattern to fuse a tensor.collapse_shape op with its consumer structured op,1002/// when the reshape op is collapsing dimensions. The dimensionality of the loop1003/// in the consumer is expanded.1004class FoldWithProducerReshapeOpByExpansion1005    : public OpInterfaceRewritePattern<LinalgOp> {1006public:1007  FoldWithProducerReshapeOpByExpansion(MLIRContext *context,1008                                       ControlFusionFn foldReshapes,1009                                       PatternBenefit benefit = 1)1010      : OpInterfaceRewritePattern<LinalgOp>(context, benefit),1011        controlFoldingReshapes(std::move(foldReshapes)) {}1012 1013  LogicalResult matchAndRewrite(LinalgOp linalgOp,1014                                PatternRewriter &rewriter) const override {1015    for (OpOperand *opOperand : linalgOp.getDpsInputOperands()) {1016      tensor::CollapseShapeOp reshapeOp =1017          opOperand->get().getDefiningOp<tensor::CollapseShapeOp>();1018      if (!reshapeOp)1019        continue;1020      // Fold only if1021      // - The tensor reshape op is folding.1022      // - All constraints of fusing with reshape by expansion are met.1023      if (!isFusableWithReshapeByDimExpansion(linalgOp, opOperand) ||1024          (!controlFoldingReshapes(opOperand)))1025        continue;1026 1027      std::optional<SmallVector<Value>> replacementValues =1028          fuseWithReshapeByExpansion(linalgOp, reshapeOp, opOperand, rewriter);1029      if (!replacementValues)1030        return failure();1031      rewriter.replaceOp(linalgOp, *replacementValues);1032      return success();1033    }1034    return failure();1035  }1036 1037private:1038  ControlFusionFn controlFoldingReshapes;1039};1040 1041/// Carries information about a padded dimension.1042struct PadDimInfo {1043  // The resulting shape after padding each dimension.1044  SmallVector<int64_t> paddedShape;1045 1046  // Low and high padding amounts for each dimension.1047  SmallVector<OpFoldResult> lowPad;1048  SmallVector<OpFoldResult> highPad;1049};1050 1051/// Computes the expanded padding information for the given pad operation based1052/// on the provided expanded shape and reassociation indices. Returns a list of1053/// PadDimInfo containing the low and high padding amounts and the padded1054/// size for each dimension, or failure if the expansion is not possible.1055static FailureOr<PadDimInfo>1056computeExpandedPadding(tensor::PadOp padOp, ArrayRef<int64_t> expandedShape,1057                       ArrayRef<ReassociationIndices> reassociations,1058                       PatternRewriter &rewriter) {1059  // If the padding value depends on the index values of the pad operation,1060  // then it may not be valid to expand the dimensions, since it will change1061  // the index values on which the padding value depends. This is not currently1062  // supported by the pad expansion patterns, but it could be implemented1063  // similarly to the expansion of linalg.generic ops with linalg.index ops in1064  // the body, as is done in `updateExpandedGenericOpRegion`.1065  if (!padOp.getConstantPaddingValue())1066    return failure();1067 1068  // Expanded dimensions cannot have padding because the resulting padding may1069  // not be representable by a tensor.pad op. There are some special cases where1070  // it is possible (like expanding unit dims), but supporting these cases is1071  // NYI, so disallow it for now.1072  ArrayRef<int64_t> low = padOp.getStaticLow();1073  ArrayRef<int64_t> high = padOp.getStaticHigh();1074  for (auto [reInd, l, h] : llvm::zip_equal(reassociations, low, high)) {1075    if (reInd.size() != 1 && (l != 0 || h != 0))1076      return failure();1077  }1078 1079  SmallVector<OpFoldResult> mixedLowPad(padOp.getMixedLowPad());1080  SmallVector<OpFoldResult> mixedHighPad(padOp.getMixedHighPad());1081  ArrayRef<int64_t> paddedShape = padOp.getResultType().getShape();1082  PadDimInfo padDimInfo;1083  padDimInfo.paddedShape.assign(expandedShape);1084  padDimInfo.lowPad.assign(expandedShape.size(), rewriter.getIndexAttr(0));1085  padDimInfo.highPad.assign(expandedShape.size(), rewriter.getIndexAttr(0));1086  for (auto [idx, reInd] : llvm::enumerate(reassociations)) {1087    if (reInd.size() == 1) {1088      padDimInfo.paddedShape[reInd[0]] = paddedShape[idx];1089      padDimInfo.lowPad[reInd[0]] = mixedLowPad[idx];1090      padDimInfo.highPad[reInd[0]] = mixedHighPad[idx];1091    }1092  }1093 1094  return padDimInfo;1095}1096 1097class FoldPadWithProducerReshapeOpByExpansion1098    : public OpRewritePattern<tensor::PadOp> {1099public:1100  FoldPadWithProducerReshapeOpByExpansion(MLIRContext *context,1101                                          ControlFusionFn foldReshapes,1102                                          PatternBenefit benefit = 1)1103      : OpRewritePattern<tensor::PadOp>(context, benefit),1104        controlFoldingReshapes(std::move(foldReshapes)) {}1105 1106  LogicalResult matchAndRewrite(tensor::PadOp padOp,1107                                PatternRewriter &rewriter) const override {1108    tensor::CollapseShapeOp reshapeOp =1109        padOp.getSource().getDefiningOp<tensor::CollapseShapeOp>();1110    if (!reshapeOp)1111      return failure();1112 1113    if (!controlFoldingReshapes(&padOp.getSourceMutable())) {1114      return rewriter.notifyMatchFailure(padOp,1115                                         "fusion blocked by control function");1116    }1117 1118    RankedTensorType expandedType = reshapeOp.getSrcType();1119    SmallVector<ReassociationIndices> reassociations =1120        reshapeOp.getReassociationIndices();1121    FailureOr<PadDimInfo> maybeExpandedPadding = computeExpandedPadding(1122        padOp, expandedType.getShape(), reassociations, rewriter);1123    if (failed(maybeExpandedPadding))1124      return failure();1125    PadDimInfo &expandedPadding = maybeExpandedPadding.value();1126 1127    Location loc = padOp->getLoc();1128    RankedTensorType expandedPaddedType =1129        padOp.getResultType().clone(expandedPadding.paddedShape);1130 1131    auto newPadOp = tensor::PadOp::create(1132        rewriter, loc, expandedPaddedType, reshapeOp.getSrc(),1133        expandedPadding.lowPad, expandedPadding.highPad,1134        padOp.getConstantPaddingValue(), padOp.getNofold());1135 1136    rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(1137        padOp, padOp.getResultType(), newPadOp.getResult(), reassociations);1138 1139    return success();1140  }1141 1142private:1143  ControlFusionFn controlFoldingReshapes;1144};1145 1146class FoldReshapeWithProducerPadOpByExpansion1147    : public OpRewritePattern<tensor::ExpandShapeOp> {1148public:1149  FoldReshapeWithProducerPadOpByExpansion(MLIRContext *context,1150                                          ControlFusionFn foldReshapes,1151                                          PatternBenefit benefit = 1)1152      : OpRewritePattern<tensor::ExpandShapeOp>(context, benefit),1153        controlFoldingReshapes(std::move(foldReshapes)) {}1154 1155  LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp,1156                                PatternRewriter &rewriter) const override {1157    tensor::PadOp padOp = expandOp.getSrc().getDefiningOp<tensor::PadOp>();1158    if (!padOp)1159      return failure();1160 1161    if (!controlFoldingReshapes(&expandOp.getSrcMutable())) {1162      return rewriter.notifyMatchFailure(expandOp,1163                                         "fusion blocked by control function");1164    }1165 1166    RankedTensorType expandedType = expandOp.getResultType();1167    SmallVector<ReassociationIndices> reassociations =1168        expandOp.getReassociationIndices();1169    FailureOr<PadDimInfo> maybeExpandedPadding = computeExpandedPadding(1170        padOp, expandedType.getShape(), reassociations, rewriter);1171    if (failed(maybeExpandedPadding))1172      return failure();1173    PadDimInfo &expandedPadding = maybeExpandedPadding.value();1174 1175    Location loc = expandOp->getLoc();1176    SmallVector<OpFoldResult> newExpandedSizes = expandOp.getMixedOutputShape();1177    SmallVector<int64_t> newExpandedShape(expandedType.getShape());1178    rewriter.setInsertionPointAfterValue(padOp.getSource());1179    SmallVector<OpFoldResult> padSrcSizes =1180        tensor::getMixedSizes(rewriter, loc, padOp.getSource());1181    for (auto [idx, reInd] : llvm::enumerate(reassociations)) {1182      // We know that any reassociation with multiple dims is not padded because1183      // of the requirements of computeExpandedPadding.1184      if (reInd.size() == 1) {1185        newExpandedShape[reInd[0]] = padOp.getSourceType().getDimSize(idx);1186        newExpandedSizes[reInd[0]] = padSrcSizes[idx];1187      }1188    }1189    RankedTensorType newExpandedType = expandedType.clone(newExpandedShape);1190    auto newExpandOp = tensor::ExpandShapeOp::create(1191        rewriter, loc, newExpandedType, padOp.getSource(), reassociations,1192        newExpandedSizes);1193    RankedTensorType expandedPaddedType =1194        padOp.getResultType().clone(expandedPadding.paddedShape);1195    rewriter.setInsertionPoint(expandOp);1196    auto newPadOp = tensor::PadOp::create(1197        rewriter, loc, expandedPaddedType, newExpandOp.getResult(),1198        expandedPadding.lowPad, expandedPadding.highPad,1199        padOp.getConstantPaddingValue(), padOp.getNofold());1200 1201    rewriter.replaceOp(expandOp, newPadOp.getResult());1202 1203    return success();1204  }1205 1206private:1207  ControlFusionFn controlFoldingReshapes;1208};1209 1210/// Pattern to fold a tensor.expand_shape op with its producer generic op1211/// by expanding the dimensionality of the loop in the producer op.1212struct FoldReshapeWithGenericOpByExpansion1213    : public OpRewritePattern<tensor::ExpandShapeOp> {1214 1215  FoldReshapeWithGenericOpByExpansion(MLIRContext *context,1216                                      ControlFusionFn foldReshapes,1217                                      PatternBenefit benefit = 1)1218      : OpRewritePattern<tensor::ExpandShapeOp>(context, benefit),1219        controlFoldingReshapes(std::move(foldReshapes)) {}1220 1221  LogicalResult matchAndRewrite(tensor::ExpandShapeOp reshapeOp,1222                                PatternRewriter &rewriter) const override {1223    // Fold only if all constraints of fusing with reshape by expansion are met.1224    auto producerResult = dyn_cast<OpResult>(reshapeOp.getSrc());1225    if (!producerResult) {1226      return rewriter.notifyMatchFailure(reshapeOp,1227                                         "source not produced by an operation");1228    }1229 1230    auto producer = dyn_cast<LinalgOp>(producerResult.getOwner());1231    if (!producer) {1232      return rewriter.notifyMatchFailure(reshapeOp,1233                                         "producer not a generic op");1234    }1235 1236    if (!isFusableWithReshapeByDimExpansion(1237            producer,1238            producer.getDpsInitOperand(producerResult.getResultNumber()))) {1239      return rewriter.notifyMatchFailure(1240          reshapeOp, "failed preconditions of fusion with producer generic op");1241    }1242 1243    if (!controlFoldingReshapes(&reshapeOp.getSrcMutable())) {1244      return rewriter.notifyMatchFailure(reshapeOp,1245                                         "fusion blocked by control function");1246    }1247 1248    std::optional<SmallVector<Value>> replacementValues =1249        fuseWithReshapeByExpansion(1250            producer, reshapeOp,1251            producer.getDpsInitOperand(producerResult.getResultNumber()),1252            rewriter);1253    if (!replacementValues) {1254      return rewriter.notifyMatchFailure(reshapeOp,1255                                         "fusion by expansion failed");1256    }1257 1258    // Find the replacement for the reshape op. Since the replacements have the1259    // same type as the returns of the original generic op, the consumer reshape1260    // op can be replaced by the source of the collapse_shape op that defines1261    // the replacement.1262    Value reshapeReplacement =1263        (*replacementValues)[cast<OpResult>(reshapeOp.getSrc())1264                                 .getResultNumber()];1265    if (auto collapseOp =1266            reshapeReplacement.getDefiningOp<tensor::CollapseShapeOp>()) {1267      reshapeReplacement = collapseOp.getSrc();1268    }1269    rewriter.replaceOp(reshapeOp, reshapeReplacement);1270    rewriter.replaceOp(producer, *replacementValues);1271    return success();1272  }1273 1274private:1275  ControlFusionFn controlFoldingReshapes;1276};1277} // namespace1278 1279//===---------------------------------------------------------------------===//1280// Methods and patterns to fuse reshape with linalg.generic operations by1281// contraction of dimensions.1282//===---------------------------------------------------------------------===//1283 1284/// For a given list of indices in the range of the `indexingMap` that are1285/// folded, return the indices of the corresponding domain. Return1286/// `std::nullopt` on failure. Ensures that all the elements of the returned1287/// reassociation are distinct.1288static ReassociationIndices1289getDomainReassociation(AffineMap indexingMap,1290                       ReassociationIndicesRef rangeReassociation) {1291  assert(indexingMap.isProjectedPermutation() &&1292         "expected projected permutation");1293 1294  ReassociationIndices domainReassociation = llvm::to_vector<4>(1295      llvm::map_range(rangeReassociation, [&](int64_t pos) -> int64_t {1296        return cast<AffineDimExpr>(indexingMap.getResults()[pos]).getPosition();1297      }));1298  // The projected permutation semantics ensures that there is no repetition of1299  // the domain indices.1300  return domainReassociation;1301}1302 1303/// For a given `dimSequence`, check if the sequence is conserved in the1304/// `indexingMap`. `indexingMap` is expected to be a projected permutation.1305/// Non-existence of the sequence returns true as well.1306bool mlir::linalg::isDimSequencePreserved(AffineMap indexingMap,1307                                          ReassociationIndicesRef dimSequence) {1308  assert(!dimSequence.empty() &&1309         "expected non-empty list for dimension sequence");1310  assert(indexingMap.isProjectedPermutation() &&1311         "expected indexing map to be projected permutation");1312 1313  llvm::SmallDenseSet<unsigned, 4> sequenceElements;1314  sequenceElements.insert_range(dimSequence);1315 1316  unsigned dimSequenceStart = dimSequence[0];1317  for (const auto &expr : enumerate(indexingMap.getResults())) {1318    unsigned dimInMapStart = cast<AffineDimExpr>(expr.value()).getPosition();1319    // 1.  Check if this start of the sequence.1320    if (dimInMapStart == dimSequenceStart) {1321      if (expr.index() + dimSequence.size() > indexingMap.getNumResults())1322        return false;1323      // 1a. Check if sequence is preserved.1324      for (const auto &dimInSequence : enumerate(dimSequence)) {1325        unsigned dimInMap =1326            cast<AffineDimExpr>(1327                indexingMap.getResult(expr.index() + dimInSequence.index()))1328                .getPosition();1329        if (dimInMap != dimInSequence.value())1330          return false;1331      }1332      // Found the sequence. Projected permutation1333      // enforces that all AffineDimExprs in the result are unique, so no1334      // further checks are needed.1335      return true;1336    }1337    // 2. If position in the expr (which is of type AffineDimExpr) is part1338    // of sequence, return false here. This implies the entire sequence does not1339    // exist in the indexing map.1340    if (sequenceElements.count(dimInMapStart))1341      return false;1342  }1343  // 3. No element of sequence found. Return true.1344  return true;1345}1346 1347bool mlir::linalg::areDimSequencesPreserved(1348    ArrayRef<AffineMap> maps, ArrayRef<ReassociationIndices> dimSequences) {1349  return llvm::all_of(maps, [&](AffineMap map) {1350    return llvm::all_of(dimSequences, [&](ReassociationIndicesRef dimSequence) {1351      return isDimSequencePreserved(map, dimSequence);1352    });1353  });1354}1355 1356// Return the list of dimensions of the iteration domain that can be1357// collapsed to allow for fusion with the a producer that is an expand_shape1358// operation. If all dimensions created by expansion can be collapsed in the1359// iteration space then the reshape is defunct.1360//1361// Example:1362//1363// ```mlir1364// #map = affine_map<(d0, d1) -> (d0, d1)>1365// %1 = tensor.expand_shape %0 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32>1366// %2 = tensor.empty [..] : tensor<?x4xf32>1367// %3 = linalg.generic {1368//     indexing_maps = [#map, #map],1369//     iterator_types = ["parallel" ,"parallel"]}1370//     ins(%1 : tensor<?x4xf32>) outs(%2 : tensor<?x4xf32>) {.. }1371// ```1372//1373// can be fused by collapsing the dimensions of the iteration space.1374//1375// ```mlir1376// #map = affine_map<(d0) -> (d0)>1377// %2 = tensor.empty [..] : tensor<?xf32>1378// %3 = linalg.generic {1379//     indexing_maps = [#map, #map],1380//     iterator_types = ["parallel"]}1381//     ins(%1 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {.. }1382// %4 = tensor.expand_shape %3 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32>1383// ```1384//1385// In the following example,1386//1387// ```mlir1388// #map0 = affine_map<(d0, d1) -> (d0, d1)>1389// #map1 = affine_map<(d0, d1) -> (d1, d0)>1390// %1 = tensor.expand_shape %0 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32>1391// %2 = tensor.empty [..] : tensor<4x?xf32>1392// %2 = linalg.generic {1393//     indexing_maps = [#map0, #map1],1394//     iterator_types = ["parallel" ,"parallel"]}1395//     ins(%1 : tensor<?x4xf32>) outs(%2 : tensor<4x?xf32>) {.. }1396// ```1397//1398// the reshape cannot be fused with the generic op by collapsing the op1399// dimensions since the indexing maps will have to contain mods and divs1400// to preserve the accesses pattern. When no dimensions of the iteration1401// space are collapsable and empty vector is returned.1402static SmallVector<ReassociationIndices>1403getCollapsableIterationSpaceDims(GenericOp genericOp, OpOperand *fusableOperand,1404                                 ArrayRef<ReassociationIndices> reassociation) {1405  // Some basic checks for this fusion to be valid.1406  if (!genericOp.hasPureTensorSemantics())1407    return {};1408 1409  if (!llvm::all_of(genericOp.getIndexingMapsArray(), [](AffineMap map) {1410        return map.isProjectedPermutation();1411      })) {1412    return {};1413  }1414 1415  // Compute all the loops with the reduction iterator types.1416  SmallVector<unsigned> reductionDims;1417  genericOp.getReductionDims(reductionDims);1418 1419  llvm::SmallDenseSet<unsigned, 4> processedIterationDims;1420  AffineMap indexingMap = genericOp.getMatchingIndexingMap(fusableOperand);1421  auto iteratorTypes = genericOp.getIteratorTypesArray();1422  SmallVector<ReassociationIndices> iterationSpaceReassociation;1423  for (ReassociationIndicesRef foldedRangeDims : reassociation) {1424    assert(!foldedRangeDims.empty() && "unexpected empty reassociation");1425 1426    // Ignore dims that are not folded.1427    if (foldedRangeDims.size() == 1)1428      continue;1429 1430    ReassociationIndices foldedIterationSpaceDims =1431        getDomainReassociation(indexingMap, foldedRangeDims);1432 1433    // Check that the folded iteration dims do not contain already processed1434    // dims.1435    if (llvm::any_of(foldedIterationSpaceDims, [&](int64_t dim) {1436          return processedIterationDims.count(dim);1437        }))1438      continue;1439 1440    // Check that all folded iterator types are all parallel or all reductions.1441    utils::IteratorType startIteratorType =1442        iteratorTypes[foldedIterationSpaceDims[0]];1443    if (!isParallelIterator(startIteratorType) &&1444        !isReductionIterator(startIteratorType))1445      continue;1446    if (llvm::any_of(foldedIterationSpaceDims, [&](int64_t dim) {1447          return iteratorTypes[dim] != startIteratorType;1448        }))1449      continue;1450 1451    // If the folded dimensions correspond to a "reduction" iterator type,1452    // the folded dimensions need to be "in-order". Strictly speaking this is1453    // not necessary, for reductions that are associative and commutative,  but1454    // using a more strict definition of reduction for now.1455    if (isReductionIterator(startIteratorType)) {1456      bool isContiguous = false;1457      for (const auto &startDim : llvm::enumerate(reductionDims)) {1458        // Move window in `reductionDims` to start of the folded iteration dims.1459        if (startDim.value() != foldedIterationSpaceDims[0])1460          continue;1461        // If sizes doesnt match, trivial not contiguous. This condition should1462        // not be hit.1463        if (startDim.index() + foldedIterationSpaceDims.size() >1464            reductionDims.size())1465          break;1466        // Check that the contiguity is maintained.1467        isContiguous = true;1468        for (const auto &foldedDim :1469             llvm::enumerate(foldedIterationSpaceDims)) {1470          if (reductionDims[foldedDim.index() + startDim.index()] !=1471              foldedDim.value()) {1472            isContiguous = false;1473            break;1474          }1475        }1476        break;1477      }1478      if (!isContiguous)1479        continue;1480    }1481 1482    // Check that the sequence is preserved in all indexing maps.1483    if (llvm::any_of(genericOp.getIndexingMapsArray(),1484                     [&](AffineMap indexingMap) {1485                       return !isDimSequencePreserved(indexingMap,1486                                                      foldedIterationSpaceDims);1487                     }))1488      continue;1489 1490    processedIterationDims.insert_range(foldedIterationSpaceDims);1491    iterationSpaceReassociation.emplace_back(1492        std::move(foldedIterationSpaceDims));1493  }1494 1495  return iterationSpaceReassociation;1496}1497 1498/// Helper class to carry state while collapsing the `linalg.generic` op.1499namespace {1500class CollapsingInfo {1501public:1502  LogicalResult initialize(unsigned origNumLoops,1503                           ArrayRef<ReassociationIndices> foldedIterationDims) {1504    llvm::SmallDenseSet<int64_t, 4> processedDims;1505    // Find all the dims that are folded.1506    for (ReassociationIndicesRef foldedIterationDim : foldedIterationDims) {1507      if (foldedIterationDim.empty())1508        continue;1509      // If the folded dims contain dims already folded, that's illegal1510      // specification. Repetition within a list is also illegal.1511      for (auto dim : foldedIterationDim) {1512        if (dim >= origNumLoops)1513          return failure();1514        if (processedDims.count(dim))1515          return failure();1516        processedDims.insert(dim);1517      }1518      collapsedOpToOrigOpIterationDim.emplace_back(foldedIterationDim.begin(),1519                                                   foldedIterationDim.end());1520    }1521    if (processedDims.size() > origNumLoops)1522      return failure();1523 1524    // Add all the preserved dims of the original op as single1525    // elements to `collapsedOpToOrigOpIterationDim`.1526    for (auto dim : llvm::seq<int64_t>(0, origNumLoops)) {1527      if (processedDims.count(dim))1528        continue;1529      collapsedOpToOrigOpIterationDim.emplace_back(ReassociationIndices{dim});1530    }1531 1532    llvm::sort(collapsedOpToOrigOpIterationDim,1533               [&](ReassociationIndicesRef lhs, ReassociationIndicesRef rhs) {1534                 return lhs[0] < rhs[0];1535               });1536    origOpToCollapsedOpIterationDim.resize(origNumLoops);1537    for (const auto &foldedDims :1538         llvm::enumerate(collapsedOpToOrigOpIterationDim)) {1539      for (const auto &dim : enumerate(foldedDims.value()))1540        origOpToCollapsedOpIterationDim[dim.value()] =1541            std::make_pair<int64_t, unsigned>(foldedDims.index(), dim.index());1542    }1543    return success();1544  }1545 1546  /// Return mapping from collapsed loop domain to original loop domain.1547  ArrayRef<ReassociationIndices> getCollapsedOpToOrigOpMapping() const {1548    return collapsedOpToOrigOpIterationDim;1549  }1550 1551  /// Return mapping from original loop domain to collapsed loop domain. The1552  /// mapping is a pair. First value is the dimension in the collapsed loop that1553  /// the original loop is mapped to. Second is the relative position in folded1554  /// list of this domain. For example if the original loop domain is 3D, and1555  /// the collapsed loop domain is folding all of it, i.e.1556  ///1557  /// ```1558  /// collapsedOpToOrigOpMapping = [[0, 1, 2] [3, 4]]`1559  /// ```1560  ///1561  /// then1562  ///1563  /// ```1564  ///  origOpToCollapsedOpMapping[0] = {0, 0};1565  ///  origOpToCollapsedOpMapping[1] = {0, 1};1566  ///  origOpToCollapsedOpMapping[2] = {0, 2};1567  ///  origOpToCollapsedOpMapping[3] = {1, 0};1568  ///  origOpToCollapsedOpMapping[4] = {1, 1};1569  /// ```1570  ///1571  ArrayRef<std::pair<int64_t, unsigned>> getOrigOpToCollapsedOpMapping() const {1572    return origOpToCollapsedOpIterationDim;1573  }1574 1575  /// Return the collapsed op iteration domain rank.1576  unsigned getCollapsedOpIterationRank() const {1577    return collapsedOpToOrigOpIterationDim.size();1578  }1579 1580private:1581  /// Map from the iteration domain index in collapsed op to the iteration1582  /// domain indices in the original op.1583  SmallVector<ReassociationIndices> collapsedOpToOrigOpIterationDim;1584 1585  /// Map from iteration domain index in the original op to the iteration domain1586  /// index in the collapsed op.1587  SmallVector<std::pair<int64_t, unsigned>> origOpToCollapsedOpIterationDim;1588};1589} // namespace1590 1591/// Get the iterator types for the collapsed operation given the original1592/// iterator types and collapsed dimensions.1593static SmallVector<utils::IteratorType>1594getCollapsedOpIteratorTypes(ArrayRef<utils::IteratorType> iteratorTypes,1595                            const CollapsingInfo &collapsingInfo) {1596  SmallVector<utils::IteratorType> collapsedIteratorTypes;1597  for (ReassociationIndicesRef foldedIterDims :1598       collapsingInfo.getCollapsedOpToOrigOpMapping()) {1599    assert(!foldedIterDims.empty() &&1600           "reassociation indices expected to have non-empty sets");1601    // Just pick the iterator type of the first folded dim. Pre-condition checks1602    // expected to have checked that iterator types of all folded dimensions are1603    // the same.1604    collapsedIteratorTypes.push_back(iteratorTypes[foldedIterDims[0]]);1605  }1606  return collapsedIteratorTypes;1607}1608 1609/// Compute the indexing map in the collapsed op that corresponds to the given1610/// `indexingMap` of the original operation.1611static AffineMap1612getCollapsedOpIndexingMap(AffineMap indexingMap,1613                          const CollapsingInfo &collapsingInfo) {1614  MLIRContext *context = indexingMap.getContext();1615  assert(indexingMap.isProjectedPermutation() &&1616         "expected indexing map to be projected permutation");1617  SmallVector<AffineExpr> resultExprs;1618  auto origOpToCollapsedOpMapping =1619      collapsingInfo.getOrigOpToCollapsedOpMapping();1620  for (auto expr : indexingMap.getResults()) {1621    unsigned dim = cast<AffineDimExpr>(expr).getPosition();1622    // If the dim is not the first of the collapsed dim, do nothing.1623    if (origOpToCollapsedOpMapping[dim].second != 0)1624      continue;1625    // The next n-dims are guaranteed to be collapsed. So just use the1626    // iteration dimension of the collapsed op.1627    resultExprs.push_back(1628        getAffineDimExpr(origOpToCollapsedOpMapping[dim].first, context));1629  }1630  return AffineMap::get(collapsingInfo.getCollapsedOpIterationRank(), 0,1631                        resultExprs, context);1632}1633 1634/// Return the `reassociation` indices to use to collapse the operand when the1635/// iteration space of a generic op is collapsed.1636static SmallVector<ReassociationIndices>1637getOperandReassociation(AffineMap indexingMap,1638                        const CollapsingInfo &collapsingInfo) {1639  unsigned counter = 0;1640  SmallVector<ReassociationIndices> operandReassociation;1641  auto origOpToCollapsedOpMapping =1642      collapsingInfo.getOrigOpToCollapsedOpMapping();1643  auto collapsedOpToOrigOpMapping =1644      collapsingInfo.getCollapsedOpToOrigOpMapping();1645  while (counter < indexingMap.getNumResults()) {1646    unsigned dim =1647        cast<AffineDimExpr>(indexingMap.getResult(counter)).getPosition();1648    // This is the start of a collapsed dimensions of the iteration that1649    // is gauranteed to be preserved in the indexing map. The number of folded1650    // dims is obtained from the collapsed op to original op mapping.1651    unsigned numFoldedDims =1652        collapsedOpToOrigOpMapping[origOpToCollapsedOpMapping[dim].first]1653            .size();1654    if (origOpToCollapsedOpMapping[dim].second == 0) {1655      auto range = llvm::seq<unsigned>(counter, counter + numFoldedDims);1656      operandReassociation.emplace_back(range.begin(), range.end());1657    }1658    counter += numFoldedDims;1659  }1660  return operandReassociation;1661}1662 1663/// Get the new value to use for a given `OpOperand` in the collapsed operation.1664static Value getCollapsedOpOperand(Location loc, LinalgOp op,1665                                   OpOperand *opOperand,1666                                   const CollapsingInfo &collapsingInfo,1667                                   OpBuilder &builder) {1668  AffineMap indexingMap = op.getMatchingIndexingMap(opOperand);1669  SmallVector<ReassociationIndices> operandReassociation =1670      getOperandReassociation(indexingMap, collapsingInfo);1671 1672  // If the number of entries in the reassociation for the operand is same as1673  // the number of results of the indexing map, then nothing to do for this1674  // operand.1675  Value operand = opOperand->get();1676  if (operandReassociation.size() == indexingMap.getNumResults())1677    return operand;1678 1679  // Insert a reshape to collapse the dimensions.1680  if (isa<MemRefType>(operand.getType())) {1681    return memref::CollapseShapeOp::create(builder, loc, operand,1682                                           operandReassociation)1683        .getResult();1684  }1685  return tensor::CollapseShapeOp::create(builder, loc, operand,1686                                         operandReassociation)1687      .getResult();1688}1689 1690/// Modify the `linalg.index` operations in the original generic op, to its1691/// value in the collapsed operation.1692static void generateCollapsedIndexingRegion(1693    Location loc, Block *block, const CollapsingInfo &collapsingInfo,1694    ArrayRef<OpFoldResult> loopRange, RewriterBase &rewriter) {1695  OpBuilder::InsertionGuard g(rewriter);1696  rewriter.setInsertionPointToStart(block);1697 1698  // Collect all the original index ops.1699  auto indexOps = llvm::to_vector(block->getOps<linalg::IndexOp>());1700 1701  // For each folded dimension list resolve the original induction variable1702  // values in terms of the folded dimension induction variable.1703  //   i_{folded} = (i_0 * d1 + i1) * d2 + i2.1704  // can be inverted to1705  //   i2 = i_{folded} % d21706  //   i1 = (i_{folded} / d2) % d11707  //   i0 = i_{folded} / (d1 * d2)1708  llvm::DenseMap<unsigned, Value> indexReplacementVals;1709  for (auto foldedDims :1710       enumerate(collapsingInfo.getCollapsedOpToOrigOpMapping())) {1711    ReassociationIndicesRef foldedDimsRef(foldedDims.value());1712    Value newIndexVal =1713        linalg::IndexOp::create(rewriter, loc, foldedDims.index());1714    for (auto dim : llvm::reverse(foldedDimsRef.drop_front())) {1715      Value loopDim =1716          getValueOrCreateConstantIndexOp(rewriter, loc, loopRange[dim]);1717      indexReplacementVals[dim] =1718          rewriter.createOrFold<arith::RemSIOp>(loc, newIndexVal, loopDim);1719      newIndexVal =1720          rewriter.createOrFold<arith::DivSIOp>(loc, newIndexVal, loopDim);1721    }1722    indexReplacementVals[foldedDims.value().front()] = newIndexVal;1723  }1724 1725  for (auto indexOp : indexOps) {1726    auto dim = indexOp.getDim();1727    rewriter.replaceOp(indexOp, indexReplacementVals[dim]);1728  }1729}1730 1731static void collapseOperandsAndResults(LinalgOp op,1732                                       const CollapsingInfo &collapsingInfo,1733                                       RewriterBase &rewriter,1734                                       SmallVectorImpl<Value> &inputOperands,1735                                       SmallVectorImpl<Value> &outputOperands,1736                                       SmallVectorImpl<Type> &resultTypes) {1737  Location loc = op->getLoc();1738  inputOperands =1739      llvm::map_to_vector(op.getDpsInputOperands(), [&](OpOperand *opOperand) {1740        return getCollapsedOpOperand(loc, op, opOperand, collapsingInfo,1741                                     rewriter);1742      });1743 1744  // Get the output operands and result types.1745  resultTypes.reserve(op.getNumDpsInits());1746  outputOperands.reserve(op.getNumDpsInits());1747  for (OpOperand &output : op.getDpsInitsMutable()) {1748    Value newOutput =1749        getCollapsedOpOperand(loc, op, &output, collapsingInfo, rewriter);1750    outputOperands.push_back(newOutput);1751    // If the op has "buffer semantics", then the init operands are ranked1752    // memrefs and the op has no results.1753    if (!op.hasPureBufferSemantics())1754      resultTypes.push_back(newOutput.getType());1755  }1756}1757 1758/// Clone a `LinalgOp` to a collapsed version of same name1759template <typename OpTy>1760static OpTy cloneToCollapsedOp(RewriterBase &rewriter, OpTy origOp,1761                               const CollapsingInfo &collapsingInfo) {1762  return nullptr;1763}1764 1765/// Collapse any `LinalgOp` that does not require any specialization such as1766/// indexing_maps, iterator_types, etc.1767template <>1768LinalgOp cloneToCollapsedOp<LinalgOp>(RewriterBase &rewriter, LinalgOp origOp,1769                                      const CollapsingInfo &collapsingInfo) {1770  SmallVector<Value> inputOperands, outputOperands;1771  SmallVector<Type> resultTypes;1772  collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,1773                             outputOperands, resultTypes);1774 1775  return clone(1776      rewriter, origOp, resultTypes,1777      llvm::to_vector(llvm::concat<Value>(inputOperands, outputOperands)));1778}1779 1780/// Collapse a `GenericOp`1781template <>1782GenericOp cloneToCollapsedOp<GenericOp>(RewriterBase &rewriter,1783                                        GenericOp origOp,1784                                        const CollapsingInfo &collapsingInfo) {1785  SmallVector<Value> inputOperands, outputOperands;1786  SmallVector<Type> resultTypes;1787  collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,1788                             outputOperands, resultTypes);1789  SmallVector<AffineMap> indexingMaps(1790      llvm::map_range(origOp.getIndexingMapsArray(), [&](AffineMap map) {1791        return getCollapsedOpIndexingMap(map, collapsingInfo);1792      }));1793 1794  SmallVector<utils::IteratorType> iteratorTypes(getCollapsedOpIteratorTypes(1795      origOp.getIteratorTypesArray(), collapsingInfo));1796 1797  GenericOp collapsedOp = linalg::GenericOp::create(1798      rewriter, origOp.getLoc(), resultTypes, inputOperands, outputOperands,1799      indexingMaps, iteratorTypes,1800      [](OpBuilder &builder, Location loc, ValueRange args) {});1801  Block *origOpBlock = &origOp->getRegion(0).front();1802  Block *collapsedOpBlock = &collapsedOp->getRegion(0).front();1803  rewriter.mergeBlocks(origOpBlock, collapsedOpBlock,1804                       collapsedOpBlock->getArguments());1805  return collapsedOp;1806}1807 1808static LinalgOp createCollapsedOp(LinalgOp op,1809                                  const CollapsingInfo &collapsingInfo,1810                                  RewriterBase &rewriter) {1811  if (GenericOp genericOp = dyn_cast<GenericOp>(op.getOperation())) {1812    return cloneToCollapsedOp(rewriter, genericOp, collapsingInfo);1813  } else {1814    return cloneToCollapsedOp(rewriter, op, collapsingInfo);1815  }1816}1817 1818/// Implementation of fusion with reshape operation by collapsing dimensions.1819FailureOr<CollapseResult> mlir::linalg::collapseOpIterationDims(1820    LinalgOp op, ArrayRef<ReassociationIndices> foldedIterationDims,1821    RewriterBase &rewriter) {1822  // Bail on trivial no-op cases.1823  if (op.getNumLoops() <= 1 || foldedIterationDims.empty() ||1824      llvm::all_of(foldedIterationDims, [](ReassociationIndicesRef foldedDims) {1825        return foldedDims.size() <= 1;1826      }))1827    return failure();1828 1829  CollapsingInfo collapsingInfo;1830  if (failed(1831          collapsingInfo.initialize(op.getNumLoops(), foldedIterationDims))) {1832    return rewriter.notifyMatchFailure(1833        op, "illegal to collapse specified dimensions");1834  }1835 1836  bool hasPureBufferSemantics = op.hasPureBufferSemantics();1837  if (hasPureBufferSemantics &&1838      !llvm::all_of(op->getOpOperands(), [&](OpOperand &opOperand) -> bool {1839        MemRefType memRefToCollapse =1840            dyn_cast<MemRefType>(opOperand.get().getType());1841        if (!memRefToCollapse)1842          return true;1843 1844        AffineMap indexingMap = op.getMatchingIndexingMap(&opOperand);1845        SmallVector<ReassociationIndices> operandReassociation =1846            getOperandReassociation(indexingMap, collapsingInfo);1847        return memref::CollapseShapeOp::isGuaranteedCollapsible(1848            memRefToCollapse, operandReassociation);1849      }))1850    return rewriter.notifyMatchFailure(op,1851                                       "memref is not guaranteed collapsible");1852 1853  // Bail on non-canonical ranges.1854  SmallVector<Range> loopRanges = op.createLoopRanges(rewriter, op.getLoc());1855  auto opFoldIsConstantValue = [](OpFoldResult ofr, int64_t value) {1856    if (auto attr = llvm::dyn_cast_if_present<Attribute>(ofr))1857      return cast<IntegerAttr>(attr).getInt() == value;1858    llvm::APInt actual;1859    return matchPattern(cast<Value>(ofr), m_ConstantInt(&actual)) &&1860           actual.getSExtValue() == value;1861  };1862  if (!llvm::all_of(loopRanges, [&](Range range) {1863        return opFoldIsConstantValue(range.offset, 0) &&1864               opFoldIsConstantValue(range.stride, 1);1865      })) {1866    return rewriter.notifyMatchFailure(1867        op, "expected all loop ranges to have zero start and unit stride");1868  }1869 1870  LinalgOp collapsedOp = createCollapsedOp(op, collapsingInfo, rewriter);1871 1872  Location loc = op->getLoc();1873  SmallVector<OpFoldResult> loopBound =1874      llvm::map_to_vector(loopRanges, [](Range range) { return range.size; });1875 1876  if (collapsedOp.hasIndexSemantics()) {1877    // Collect the loop range of the generic op.1878    OpBuilder::InsertionGuard g(rewriter);1879    rewriter.setInsertionPoint(collapsedOp);1880    generateCollapsedIndexingRegion(loc, &collapsedOp->getRegion(0).front(),1881                                    collapsingInfo, loopBound, rewriter);1882  }1883 1884  // Insert expanding reshape for the result to get back the original result1885  // type.1886  SmallVector<Value> results;1887  for (const auto &originalResult : llvm::enumerate(op->getResults())) {1888    Value collapsedOpResult = collapsedOp->getResult(originalResult.index());1889    auto originalResultType =1890        cast<ShapedType>(originalResult.value().getType());1891    auto collapsedOpResultType = cast<ShapedType>(collapsedOpResult.getType());1892    if (collapsedOpResultType.getRank() != originalResultType.getRank()) {1893      AffineMap indexingMap =1894          op.getIndexingMapMatchingResult(originalResult.value());1895      SmallVector<ReassociationIndices> reassociation =1896          getOperandReassociation(indexingMap, collapsingInfo);1897      assert(1898          indexingMap.isProjectedPermutation() &&1899          "Expected indexing map to be a projected permutation for collapsing");1900      SmallVector<OpFoldResult> resultShape =1901          applyPermutationMap(indexingMap, ArrayRef(loopBound));1902      Value result;1903      if (isa<MemRefType>(collapsedOpResult.getType())) {1904        MemRefType expandShapeResultType = MemRefType::get(1905            originalResultType.getShape(), originalResultType.getElementType());1906        result = memref::ExpandShapeOp::create(1907            rewriter, loc, expandShapeResultType, collapsedOpResult,1908            reassociation, resultShape);1909      } else {1910        result = tensor::ExpandShapeOp::create(1911            rewriter, loc, originalResultType, collapsedOpResult, reassociation,1912            resultShape);1913      }1914      results.push_back(result);1915    } else {1916      results.push_back(collapsedOpResult);1917    }1918  }1919  return CollapseResult{results, collapsedOp};1920}1921 1922namespace {1923 1924/// Pattern to fuse a tensor.expand_shape op with its consumer generic op by1925/// contracting dimensions of the loop.1926class FoldWithProducerReshapeOpByCollapsing1927    : public OpRewritePattern<GenericOp> {1928public:1929  // TODO : support fusion with all linalg ops, not just generic.1930  FoldWithProducerReshapeOpByCollapsing(MLIRContext *context,1931                                        ControlFusionFn foldReshapes,1932                                        PatternBenefit benefit = 1)1933      : OpRewritePattern<GenericOp>(context, benefit),1934        controlFoldingReshapes(std::move(foldReshapes)) {}1935 1936  LogicalResult matchAndRewrite(GenericOp genericOp,1937                                PatternRewriter &rewriter) const override {1938    for (OpOperand &opOperand : genericOp->getOpOperands()) {1939      tensor::ExpandShapeOp reshapeOp =1940          opOperand.get().getDefiningOp<tensor::ExpandShapeOp>();1941      if (!reshapeOp)1942        continue;1943 1944      SmallVector<ReassociationIndices> collapsableIterationDims =1945          getCollapsableIterationSpaceDims(genericOp, &opOperand,1946                                           reshapeOp.getReassociationIndices());1947      if (collapsableIterationDims.empty() ||1948          !controlFoldingReshapes(&opOperand)) {1949        continue;1950      }1951 1952      std::optional<CollapseResult> collapseResult = collapseOpIterationDims(1953          genericOp, collapsableIterationDims, rewriter);1954      if (!collapseResult) {1955        return rewriter.notifyMatchFailure(1956            genericOp, "failed to do the fusion by collapsing transformation");1957      }1958 1959      rewriter.replaceOp(genericOp, collapseResult->results);1960      return success();1961    }1962    return failure();1963  }1964 1965private:1966  ControlFusionFn controlFoldingReshapes;1967};1968 1969/// Pattern to fold a tensor.collapse_shape op with its producer generic op1970/// by expanding the dimensionality of the loop in the producer op.1971struct FoldReshapeWithGenericOpByCollapsing1972    : public OpRewritePattern<tensor::CollapseShapeOp> {1973 1974  FoldReshapeWithGenericOpByCollapsing(MLIRContext *context,1975                                       ControlFusionFn foldReshapes,1976                                       PatternBenefit benefit = 1)1977      : OpRewritePattern<tensor::CollapseShapeOp>(context, benefit),1978        controlFoldingReshapes(std::move(foldReshapes)) {}1979 1980  LogicalResult matchAndRewrite(tensor::CollapseShapeOp reshapeOp,1981                                PatternRewriter &rewriter) const override {1982    // Fold only if all constraints of fusing with reshape by collapsing are1983    // met.1984    auto producerResult = dyn_cast<OpResult>(reshapeOp.getSrc());1985    if (!producerResult) {1986      return rewriter.notifyMatchFailure(reshapeOp,1987                                         "source not produced by an operation");1988    }1989 1990    // TODO : support fusion with all linalg producers, not just generic.1991    auto producer = dyn_cast<GenericOp>(producerResult.getOwner());1992    if (!producer) {1993      return rewriter.notifyMatchFailure(reshapeOp,1994                                         "producer not a generic op");1995    }1996 1997    SmallVector<ReassociationIndices> collapsableIterationDims =1998        getCollapsableIterationSpaceDims(1999            producer,2000            producer.getDpsInitOperand(producerResult.getResultNumber()),2001            reshapeOp.getReassociationIndices());2002    if (collapsableIterationDims.empty()) {2003      return rewriter.notifyMatchFailure(2004          reshapeOp, "failed preconditions of fusion with producer generic op");2005    }2006 2007    if (!controlFoldingReshapes(&reshapeOp.getSrcMutable())) {2008      return rewriter.notifyMatchFailure(reshapeOp,2009                                         "fusion blocked by control function");2010    }2011 2012    // Set the insertion point after `producer` because there could be uses2013    // of `producer` between it and the `tensor.collapse_shape` op.2014    rewriter.setInsertionPointAfter(producer);2015    std::optional<CollapseResult> collapseResult =2016        collapseOpIterationDims(producer, collapsableIterationDims, rewriter);2017    if (!collapseResult) {2018      return rewriter.notifyMatchFailure(2019          producer, "failed to do the fusion by collapsing transformation");2020    }2021 2022    rewriter.replaceOp(producer, collapseResult->results);2023    return success();2024  }2025 2026private:2027  ControlFusionFn controlFoldingReshapes;2028};2029 2030/// Computes the collapsed padding information for the given pad operation based2031/// on the provided collapsed shape and reassociation indices. Returns a2032/// PadDimInfo containing the low and high padding amounts and the collapsed2033/// shape for each dimension, or failure if the collapse is not possible.2034static FailureOr<PadDimInfo>2035computeCollapsedPadding(tensor::PadOp padOp,2036                        ArrayRef<ReassociationIndices> reassociations,2037                        PatternRewriter &rewriter) {2038  // If the padding value depends on the index values of the pad operation,2039  // then it may not be valid to collapse the dimensions, since it will change2040  // the index values on which the padding value depends. This is not currently2041  // supported by the pad collapsing patterns, but it could be implemented2042  // similarly to the collapsing of linalg.generic ops with linalg.index ops in2043  // the body, as is done in `generateCollapsedIndexingRegion`.2044  if (!padOp.getConstantPaddingValue())2045    return failure();2046 2047  // Collapsed dimensions cannot have padding because this can produce strided2048  // padding that isn't representable by a tensor.pad op. There are some special2049  // cases where it is possible (like collapsing unit dims), but supporting2050  // these cases is NYI, so disallow it for now.2051  ArrayRef<int64_t> low = padOp.getStaticLow();2052  ArrayRef<int64_t> high = padOp.getStaticHigh();2053  for (auto [idx, reInd] : llvm::enumerate(reassociations)) {2054    for (int64_t dim : reInd) {2055      if ((low[dim] != 0 || high[dim] != 0) && reInd.size() != 1)2056        return failure();2057    }2058  }2059 2060  // Initialize padding values for collapsed tensors with zeros2061  ArrayRef<int64_t> expandedPaddedShape = padOp.getType().getShape();2062  PadDimInfo padDimInfo;2063  padDimInfo.lowPad.assign(reassociations.size(), rewriter.getIndexAttr(0));2064  padDimInfo.highPad.assign(reassociations.size(), rewriter.getIndexAttr(0));2065 2066  // Update padding for dimensions that are not being collapsed, and compute2067  // the collapsed padded shape.2068  SmallVector<OpFoldResult> mixedLowPad(padOp.getMixedLowPad());2069  SmallVector<OpFoldResult> mixedHighPad(padOp.getMixedHighPad());2070  for (auto [idx, reInd] : llvm::enumerate(reassociations)) {2071    if (reInd.size() == 1) {2072      padDimInfo.lowPad[idx] = mixedLowPad[reInd[0]];2073      padDimInfo.highPad[idx] = mixedHighPad[reInd[0]];2074    }2075    SaturatedInteger collapsedSize = SaturatedInteger::wrap(1);2076    for (int64_t dim : reInd) {2077      collapsedSize =2078          collapsedSize * SaturatedInteger::wrap(expandedPaddedShape[dim]);2079    }2080    padDimInfo.paddedShape.push_back(collapsedSize.asInteger());2081  }2082 2083  return padDimInfo;2084}2085 2086class FoldPadWithProducerReshapeOpByCollapsing2087    : public OpRewritePattern<tensor::PadOp> {2088public:2089  FoldPadWithProducerReshapeOpByCollapsing(MLIRContext *context,2090                                           ControlFusionFn foldReshapes,2091                                           PatternBenefit benefit = 1)2092      : OpRewritePattern<tensor::PadOp>(context, benefit),2093        controlFoldingReshapes(std::move(foldReshapes)) {}2094 2095  LogicalResult matchAndRewrite(tensor::PadOp padOp,2096                                PatternRewriter &rewriter) const override {2097    tensor::ExpandShapeOp reshapeOp =2098        padOp.getSource().getDefiningOp<tensor::ExpandShapeOp>();2099    if (!reshapeOp)2100      return failure();2101 2102    if (!controlFoldingReshapes(&padOp.getSourceMutable())) {2103      return rewriter.notifyMatchFailure(padOp,2104                                         "fusion blocked by control function");2105    }2106 2107    SmallVector<ReassociationIndices> reassociations =2108        reshapeOp.getReassociationIndices();2109    FailureOr<PadDimInfo> maybeCollapsedPadding =2110        computeCollapsedPadding(padOp, reassociations, rewriter);2111    if (failed(maybeCollapsedPadding))2112      return failure();2113    PadDimInfo &collapsedPadding = maybeCollapsedPadding.value();2114 2115    SmallVector<OpFoldResult> expandedPaddedSizes =2116        reshapeOp.getMixedOutputShape();2117    AffineExpr d0, d1, d2;2118    bindDims(rewriter.getContext(), d0, d1, d2);2119    auto addMap = AffineMap::get(3, 0, {d0 + d1 + d2});2120    Location loc = reshapeOp->getLoc();2121    for (auto [reInd, l, h] :2122         llvm::zip_equal(reassociations, collapsedPadding.lowPad,2123                         collapsedPadding.highPad)) {2124      if (reInd.size() == 1) {2125        expandedPaddedSizes[reInd[0]] = affine::makeComposedFoldedAffineApply(2126            rewriter, loc, addMap, {l, h, expandedPaddedSizes[reInd[0]]});2127      }2128    }2129 2130    RankedTensorType collapsedPaddedType =2131        padOp.getType().clone(collapsedPadding.paddedShape);2132    auto newPadOp = tensor::PadOp::create(2133        rewriter, loc, collapsedPaddedType, reshapeOp.getSrc(),2134        collapsedPadding.lowPad, collapsedPadding.highPad,2135        padOp.getConstantPaddingValue(), padOp.getNofold());2136 2137    rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(2138        padOp, padOp.getResultType(), newPadOp.getResult(), reassociations,2139        expandedPaddedSizes);2140 2141    return success();2142  }2143 2144private:2145  ControlFusionFn controlFoldingReshapes;2146};2147 2148class FoldReshapeWithProducerPadOpByCollapsing2149    : public OpRewritePattern<tensor::CollapseShapeOp> {2150public:2151  FoldReshapeWithProducerPadOpByCollapsing(MLIRContext *context,2152                                           ControlFusionFn foldReshapes,2153                                           PatternBenefit benefit = 1)2154      : OpRewritePattern<tensor::CollapseShapeOp>(context, benefit),2155        controlFoldingReshapes(std::move(foldReshapes)) {}2156 2157  LogicalResult matchAndRewrite(tensor::CollapseShapeOp reshapeOp,2158                                PatternRewriter &rewriter) const override {2159    tensor::PadOp padOp = reshapeOp.getSrc().getDefiningOp<tensor::PadOp>();2160    if (!padOp)2161      return failure();2162 2163    if (!controlFoldingReshapes(&reshapeOp.getSrcMutable())) {2164      return rewriter.notifyMatchFailure(padOp,2165                                         "fusion blocked by control function");2166    }2167 2168    SmallVector<ReassociationIndices> reassociations =2169        reshapeOp.getReassociationIndices();2170    RankedTensorType collapsedPaddedType = reshapeOp.getResultType();2171    FailureOr<PadDimInfo> maybeCollapsedPadding =2172        computeCollapsedPadding(padOp, reassociations, rewriter);2173    if (failed(maybeCollapsedPadding))2174      return failure();2175    PadDimInfo &collapsedPadding = maybeCollapsedPadding.value();2176 2177    Location loc = reshapeOp->getLoc();2178    auto newCollapseOp = tensor::CollapseShapeOp::create(2179        rewriter, loc, padOp.getSource(), reassociations);2180 2181    auto newPadOp = tensor::PadOp::create(2182        rewriter, loc, collapsedPaddedType, newCollapseOp.getResult(),2183        collapsedPadding.lowPad, collapsedPadding.highPad,2184        padOp.getConstantPaddingValue(), padOp.getNofold());2185 2186    rewriter.replaceOp(reshapeOp, newPadOp.getResult());2187    return success();2188  }2189 2190private:2191  ControlFusionFn controlFoldingReshapes;2192};2193 2194/// Pattern to collapse dimensions.2195template <typename LinalgType>2196class CollapseLinalgDimensions : public OpRewritePattern<LinalgType> {2197public:2198  CollapseLinalgDimensions(MLIRContext *context,2199                           GetCollapsableDimensionsFn collapseDimensions,2200                           PatternBenefit benefit = 1)2201      : OpRewritePattern<LinalgType>(context, benefit),2202        controlCollapseDimension(std::move(collapseDimensions)) {}2203 2204  LogicalResult matchAndRewrite(LinalgType op,2205                                PatternRewriter &rewriter) const override {2206    SmallVector<ReassociationIndices> collapsableIterationDims =2207        controlCollapseDimension(op);2208    if (collapsableIterationDims.empty())2209      return failure();2210 2211    // Check if the specified list of dimensions to collapse is a valid list.2212    if (!areDimSequencesPreserved(op.getIndexingMapsArray(),2213                                  collapsableIterationDims)) {2214      return rewriter.notifyMatchFailure(2215          op, "specified dimensions cannot be collapsed");2216    }2217 2218    std::optional<CollapseResult> collapseResult =2219        collapseOpIterationDims(op, collapsableIterationDims, rewriter);2220    if (!collapseResult) {2221      return rewriter.notifyMatchFailure(op, "failed to collapse dimensions");2222    }2223    rewriter.replaceOp(op, collapseResult->results);2224    return success();2225  }2226 2227private:2228  GetCollapsableDimensionsFn controlCollapseDimension;2229};2230 2231} // namespace2232 2233//===---------------------------------------------------------------------===//2234// Methods and patterns that fuse constants with linalg.generic operations.2235//===---------------------------------------------------------------------===//2236 2237namespace {2238/// Pattern to fold a generic op with a splat constant/scalar constant. Does not2239/// handle cases where the constant is not single-valued.2240class FoldScalarOrSplatConstant : public OpRewritePattern<GenericOp> {2241public:2242  FoldScalarOrSplatConstant(MLIRContext *context, PatternBenefit benefit = 1)2243      : OpRewritePattern<GenericOp>(context, benefit) {}2244 2245  LogicalResult matchAndRewrite(GenericOp genericOp,2246                                PatternRewriter &rewriter) const override {2247    if (!genericOp.hasPureTensorSemantics())2248      return failure();2249    for (OpOperand *opOperand : genericOp.getDpsInputOperands()) {2250      Operation *def = opOperand->get().getDefiningOp();2251      TypedAttr constantAttr;2252      auto isScalarOrSplatConstantOp = [&constantAttr](Operation *def) -> bool {2253        {2254          DenseElementsAttr splatAttr;2255          if (matchPattern(def, m_Constant<DenseElementsAttr>(&splatAttr)) &&2256              splatAttr.isSplat() &&2257              splatAttr.getType().getElementType().isIntOrFloat()) {2258            constantAttr = splatAttr.getSplatValue<TypedAttr>();2259            return true;2260          }2261        }2262        {2263          IntegerAttr intAttr;2264          if (matchPattern(def, m_Constant<IntegerAttr>(&intAttr))) {2265            constantAttr = intAttr;2266            return true;2267          }2268        }2269        {2270          FloatAttr floatAttr;2271          if (matchPattern(def, m_Constant<FloatAttr>(&floatAttr))) {2272            constantAttr = floatAttr;2273            return true;2274          }2275        }2276        return false;2277      };2278 2279      auto resultValue = dyn_cast<OpResult>(opOperand->get());2280      if (!def || !resultValue || !isScalarOrSplatConstantOp(def))2281        continue;2282 2283      // The operands and the indexing_maps of the fused operation the same as2284      // the operands and indexing_maps of the generic operations with the2285      // values at the constant index dropped.2286      SmallVector<AffineMap> fusedIndexMaps;2287      SmallVector<Value> fusedOperands;2288      SmallVector<Location> fusedLocs{genericOp.getLoc()};2289      fusedIndexMaps.reserve(genericOp->getNumOperands());2290      fusedOperands.reserve(genericOp.getNumDpsInputs());2291      fusedLocs.reserve(fusedLocs.size() + genericOp.getNumDpsInputs());2292      for (OpOperand *inputOperand : genericOp.getDpsInputOperands()) {2293        if (inputOperand == opOperand)2294          continue;2295        Value inputValue = inputOperand->get();2296        fusedIndexMaps.push_back(2297            genericOp.getMatchingIndexingMap(inputOperand));2298        fusedOperands.push_back(inputValue);2299        fusedLocs.push_back(inputValue.getLoc());2300      }2301      for (OpOperand &outputOperand : genericOp.getDpsInitsMutable())2302        fusedIndexMaps.push_back(2303            genericOp.getMatchingIndexingMap(&outputOperand));2304 2305      // Check if the operation shapes to loops map is computable.2306      if (!inversePermutation(2307              concatAffineMaps(fusedIndexMaps, rewriter.getContext()))) {2308        return rewriter.notifyMatchFailure(2309            genericOp, "fused op loop bound computation failed");2310      }2311 2312      // Create a constant scalar value from the splat constant.2313      Value scalarConstant =2314          arith::ConstantOp::create(rewriter, def->getLoc(), constantAttr);2315 2316      SmallVector<Value> outputOperands = genericOp.getOutputs();2317      auto fusedOp =2318          GenericOp::create(rewriter, rewriter.getFusedLoc(fusedLocs),2319                            genericOp->getResultTypes(),2320                            /*inputs=*/fusedOperands,2321                            /*outputs=*/outputOperands,2322                            rewriter.getAffineMapArrayAttr(fusedIndexMaps),2323                            genericOp.getIteratorTypes(),2324                            /*doc=*/nullptr,2325                            /*library_call=*/nullptr);2326 2327      // Map the block argument corresponding to the replaced argument with the2328      // scalar constant.2329      Region &region = genericOp->getRegion(0);2330      Block &entryBlock = *region.begin();2331      IRMapping mapping;2332      mapping.map(entryBlock.getArgument(opOperand->getOperandNumber()),2333                  scalarConstant);2334      Region &fusedRegion = fusedOp->getRegion(0);2335      rewriter.cloneRegionBefore(region, fusedRegion, fusedRegion.begin(),2336                                 mapping);2337      rewriter.replaceOp(genericOp, fusedOp->getResults());2338      return success();2339    }2340    return failure();2341  }2342};2343 2344} // namespace2345 2346//===---------------------------------------------------------------------===//2347// Miscellaneous patterns that help fusion.2348//===---------------------------------------------------------------------===//2349 2350namespace {2351/// Forces `outs` operands of linalg operations to use `tensor.empty` if the2352/// value of the `outs` operand is not used within the op.  This is only2353/// implemented for `linalg.generic` operations for now, but should hold for all2354/// linalg structured ops.2355struct RemoveOutsDependency : public OpRewritePattern<GenericOp> {2356  using OpRewritePattern<GenericOp>::OpRewritePattern;2357 2358  LogicalResult matchAndRewrite(GenericOp op,2359                                PatternRewriter &rewriter) const override {2360    rewriter.startOpModification(op);2361    bool modifiedOutput = false;2362    Location loc = op.getLoc();2363    for (OpOperand &opOperand : op.getDpsInitsMutable()) {2364      if (!op.payloadUsesValueFromOperand(&opOperand)) {2365        Value operandVal = opOperand.get();2366        auto operandType = dyn_cast<RankedTensorType>(operandVal.getType());2367        if (!operandType)2368          continue;2369 2370        // If outs is sparse, leave it to the sparsifier.2371        if (sparse_tensor::getSparseTensorEncoding(operandVal.getType()))2372          continue;2373 2374        // If outs is already an `empty` operation, nothing to do.2375        auto definingOp = operandVal.getDefiningOp<tensor::EmptyOp>();2376        if (definingOp)2377          continue;2378        modifiedOutput = true;2379        SmallVector<OpFoldResult> mixedSizes =2380            tensor::getMixedSizes(rewriter, loc, operandVal);2381        Value emptyTensor = tensor::EmptyOp::create(2382            rewriter, loc, mixedSizes, operandType.getElementType());2383        op->setOperand(opOperand.getOperandNumber(), emptyTensor);2384      }2385    }2386    if (!modifiedOutput) {2387      rewriter.cancelOpModification(op);2388      return failure();2389    }2390    rewriter.finalizeOpModification(op);2391    return success();2392  }2393};2394 2395/// Fold linalg.fill into linalg.generic2396struct FoldFillWithGenericOp : public OpRewritePattern<GenericOp> {2397  using OpRewritePattern<GenericOp>::OpRewritePattern;2398 2399  LogicalResult matchAndRewrite(GenericOp genericOp,2400                                PatternRewriter &rewriter) const override {2401    if (!genericOp.hasPureTensorSemantics())2402      return failure();2403    bool fillFound = false;2404    Block &payload = genericOp.getRegion().front();2405    for (OpOperand *opOperand : genericOp.getDpsInputOperands()) {2406      if (!genericOp.payloadUsesValueFromOperand(opOperand))2407        continue;2408      FillOp fillOp = opOperand->get().getDefiningOp<FillOp>();2409      if (!fillOp)2410        continue;2411      fillFound = true;2412      Value fillVal = fillOp.value();2413      auto resultType =2414          cast<RankedTensorType>(fillOp.result().getType()).getElementType();2415      Value convertedVal =2416          convertScalarToDtype(rewriter, fillOp.getLoc(), fillVal, resultType,2417                               /*isUnsignedCast =*/false);2418      rewriter.replaceAllUsesWith(2419          payload.getArgument(opOperand->getOperandNumber()), convertedVal);2420    }2421    return success(fillFound);2422  }2423};2424} // namespace2425 2426void mlir::linalg::populateFoldReshapeOpsByExpansionPatterns(2427    RewritePatternSet &patterns,2428    const ControlFusionFn &controlFoldingReshapes) {2429  patterns.add<FoldReshapeWithGenericOpByExpansion>(patterns.getContext(),2430                                                    controlFoldingReshapes);2431  patterns.add<FoldPadWithProducerReshapeOpByExpansion>(patterns.getContext(),2432                                                        controlFoldingReshapes);2433  patterns.add<FoldReshapeWithProducerPadOpByExpansion>(patterns.getContext(),2434                                                        controlFoldingReshapes);2435  patterns.add<FoldWithProducerReshapeOpByExpansion>(patterns.getContext(),2436                                                     controlFoldingReshapes);2437}2438 2439void mlir::linalg::populateFoldReshapeOpsByCollapsingPatterns(2440    RewritePatternSet &patterns,2441    const ControlFusionFn &controlFoldingReshapes) {2442  patterns.add<FoldWithProducerReshapeOpByCollapsing>(patterns.getContext(),2443                                                      controlFoldingReshapes);2444  patterns.add<FoldPadWithProducerReshapeOpByCollapsing>(2445      patterns.getContext(), controlFoldingReshapes);2446  patterns.add<FoldReshapeWithProducerPadOpByCollapsing>(2447      patterns.getContext(), controlFoldingReshapes);2448  patterns.add<FoldReshapeWithGenericOpByCollapsing>(patterns.getContext(),2449                                                     controlFoldingReshapes);2450}2451 2452void mlir::linalg::populateElementwiseOpsFusionPatterns(2453    RewritePatternSet &patterns,2454    const ControlFusionFn &controlElementwiseOpsFusion) {2455  auto *context = patterns.getContext();2456  patterns.add<FuseElementwiseOps>(context, controlElementwiseOpsFusion);2457  patterns.add<FoldFillWithGenericOp, FoldScalarOrSplatConstant,2458               RemoveOutsDependency>(context);2459  // Add the patterns that clean up dead operands and results.2460  populateEraseUnusedOperandsAndResultsPatterns(patterns);2461}2462 2463void mlir::linalg::populateCollapseDimensions(2464    RewritePatternSet &patterns,2465    const GetCollapsableDimensionsFn &controlCollapseDimensions) {2466  patterns.add<CollapseLinalgDimensions<linalg::GenericOp>,2467               CollapseLinalgDimensions<linalg::CopyOp>>(2468      patterns.getContext(), controlCollapseDimensions);2469}2470 2471//===---------------------------------------------------------------------===//2472// Passes2473//===---------------------------------------------------------------------===//2474 2475namespace {2476 2477/// Pass that fuses generic ops on tensors. Used only for testing.2478// TODO(ravishankarm): This pass is to be deprecated. The efficacy of the2479// patterns added here heavily depends on the cost function used. Having an2480// opinionated pass of this form is not recommended. Deprecate this pass in2481// favor of test passes that check the functionality of each of the patterns2482// added here individually.2483struct LinalgElementwiseOpFusionPass2484    : public impl::LinalgElementwiseOpFusionPassBase<2485          LinalgElementwiseOpFusionPass> {2486  using impl::LinalgElementwiseOpFusionPassBase<2487      LinalgElementwiseOpFusionPass>::LinalgElementwiseOpFusionPassBase;2488  void runOnOperation() override {2489    Operation *op = getOperation();2490    MLIRContext *context = op->getContext();2491    RewritePatternSet patterns(context);2492 2493    // Add folding with reshape by expansion patterns.2494    ControlFusionFn defaultControlFn = [](OpOperand *fusedOperand) {2495      Operation *producer = fusedOperand->get().getDefiningOp();2496      return producer && producer->hasOneUse();2497    };2498 2499    // Add elementwise op fusion patterns.2500    populateElementwiseOpsFusionPatterns(patterns, defaultControlFn);2501    populateFoldReshapeOpsByExpansionPatterns(patterns, defaultControlFn);2502    tensor::populateBubbleUpExpandShapePatterns(patterns);2503 2504    // General canonicalization patterns.2505    affine::AffineApplyOp::getCanonicalizationPatterns(patterns, context);2506    GenericOp::getCanonicalizationPatterns(patterns, context);2507    tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);2508    tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);2509    context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(2510        patterns);2511 2512    // Add constant folding patterns.2513    populateConstantFoldLinalgOperations(patterns, defaultControlFn);2514 2515    // Use TopDownTraversal for compile time reasons.2516    (void)applyPatternsGreedily(op, std::move(patterns),2517                                GreedyRewriteConfig().setUseTopDownTraversal());2518  }2519};2520 2521} // namespace2522