TileUsingInterface.h

Go to the documentation of this file.

//===- TileUsingInterface.h - Tiling ops using TilingInterface --*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_DIALECT_SCF_TRANSFORMS_TILEUSINGINTERFACE_H
#define MLIR_DIALECT_SCF_TRANSFORMS_TILEUSINGINTERFACE_H
 
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "mlir/Rewrite/FrozenRewritePatternSet.h"
 
#include <deque>
 
namespace mlir {
class Operation;
class RewriterBase;
class TilingInterface;
} // namespace mlir
 
namespace mlir {
namespace scf {
 
using SCFTileSizeComputationFunction =
    std::function<SmallVector<OpFoldResult>(OpBuilder &, Operation *)>;
 
/// Computation function returning, for the op currently being tiled or fused,
/// the per-iteration-domain-dimension `InnerTileAlignment` array.
///`tileSizes` holds the given tile sizes for pure tiling
/// and is empty for fusion; `slices` holds the `tensor.extract_slice` (producer
/// fusion) or `tensor.insert_slice`/`tensor.parallel_insert_slice`s (consumer
/// fusion) being fused, and is empty for pure tiling. Only pack/unpack
/// implementations consult the result; every other op ignores it.
using InnerTileAlignmentFnTy = std::function<SmallVector<InnerTileAlignment>(
    TilingInterface op, ArrayRef<OpFoldResult> tileSizes,
    ArrayRef<Operation *> slices)>;
 
/// Options to use to control tiling.
struct SCFTilingOptions {
  /// Specify which loop construct to use for tile and fuse.
  enum class LoopType { ForOp, ForallOp, CustomOp };
  LoopType loopType = LoopType::ForOp;
  SCFTilingOptions &setLoopType(LoopType type) {
    loopType = type;
    return *this;
  }
 
  /// Computation function that returns the tile sizes to use for each loop.
  /// Returning a tile size of zero implies no tiling for that loop. If the
  /// size of the returned vector is smaller than the number of loops, the inner
  /// loops are not tiled. If the size of the returned vector is larger, then
  /// the vector is truncated to number of loops.
  SCFTileSizeComputationFunction tileSizeComputationFunction = nullptr;
 
  SCFTilingOptions &
  setTileSizeComputationFunction(SCFTileSizeComputationFunction fun) {
    tileSizeComputationFunction = std::move(fun);
    return *this;
  }
  /// Convenience function to set the `tileSizeComputationFunction` to a
  /// function that computes tile sizes at the point they are needed. Allows
  /// proper interaction with folding.
  SCFTilingOptions &setTileSizes(ArrayRef<OpFoldResult> tileSizes);
 
  /// The interchange vector to reorder the tiled loops.
  SmallVector<int64_t> interchangeVector = {};
  SCFTilingOptions &setInterchange(ArrayRef<int64_t> interchange) {
    interchangeVector = llvm::to_vector(interchange);
    return *this;
  }
 
  /// Optional control function returning, per tiled/fused op, the inner tile
  /// alignment hints.
  InnerTileAlignmentFnTy innerTileAlignmentFn = nullptr;
  SCFTilingOptions &setInnerTileAlignmentFn(InnerTileAlignmentFnTy fn) {
    innerTileAlignmentFn = std::move(fn);
    return *this;
  }
 
  /// Sets a constant `innerTileAlignmentFn` that returns `alignments` for every
  /// op, ignoring which op is being tiled or fused. Use when a single fixed
  /// array is correct for the whole tiling/fusion (e.g. tiling a single op).
  SCFTilingOptions &
  setInnerTileAlignments(ArrayRef<InnerTileAlignment> alignments) {
    SmallVector<InnerTileAlignment> fixed = llvm::to_vector(alignments);
    innerTileAlignmentFn =
        [fixed = std::move(fixed)](TilingInterface, ArrayRef<OpFoldResult>,
                                   ArrayRef<Operation *>) { return fixed; };
    return *this;
  }
 
  //-------------------------------------------------------------------------//
  // Options related to tiling using `scf.forall`.
  //-------------------------------------------------------------------------//
 
  /// Computation function that returns the number of threads to use for
  /// each loop. Returning a num threads of zero implies no tiling for that
  /// loop. If the size of the returned vector is smaller than the number of
  /// loops, the inner loops are not tiled. If the size of the returned vector
  /// is larger, then the vector is truncated to number of loops. Note: This
  /// option is only supported with loopType set to `LoopType::ForallOp`. If the
  /// tile size function is not specified while the num threads computation is,
  /// then the tile size is determined automatically to map at most one tile per
  /// thread.
  SCFTileSizeComputationFunction numThreadsComputationFunction = nullptr;
 
  SCFTilingOptions &
  setNumThreadsComputationFunction(SCFTileSizeComputationFunction fun) {
    numThreadsComputationFunction = std::move(fun);
    return *this;
  }
  /// Convenience function to set the `numThreadsComputationFunction` to a
  /// function that computes num threads at the point they are needed.
  SCFTilingOptions &setNumThreads(ArrayRef<OpFoldResult> numThreads);
 
  /// Specify mapping of loops to devices. This is only respected when the loop
  /// constructs support such a mapping (like `scf.forall`). Will be ignored
  /// when using loop constructs that dont support such a mapping (like
  /// `scf.for`)
  SmallVector<Attribute> mappingVector = {};
  SCFTilingOptions &setMapping(ArrayRef<Attribute> mapping) {
    mappingVector = llvm::to_vector(mapping);
    return *this;
  }
 
  //-------------------------------------------------------------------------//
  // Options related reduction tiling
  //-------------------------------------------------------------------------//
 
  /// Specify how reduction dimensions should be tiled.
  ReductionTilingStrategy reductionStrategy =
      ReductionTilingStrategy::FullReduction;
  SCFTilingOptions &
  setReductionTilingStrategy(ReductionTilingStrategy strategy) {
    reductionStrategy = strategy;
    return *this;
  }
 
  /// Specify the reduction dimensions to be tiled. Note that this needs to be
  /// specified. If left unspecified, then none of the reduction dimensions are
  /// tiled.
  SetVector<unsigned> reductionDims;
  SCFTilingOptions &setReductionDims(ArrayRef<unsigned> dims) {
    reductionDims.clear();
    reductionDims.insert(dims.begin(), dims.end());
    return *this;
  }
 
  //-------------------------------------------------------------------------//
  // Options related to tiling using custom loop.
  //-------------------------------------------------------------------------//
 
  // For generating the inter-tile loops using a custom loop, two callback
  // functions are needed
  // 1. That generates the "loop header", i.e. the loop that iterates over the
  //    different tiles.
  // 2. That generates the loop terminator
  //
  // For `scf.forall` case the call back to generate loop header would generate
  //
  // ```mlir
  // scf.forall (...) = ... {
  //   ..
  // }
  // ```
  //
  // and the call back to generate the loop terminator would generate the
  // `scf.in_parallel` region
  //
  // ```mlir
  // scf.forall (...) = ... {
  //   scf.in_parallel {
  //      tensor.parallel_insert_slice ...
  //   }
  // }
  // ```
  //
 
  // Information that is to be returned by loop header callback needed for the
  // rest of the tiled codegeneration.
  // - `loops`: The generated loops
  // - `tileOffset`: The values that represent the offset of the iteration space
  //                 tile.
  // - `tileSizes` : The values that represent the size of the iteration space
  //                 tile.
  // - `destinationTensors` : The tensors to use as destinations during tiling.
  struct CustomLoopHeaderInfo {
    SmallVector<LoopLikeOpInterface> loops;
    SmallVector<OpFoldResult> tileOffset;
    SmallVector<OpFoldResult> tileSizes;
    SmallVector<Value> destinationTensors;
  };
 
  // Type of the callback function that generates the loop headers.
  // - `loopRanges` : Values that represent the full size of the iteration space
  //                  being tiled.
  // - `givenTileSizes` : The tile sizes that are to be used to tile the
  //                      iteration space.
  // - `destinationTensors` : The tensors to use as destinations for the results
  //                          of the tiled loop for loops that implement
  //                          `DestinationStyleOpInterface`.
  // Returns the `CustomLoopHeaderInfo` object (described above). it is expected
  // that this function sets the insertion point of `rewriter` to the program
  // point where the intra-tile loop computation is to be generated.
  using GenerateLoopHeaderFn = std::function<FailureOr<CustomLoopHeaderInfo>(
      RewriterBase &rewriter, Location loc, ArrayRef<Range> loopRanges,
      ArrayRef<OpFoldResult> givenTileSizes, ValueRange destinationTensors)>;
 
  // Type of the callback function that generates the loop terminator.
  // - `loops` : generated loops from the GenerateLoopHeaderFn callback
  // - `tiledResults` : Tiles of the result computed for the iteration space
  //                    tile.
  // - `resultOffsets` : For each of the `tiledResults`, the offset at which
  //                     the result tile is to be "inserted" back into the
  //                     destination tensor.
  // - `resultSizes` : For each of the `tiledResults`, the size of the result
  //                   tile that is to be "inserted" back into the destination
  //                   tensor.
  // Returns the `CustomLoopHeaderInfo` object (described above)
  using GenerateLoopTerminatorFn = std::function<LogicalResult(
      RewriterBase &rewriter, Location loc, ArrayRef<LoopLikeOpInterface> loops,
      ValueRange tiledResults,
      ArrayRef<SmallVector<OpFoldResult>> resultOffsets,
      ArrayRef<SmallVector<OpFoldResult>> resultSizes,
      ValueRange destinationTensors)>;
 
  // Callback function to generate the inter-tile loop header.
  GenerateLoopHeaderFn generateLoopHeaderFn = nullptr;
  // Callback function to generate the inter-tile loop terminator.
  GenerateLoopTerminatorFn generateLoopTerminatorFn = nullptr;
  // Helper function to set the callbacks for inter-tile loop header and
  // terminator functions when using a custom operation for the loop.
  SCFTilingOptions &
  setCustomLoopGenerationFns(GenerateLoopHeaderFn headerFn,
                             GenerateLoopTerminatorFn terminatorFn) {
    generateLoopHeaderFn = std::move(headerFn);
    generateLoopTerminatorFn = std::move(terminatorFn);
    return *this;
  }
};
 
/// Transformation information returned after tiling.
struct SCFTilingResult {
  /// Tiled operations that are generated during tiling. The order does not
  /// matter except the last op. The replacements are expected to be the results
  /// of the last op.
  SmallVector<Operation *> tiledOps;
  /// The initial destination values passed to the tiled operations.
  SmallVector<Value> initialValues;
  /// The `scf.for` operations that iterate over the tiles.
  SmallVector<LoopLikeOpInterface> loops;
  /// Values to use as replacements for the untiled op. Is the same size as the
  /// number of results of the untiled op.
  SmallVector<Value> replacements;
  /// Slices generated after tiling that can be used for fusing with the tiled
  /// producer.
  SmallVector<Operation *> generatedSlices;
  /// In cases where there as an additional merge step after tiling
  /// return the merged ops after tiling. This list is empty when reduction
  /// tiling strategy is
  /// `scf::SCFTilingOptions::ReductionTilingStrategy::FullReduction.
  SmallVector<Operation *> mergeOps;
};
 
/// Method to tile an op that implements the `TilingInterface` using
/// `scf.for` for iterating over the tiles.
FailureOr<SCFTilingResult> tileUsingSCF(RewriterBase &rewriter,
                                        TilingInterface op,
                                        const SCFTilingOptions &options);
 
/// Options used to control tile + fuse.
struct SCFTileAndFuseOptions {
  /// The tiling options used to control the tiling of the consumer.
  SCFTilingOptions tilingOptions;
  SCFTileAndFuseOptions &setTilingOptions(SCFTilingOptions options) {
    tilingOptions = options;
    return *this;
  }
 
  /// Control function to check if a slice needs to be fused or not,
  /// The control function receives
  /// 1) the slice along which fusion is to be done,
  /// 2) the producer value that is to be fused
  /// 3) a boolean value set to `true` if the fusion is from
  ///    a destination operand.
  /// The control function returns an `std::optiona<ControlFnResult>`.
  /// If the return value is `std::nullopt`, that implies no fusion
  /// is to be performed along that slice.
  struct ControlFnResult {
    /// Set to true if the loop nest has to return a replacement value
    /// for the fused producer.
    bool yieldProducerReplacement = false;
  };
  using ControlFnTy = std::function<std::optional<ControlFnResult>(
      tensor::ExtractSliceOp candidateSliceOp, OpResult originalProducer,
      bool isDestinationOperand)>;
  /// The default control function implements greedy fusion without yielding
  /// a replacement for any of the fused results.
  ControlFnTy fusionControlFn = [](tensor::ExtractSliceOp, OpResult,
                                   bool) -> std::optional<ControlFnResult> {
    return ControlFnResult{};
  };
  SCFTileAndFuseOptions &setFusionControlFn(ControlFnTy controlFn) {
    fusionControlFn = controlFn;
    return *this;
  }
 
  /// An optional set of rewrite patterns to apply to the results of tiling
  /// before fusion. This will track deleted and newly inserted
  /// `tensor.extract_slice` ops and update the worklist.
  std::optional<FrozenRewritePatternSet> cleanupPatterns = std::nullopt;
};
 
/// Result of fusing the producer of the source of a `tensor.extract_slice`.
struct SCFFuseProducerOfSliceResult {
  OpResult origProducer;       // Original untiled producer.
  Value tiledAndFusedProducer; // Tile and fused producer value.
  SmallVector<Operation *> tiledOps;
  SmallVector<Operation *> generatedSlices;
};
/// Fuse the producer of the source of `candidateSliceOp` by computing the
/// required slice of the producer in-place.  Note that the method
/// replaces the uses of `candidateSliceOp` with the tiled and fused producer
/// value but does not delete the slice operation.
///
/// When the fused producer is a `linalg.pack`/`linalg.unpack`, `fn` (if
/// non-null) is invoked with the producer and `candidateSliceOp` to obtain the
/// inner-tile alignment hint in the producer's own iteration domain.
std::optional<SCFFuseProducerOfSliceResult>
tileAndFuseProducerOfSlice(RewriterBase &rewriter,
                           tensor::ExtractSliceOp candidateSliceOp,
                           MutableArrayRef<LoopLikeOpInterface> loops,
                           const InnerTileAlignmentFnTy &fn = nullptr);
 
/// Reconstruct the fused producer from within the tiled-and-fused code. Based
/// on the slice of the producer computed in place it is possible that within
/// the loop nest same slice of the producer is computed multiple times. It is
/// in general not possible to recompute the value of the fused producer from
/// the tiled loop code in such cases. For the cases where no slice of the
/// producer is computed in a redundant fashion it is possible to reconstruct
/// the value of the original producer from within the tiled loop. It is upto
/// the caller to ensure that the producer is not computed redundantly within
/// the tiled loop nest. For example, consider
///
/// ```mlir
/// %0 = linalg.matmul ins(...) outs(...) -> tensor<?x?xf32>
/// %1 = linalg.matmul ins(%0, ..) outs(...) -> tensor<?x?x?f32>
/// ```
///
/// If `%1` is tiled in a 2D fashion and `%0` is fused with it, the resulting IR
/// is,
///
/// ```mlir
/// %t1_0 = scf.for .... iter_args(%arg0 = ...) {
///   %t1_1 = scf.for ... iter_args(%arg1 = %arg0) {
///     ...
///     %t1_2 = linalg.matmul ins(...) outs(...) -> tensor<?x?xf32>
///     %t1_3 = linalg.matmul ins(%t1_2, ...)
///     %t1_4 = tensor.insert_slice %t1_3 into %arg1 ...
///     scf.yield %t1_4
///   }
///   scf.yield %t1_1
/// }
/// ```
///
/// Here `%t1_2` is the same for all iterations of the inner `scf.for`. Instead
/// if `%1` were tiled only along the rows, the resultant code would be
///
/// ```mlir
/// %t2_0 = scf.for .... iter_args(%arg0 = ...) {
///   ...
///   %t2_1 = linalg.matmul ins(...) outs(...) -> tensor<?x?xf32>
///   %t2_2 = linalg.matmul ins(%t2_1, ...)
///   %t2_3 = tensor.insert_slice %t2_2 into %arg0 ...
///   scf.yield %t2_3
/// }
/// ```
///
/// Here there is no intersection in the different slices of `%t2_1` computed
/// across iterations of the `scf.for`. In such cases, the value of the original
/// `%0` can be reconstructed from within the loop body. This is useful in cases
/// where `%0` had other uses as well. If not reconstructed from within the loop
/// body, uses of `%0` could not be replaced, making it still live and the
/// fusion immaterial.
///
/// The @param `yieldResultNumber` decides which result would be yield. If not
/// given, yield all `opResult` of fused producer.
///
/// The method returns the list of new slices added during the process (which
/// can be used to fuse along).
FailureOr<SmallVector<Operation *>> yieldReplacementForFusedProducer(
    RewriterBase &rewriter, tensor::ExtractSliceOp sliceOp,
    scf::SCFFuseProducerOfSliceResult fusedProducerInfo,
    MutableArrayRef<LoopLikeOpInterface> loops,
    ArrayRef<unsigned> yieldResultNumber = ArrayRef<unsigned>{});
 
/// Transformation information returned after tile and fuse.
struct SCFTileAndFuseResult {
  /// List of untiled operations that were fused with the tiled consumer.
  llvm::SetVector<Operation *> fusedProducers;
  /// List of tiled and fused operations generated. The first element is always
  /// the tiled version of the original consumer operation processed by
  /// `tileConsumerAndFuseProducersUsingSCF`, followed by any operations that
  /// were fused with it.
  llvm::SetVector<Operation *> tiledAndFusedOps;
  /// The `scf.for` operations that iterate over the tiles.
  SmallVector<LoopLikeOpInterface> loops;
  /// The replacement values to use for the tiled and fused operations.
  llvm::DenseMap<Value, Value> replacements;
};
 
/// Method to tile and fuse a sequence of operations, by tiling the consumer
/// and fusing its producers. Note that this assumes that it is valid to
/// tile+fuse the producer into the innermost tiled loop. Its up to the caller
/// to ensure that the tile sizes provided make this fusion valid.
///
/// For example, for the following sequence
///
/// ```mlir
/// %0 =
/// %1 = linalg.fill ... outs(%0 : ... )
/// %2 = linalg.matmul ... outs(%1 : ...) ...
/// ```
///
/// it is legal to fuse the fill with the matmul only if the matmul is tiled
/// along the parallel dimensions and not the reduction dimension, i.e. the tile
/// size for the reduction dimension should be 0. The resulting fused
/// transformation is
///
/// ```mlir
/// %1 = scf.for ... iter_args(%arg0 = %0)
///   %2 = tensor.extract_slice %arg0
///   %3 = linalg.fill .. outs(%2 : ... )
///   %4 = linalg.matmul .. outs(%3 : ...)
/// }
/// ```
FailureOr<SCFTileAndFuseResult>
tileConsumerAndFuseProducersUsingSCF(RewriterBase &rewriter,
                                     TilingInterface consumer,
                                     const SCFTileAndFuseOptions &options);
 
/// Fuse the consumer `candidateSlices` by computing the required slice of the
/// consumer in-place. All the entries of `candidateSlices` are expected to map
/// to the same consumer. The method returns an error if the consumer cannot be
/// tiled in a manner that is consistent for all the passed slices. Note that
/// the method replaces the uses of `candidateSlices` with the tiled and fused
/// consumer value but does not delete the slice operations.
/// TODO(MaheshRavishankar): A more natural way of exposing the consumer fusion
/// is to take the consumer operation, and find the slices to use for fusion
/// by walking its operands to the `loops` and then into the body to get the
/// slices used for fusion.
struct SCFFuseConsumerOfSliceResult {
  // Original untiled consumer operands.
  SmallVector<OpOperand *> origConsumerOperands;
  // Tiled and fused consumer operands.
  SmallVector<OpOperand *> tiledAndFusedConsumerOperands;
  SmallVector<Operation *> tiledOps;
};
/// When the consumer is a `linalg.pack`/`linalg.unpack`, `fn` (if non-null) is
/// invoked with the consumer and `candidateSlices` to obtain the inner-tile
/// alignment hint in the consumer's own iteration domain.
FailureOr<scf::SCFFuseConsumerOfSliceResult>
tileAndFuseConsumerOfSlices(RewriterBase &rewriter,
                            ArrayRef<Operation *> candidateSlices,
                            MutableArrayRef<LoopLikeOpInterface> loops,
                            const InnerTileAlignmentFnTy &fn = nullptr);
 
/// Fuse the `consumer` operation into the loop nest provided by `loops`.
/// The transformation looks for operands in the `consumer` that are defined
/// by the outermost loop of the loop nest in `loops`. The nested loop is
/// expected to have the structure of the loops generated through tiling.
///
/// When the consumer is a `linalg.pack`/`linalg.unpack`, `fn` (if non-null) is
/// invoked with the `consumer` and the fused slices to obtain the inner-tile
/// alignment hint in the consumer's own iteration domain.
FailureOr<scf::SCFFuseConsumerOfSliceResult>
tileAndFuseConsumer(RewriterBase &rewriter, Operation *consumer,
                    MutableArrayRef<LoopLikeOpInterface> loops,
                    const InnerTileAlignmentFnTy &fn = nullptr);
 
/// Method to lower an `op` that implements the `TilingInterface` to
/// loops/scalars.
FailureOr<SmallVector<scf::ForOp>>
lowerToLoopsUsingSCFForOp(RewriterBase &rewriter, TilingInterface op);
 
/// Method to tile a reduction and generate a parallel op within a serial loop.
/// Each of the partial reductions are calculated in parallel. Then after the
/// loop all the partial reduction are merged into a final reduction.
/// For example for the following sequence
///
/// ```mlir
/// %0 = linalg.generic %in ["parallel", "reduction"]
///   : tensor<7x9xf32> -> tensor<7xf32>
/// ```
///
/// into:
///
/// ```mlir
/// %0 = linalg.fill ... : tensor<7x4xf32>
/// %1 = scf.for ... iter_args(%arg0 = %0)
///   %2 = tensor.extract_slice %arg0 : tensor<7x4xf32> -> tensor<7x?xf32>
///   %3 = tensor.extract_slice %in : tensor<7x9xf32> -> tensor<7x?xf32>
///   %4 = linalg.generic %2, %3 ["parallel", "parallel"]
///     : tensor<7x?xf32> -> tensor<7x?xf32>
///   %5 = tensor.insert_slice %3, %0[0, 0] : tensor<7x4xf32>
/// }
/// %6 = linalg.generic %1 ["parallel", "reduction"]
///   : tensor<7x4xf32> -> tensor<7xf32>
/// ```
FailureOr<scf::SCFTilingResult>
tileReductionUsingScf(RewriterBase &b, PartialReductionOpInterface op,
                      ArrayRef<OpFoldResult> tileSizes);
 
} // namespace scf
} // namespace mlir
 
#endif // MLIR_DIALECT_SCF_TRANSFORMS_TILEUSINGINTERFACE_H