//===- LinalgReleayoutOps.td - Linalg relayout ops ---------*- tablegen -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file defines Pack + Unpack Ops that have been moved from the Tensor
// dialect.
//
//===----------------------------------------------------------------------===//

#ifndef LINALG_RELEAYOUT_OPS
#define LINALG_RELEAYOUT_OPS

include "mlir/Dialect/Linalg/IR/LinalgBase.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
include "mlir/Dialect/Linalg/IR/RelayoutOpInterface.td"
include "mlir/IR/OpAsmInterface.td"

//===----------------------------------------------------------------------===//
// RelayoutOp
//===----------------------------------------------------------------------===//

class Linalg_RelayoutOp<string mnemonic, list<Trait> traits = []>
    : Op<Linalg_Dialect, mnemonic,
         !listconcat(
             traits, [DeclareOpInterfaceMethods<
                          OpAsmOpInterface, ["getAsmResultNames"]>,
                      DestinationStyleOpInterface, LinalgRelayoutOpInterface,
                      ConditionallySpeculatable,
                      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
                      DeclareOpInterfaceMethods<
                          ReifyRankedShapedTypeOpInterface, [
          "reifyResultShapes"]>,
                      OptionalTypesMatchWith<"result type matches type of dest",
                                             "dest", "result", "$_self">])> {

  code commonExtraClassDeclaration = [{
    size_t getSourceRank() { return getSourceType().getRank(); };
    size_t getDestRank() { return getDestType().getRank(); };
    ShapedType getSourceType() {
      return ::llvm::cast<ShapedType>(getSource().getType()); };
    ShapedType getDestType() {
      return ::llvm::cast<ShapedType>(getDest().getType()); };

    MutableOperandRange getDpsInitsMutable() { return getDestMutable(); }

    /// Interface method for ConditionallySpeculatable.
    Speculation::Speculatability getSpeculatability();

    /// Return a mapping from positions `inner_dims_pos` to their
    /// tile factors.
    DenseMap<int64_t, OpFoldResult> getDimAndTileMapping();

    /// Return the tile sizes as OpFoldResult(s). Note, for Ops that simply
    /// define constants (e.g. `arith.constant`), this method returns
    /// Value/result of the Op (as opposed to the corresponding constant
    /// Attribute). E.g., for:
    ///   %size = arith.constant 1 : i32
    /// it will return %size, not 1. This is intended - replacing an SSA value
    /// with a constant attribute would require updating the corresponding dim
    /// from dynamic to static.
    SmallVector<OpFoldResult> getMixedTiles();

    /// Return the tile sizes as `int64_t`. If a tile size is dynamic
    /// a sentinel `kDynamic` is introduced at that position in
    /// the returned vector.
    SmallVector<int64_t> getStaticTiles();

    /// Retrieve all outer dims for this Pack/UnPack Op, i.e. all the leading
    /// dims excluding the trailing dims corresponding to `innerTiles`. Note
    /// that this will include both tiled and non-tiled dimensions. The order
    /// of the output dimensions is consistent with the shape of the packed
    /// tensor.
    ArrayRef<int64_t> getAllOuterDims();

    /// Similar to `getAllOuterDims`, but only retrieve the outer dims that
    /// have been tiled. Also, the order of the output dimensions is consistent
    /// with `inner_dims_pos` rather than the packed tensor.
    SmallVector<int64_t> getTiledOuterDims();
  }];

  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// PackOp
//===----------------------------------------------------------------------===//

def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
    AttrSizedOperandSegments]> {
  let summary = "linalg.pack operation";
  let description = [{
    The "pack" operation converts a source tensor of rank `n` into a result
    tensor of rank `n + k` with a tiled and packed layout (maybe with padding)
    and optionally transposes the tiled source tensor dimensions.

    `inner_tiles` (mandatory) specifies `k` tile sizes. These tile sizes
    correspond to the least significant ("inner") result tensor dimension sizes,
    in the same order. Tile sizes can be static or dynamic.

    `inner_dims_pos` (mandatory) specifies `k` source tensor dimensions that are
    being tiled, where `0 <= k <= n`.
     - `inner_dims_pos[i]` specifies the source tensor dimension tiled by
    `inner_tiles[i]` where `0 <= i < k`. All the values in `inner_dims_pos` are
    within [0, n).
     - The tiled dimensions (of size `inner_tiles`) are added to the end of the
     result tensor in the order in which they appear, i.e.
     `shape(result)[rank(source) + i] = inner_tiles[i]` for `0 <= i < k`.
     - The following relationship for the tiled dimensions holds:
     `shape(result)[inner_dims_pos[i]] = shape(source)[inner_dims_pos[i]] / inner_tiles[i]`,
     where (⌈/⌉ indicates CeilDiv).


    Example: If `inner_tiles = [16, 32]`, the result tensor has a shape of
    `...x16x32`. If `inner_dims_pos = [0, 1]`, the 0th source dimension is tiled
    by 16 and the 1st source dimension is tiled by 32. Other source dimensions
    (if any) are not tiled. If `inner_dims_pos = [1, 0]`, the 1st dimension is
    tiled by 16 and the 0th dimension is tiled by 32.

    Example:
    ```mlir
    // NC to NCnc
    %0 = linalg.pack %source inner_dims_pos = [0, 1] inner_tiles = [8, 32]
        into %dest : tensor<128x256xf32> -> tensor<16x8 x 8x32 xf32>
    //                                             \  /   \  /
    //                                 Outer Dims: 16x8   Inner Dims: 8x32

    // CHW to CHWhw
    %0 = linalg.pack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2]
        into %dest : tensor<3x20x24xf32> -> tensor<3x10x6 x 4x2 xf32>
    //                                              \  /    \ /
    //                                 Outer Dims: 3x10x6  Inner Dims: 4x2

    // HCW to HCWhw
    %0 = linalg.pack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2]
        into %dest : tensor<18x3x32xf32> -> tensor<9x3x8 x 4x2 xf32>
    //                                              \  /   \ /
    //                                 Outer Dims: 9x3x8  Inner Dims: 4x2
    ```

    `outer_dims_perm` (optional) specifies a permutation for the outer
    dimensions. If specified, it must have `n` elements.

    Example:
    ```mlir
    // CK to KCck
    %0 = linalg.pack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
        inner_tiles = [8, 32] into %dest
        : tensor<128x256xf32> -> tensor<8x16 x 8x32 xf32>
    //                                  \  /
    //            compare with "NC to NCnc": outer dims are transposed
    ```

    `padding_value` specifies a padding value at the boundary on non-perfectly
    divisible dimensions. Padding is optional:
    - If absent, it is assumed that for all inner tiles,
      `shape(source)[inner_dims_pos[i]] % inner_tiles[i] == 0`, i.e. all inner
      tiles divide perfectly the corresponding outer dimension in the result
      tensor. It is UB if the tile does not perfectly divide the dimension.
    - If present, it will pad along high dimensions (high-padding) to make the
      tile complete. Note that it is not allowed to have artificial padding that
      is not strictly required by linalg.pack (i.e., padding past what is needed
      to complete the last tile along each packed dimension). It is UB if extra
      padding is requested.
    It is not possible to verify the requirements statically with dynamic
    shapes, so they are treated as UB.

    Example:
    ```mlir
    %0 = linalg.pack %arg0 padding_value(%pad : f32) outer_dims_perm = [2, 1, 0]
        inner_dims_pos = [1] inner_tiles = [2] into %arg1
        : tensor<200x127x256xf32> -> tensor<256x64x200x2xf32>
    //                 \
    //                padded and tiled dim
    //
    // Source dimension 1 is tiled. 64 does not divide 127 evenly, so 1 padded
    // element is added at the end.
    //
    // Note: Only tiled dimensions can be padded.
    ```

    Invalid example that has artificial padding:
    ```mlir
    %0 = linalg.pack %src padding_value(%cst : f32) inner_dims_pos = [0]
        inner_tiles = [8] into %dest
        : tensor<9xf32> -> tensor<3x8xf32>
    //                             \
    //            expect tensor<2x8xf32> because CeilDiv(9, 8) = 2
    ```
  }];
  let arguments = (ins TensorOrMemRef<[AnyType]>:$source,
      TensorOrMemRef<[AnyType]>:$dest, 
      Optional<AnyType>:$padding_value,
      DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$outer_dims_perm,
      DenseI64ArrayAttr:$inner_dims_pos, 
      Variadic<Index>:$inner_tiles,
      DenseI64ArrayAttr:$static_inner_tiles);
  let results = (outs Optional<AnyRankedTensor>:$result);

  let builders = [
    OpBuilder<(ins "Value":$source, "Value":$dest,
      "ArrayRef<int64_t>":$innerDimsPos,
      "ArrayRef<OpFoldResult>":$innerTiles,
      CArg<"std::optional<Value>", "std::nullopt">:$paddingValue,
      CArg<"ArrayRef<int64_t>", "{}">:$outerDimsPerm)>
  ];

  let extraClassDeclaration = commonExtraClassDeclaration # [{
    // Method to get the shape of the result as `SmallVector<OpFoldResult>`.
    // This is a static method to allow getting the shape of the destination
    // expected while creating a `pack` op.
    static SmallVector<OpFoldResult> getResultShape(OpBuilder &builder,
        Location loc, ArrayRef<OpFoldResult> sourceDims,
        ArrayRef<OpFoldResult> innerTileDims, ArrayRef<int64_t> innerDimsPos,
        ArrayRef<int64_t> outerDimsPerm = {});

    // Method to get the `RankedTensorType` of the result based on the inner
    // tiles, position of the inner tiles (innerDimsPos)  and interchange vector
    // of outer loops (outerDimsPerm).
    static RankedTensorType inferPackedTensorType(RankedTensorType sourceType,
        ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
        ArrayRef<int64_t> outerDimsPerm = {});

    // Method to get the `MemRefType` of the result based on the inner
    // tiles, position of the inner tiles (innerDimsPos)  and interchange vector
    // of outer loops (outerDimsPerm).
    static MemRefType inferPackedMemRefType(MemRefType sourceType,
        ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
        ArrayRef<int64_t> outerDimsPerm = {});

    // Returns the shape of the packed type. It is a shared helper that helps
    // type inference methods in a way that ensures that they agree on which 
    // dimensions are dynamic.
    static SmallVector<int64_t> inferPackedShape(ArrayRef<int64_t> inputShape,
        ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
        ArrayRef<int64_t> outerDimsPerm = {});

    // Returns true if we have enough static information to catch undefined
    // behavior when the tile size does not divide perfectly the dimension of
    // the input tensor. Detecting UB requires that the input size and either
    // corresponding tile or output size are static.
    static bool requirePaddingValue(ArrayRef<int64_t> inputShape,
                                    ArrayRef<int64_t> innerDimsPos,
                                    ArrayRef<int64_t> outputShape,
                                    ArrayRef<int64_t> outerDimsPerm,
                                    ArrayRef<OpFoldResult> innerTiles);

    // Same as above function but here dynamic dimensions are assumed
    // to require padding.
    static bool requirePaddingValueStrict(ArrayRef<int64_t> inputShape,
                                          ArrayRef<int64_t> innerDimsPos,
                                          ArrayRef<int64_t> outputShape,
                                          ArrayRef<int64_t> outerDimsPerm,
                                          ArrayRef<OpFoldResult> innerTiles);

    static Value createDestinationTensor(OpBuilder &b, Location loc,
        Value source, ArrayRef<OpFoldResult> innerTileSizes,
        ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm);

    /// Build and return a new PackOp that is a clone of the current PackOp with
    /// (innerDimsPos, innerTiles) (resp. outerDimsPerm) are permuted by
    /// innerPermutation (resp. outerPermutation).
    /// A new `tensor.empty` of the proper shape is built in the process.
    /// Asserts that:
    ///   - At least one of innerPermutation or outerPermutation is non-empty.
    ///   - If not empty, innerPermutation is a valid permutation of size
    ///     matching innerDimPos.
    ///   - If not empty, outerPermutation is a valid permutation of size
    ///     matching outerDimsPerm.
    PackOp createTransposedClone(OpBuilder &b,
                                 Location loc,
                                 ArrayRef<int64_t> innerPermutation,
                                 ArrayRef<int64_t> outerPermutation);

    /// Check if this PackOp is like a simple pad operation.
    /// In other words, this operation:
    /// 1. adds useless dimensions (dimension of size 1),
    /// 2. pads the other ones, and
    /// 3. doesn't shuffle the dimensions
    bool isLikePad();
  }];

  let hasCanonicalizeMethod = 1;

  let hasFolder = 1;

  let hasCustomAssemblyFormat = 1;
}

//===----------------------------------------------------------------------===//
// UnPackOp
//===----------------------------------------------------------------------===//

def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
  let summary = "linalg.unpack operation";
  let description = [{
    The "unpack" operation converts a source tensor of rank `n` with a tiled and
    packed layout to a result tensor of rank `n - k`.

    `inner_tiles` (mandatory) specifies `k` tile sizes. These tile sizes
    correspond to the least significant ("inner") source tensor dimension sizes.
    The behavior of this op is undefined if:
    - `inner_tiles` do not exactly match with the corresponding source tensor
      dimension sizes.
    - Or, `inner_tiles[i]` does not divide the size of dimension
      `inner_dims_pos[i]` (assuming that `outer_dims_perm` is not specified)
      evenly.

    `inner_dims_pos` (mandatory) specifies `k` result tensor (i.e. unpacked
    tensor) dimensions that were tiled with the `inner_tiles` to create the
    packed source tensor. The source tensor (i.e. packed tensor) dimensions can
    be unpacked given `inner_dims_pos` as follows.
    - For `0 <= i < k` the following relationship holds:
    `shape(result)[inner_dims_pos[i]] <= shape(source)[n-k+i] * shape(source)[inner_dims_pos[i]]`.
    - For `0 <= j < n-k` and `j` not in `inner_dims_pos` the following relationship holds:
    `shape(result)[j] = shape(source)[j]`.

    `outer_dims_perm` (optional) specifies a permutation for the outer
    dimensions. If specified, it must have `n - k` elements. If specified, this
    permutation is applied before combining any dimensions.

    Note, the unpack operation may drop any padding introduced by the pack
    operation and hence the following holds
    `NumElementsOf(source) >= NumElementsOf(result)`.

    Examples:

    ```mlir
    // NCnc to NC:
    %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [8, 32]
        into %dest : tensor<16x8 x 8x32 xf32> -> tensor<128x256xf32>
    //                      \  /   \  /
    //          Outer Dims: 16x8  Inner Dims: 8x32

    // CK to KCck:
    %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
        inner_tiles = [8, 32]
        into %dest : tensor<8x16 x 8x32 xf32> -> tensor<128x256xf32>
    //                      \  /   \  /
    //          Outer Dims: 8x16  Inner Dims: 8x32

    // CHW to CHWhw:
    %0 = linalg.unpack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2]
        into %dest : tensor<3x10x6 x 4x2 xf32> -> tensor<3x20x24xf32>
    //                       \  /    \ /
    //          Outer Dims: 3x10x6  Inner Dims: 4x2

    // HCW to HCWhw
    %0 = linalg.unpack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2]
        into %dest : tensor<9x3x8 x 4x2 xf32> -> tensor<18x3x32xf32>
    //                       \  /   \ /
    //          Outer Dims: 9x3x8   Inner Dims: 4x2
    ```
  }];
  let arguments = (ins TensorOrMemRef<[AnyType]>:$source,
      TensorOrMemRef<[AnyType]>:$dest,
      DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$outer_dims_perm,
      DenseI64ArrayAttr:$inner_dims_pos, Variadic<Index>:$inner_tiles,
      DenseI64ArrayAttr:$static_inner_tiles);
  let results = (outs Optional<AnyRankedTensor>:$result);

  let builders = [
    OpBuilder<(ins "Value":$source, "Value":$dest,
    "ArrayRef<int64_t>":$innerDimsPos,
    "ArrayRef<OpFoldResult>":$innerTiles,
    CArg<"ArrayRef<int64_t>", "{}">:$outerDimsPerm)>
  ];

  let extraClassDeclaration = commonExtraClassDeclaration # [{
    static Value createDestinationTensor(OpBuilder &b, Location loc,
        Value source, ArrayRef<OpFoldResult> innerTileSizes,
        ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm);

    /// Build and return a new UnPackOp that is a clone of the current UnPackOp
    /// with (innerDimsPos, innerTiles) (resp. outerDimsPerm) are permuted by
    /// innerPermutation (resp. outerPermutation).
    /// Asserts that:
    ///   - At least one of innerPermutation or outerPermutation is non-empty.
    ///   - If not empty, innerPermutation is a valid permutation of size
    ///     matching innerDimPos.
    ///   - If not empty, outerPermutation is a valid permutation of size
    ///     matching outerDimsPerm.
    UnPackOp createTransposedClone(OpBuilder &b,
                                   Location loc,
                                   Value transposedSource,
                                   ArrayRef<int64_t> innerPermutation,
                                   ArrayRef<int64_t> outerPermutation);

    /// Returns true if it is statically known that the `sliceOp` result shape
    /// is compatible with the `unPackOp`. I.e., it does not drop any tile.
    bool canFoldSliceOp(tensor::ExtractSliceOp sliceOp);

    /// Check if this UnPackOp is like a simple unpad operation.
    /// In other words, this operation:
    /// 1. drops useless dimensions (dimension of size 1), and
    /// 2. reduces dimensions in place (i.e., no transpose.)
    bool isLikeUnPad();
  }];

  let hasCanonicalizeMethod = 1;

  let hasFolder = 1;

  let hasCustomAssemblyFormat = 1;
}

#endif // LINALG_RELEAYOUT_OPS