// Copyright 2023 The IREE Authors
//
// Licensed 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 IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_
#define IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Support/LLVM.h"

namespace mlir::iree_compiler::IREE::Codegen {

//===----------------------------------------------------------------------===//
// Relational operator and IOstream implementations for Layout Structs.
//===----------------------------------------------------------------------===//

bool operator==(TileSwizzle::Dim lhs, TileSwizzle::Dim rhs);
bool operator!=(TileSwizzle::Dim lhs, TileSwizzle::Dim rhs);
bool operator==(const TileSwizzle &lhs, const TileSwizzle &rhs);
bool operator!=(const TileSwizzle &lhs, const TileSwizzle &rhs);

llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                              TileSwizzle::Dim::Kind kind);

llvm::raw_ostream &operator<<(llvm::raw_ostream &os, TileSwizzle::Dim dim);

llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                              const TileSwizzle &swizzle);

bool operator==(const MaterializeEncodingInfo &lhs,
                const MaterializeEncodingInfo &rhs);
bool operator!=(const MaterializeEncodingInfo &lhs,
                const MaterializeEncodingInfo &rhs);
llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                              const MaterializeEncodingInfo &encodingInfo);

//===----------------------------------------------------------------------===//
// Layout Utilities.
//===----------------------------------------------------------------------===//

/// Conversion between TileSwizzle::Dim::Kind and string.
std::string convertSwizzleKindToString(TileSwizzle::Dim::Kind kind);
std::optional<TileSwizzle::Dim::Kind> convertStringToSwizzleKind(StringRef str);

/// Conversion between TileSwizzle::Dim::SymbolicMultiplier and string.
std::string convertSymbolicMultiplierToString(
    TileSwizzle::Dim::SymbolicMultiplier symbolicMultiplier);
std::optional<TileSwizzle::Dim::SymbolicMultiplier>
convertStringToSymbolicMultiplier(StringRef str);

/// Conversion between TileSwizzle struct and DictionaryAttr.
DictionaryAttr serializeTileSwizzle(MLIRContext *ctx,
                                    const TileSwizzle &swizzle);
std::optional<TileSwizzle> deserializeTileSwizzle(DictionaryAttr attr);

/// Conversion between MaterializeEncodingInfo struct and DictionaryAttr.
DictionaryAttr serializeEncodingInfo(MLIRContext *ctx,
                                     const MaterializeEncodingInfo &info);
std::optional<MaterializeEncodingInfo>
deserializeEncodingInfo(DictionaryAttr attr);

/// Returns true if the `info` denotes an identity layout, i.e., there is no
/// relayout requirement.
bool isIdentityLayout(const MaterializeEncodingInfo &info);

/// Concatenates the vectors.
SmallVector<int64_t>
getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape);

/// The dimensions needed for materializing an encoding for either a contraction
/// or a scaled contraction.
struct EncodingContractionLikeDimInfo {
  /// Info for each dimension. `shouldHaveDim` indicates whether or not the dim
  /// is present on the root operation's operand, and `operandIdx` is the index
  /// of the encoded tensor that belongs to the Dim. For example, the LHS
  /// operand's encoding would have `mDim = {true, 0}` for a plain matmul, while
  /// the RHS operand's encoding would have `mDim = {false, std::nullopt}`.
  /// If `shouldHaveDim` is true, and `operandIdx` is std::nullopt, this
  /// indicates that the encoded Value may be the source of a broadcast, and the
  /// Dim is not present on the Value, but it is present on the corresponding
  /// operand of the encoding root operation.
  struct DimInfo {
    bool shouldHaveDim = false;
    std::optional<unsigned> operandIdx;
  };
  DimInfo batchDim;
  DimInfo mDim;
  DimInfo nDim;
  DimInfo kDim;
  DimInfo kBDim;
};

/// Returns the dimension info needed for materializing an encoding for either a
/// contraction or a scaled contraction. This function expects at most one each
/// of M, N, K, Kb, and Batch dimensions, since multiple contraction dimensions
/// are currently not supported in data tiling.
FailureOr<EncodingContractionLikeDimInfo>
getEncodingContractionLikeDims(Encoding::EncodingAttr encoding);

struct TileMxNxKxKb {
  int64_t M = 1;
  int64_t N = 1;
  int64_t K = 1;
  int64_t KB = 1;
};

FailureOr<MaterializeEncodingInfo>
getEncodingInfoForMatmul(Encoding::EncodingAttr encoding,
                         TileMxNxKxKb tileMxNxKxKb);

struct TileMxNxK {
  int64_t M = 1;
  int64_t N = 1;
  int64_t K = 1;
};

FailureOr<MaterializeEncodingInfo>
getEncodingInfoForMatmul(Encoding::EncodingAttr encoding, TileMxNxK tileMxNxK);

/// Tile sizes for convolution encoding.
struct TileOCxIC {
  int64_t OC = 1; // k0: output channel tile
  int64_t IC = 1; // c0: input channel tile
};

/// Returns MaterializeEncodingInfo for a convolution operand.
FailureOr<MaterializeEncodingInfo>
getEncodingInfoForConv(Encoding::EncodingAttr encoding, TileOCxIC tile);

//===---------------------------------------------------------------------===//
// Data-tiled convolution utilities.
//===---------------------------------------------------------------------===//

/// Iteration space of the data-tiled convolution generic. This is the
/// canonical definition: the materialization builder and any code that
/// inspects the resulting op both derive from it.
///
///   Loop dims: (n, oc_outer, oh, ow, ic_outer, fh, fw, oc_inner, ic_inner)
///               d0  d1        d2  d3  d4        d5  d6  d7        d8
///   input:  -> (d0, d4, d2*sH + d5*dH, d3*sW + d6*dW, d8)
///   filter: -> (d1, d4, d5, d6, d8, d7)
///   output: -> (d0, d1, d2, d3, d7)
///
/// corresponding to packed operand layouts:
///   input:  [N, IC/c0, H, W, c0]
///   filter: [OC/k0, IC/c0, FH, FW, c0, k0]
///   output: [N, OC/k0, OH, OW, k0]
struct DataTiledConvIterationSpace {
  SmallVector<AffineMap> indexingMaps;
  SmallVector<utils::IteratorType> iteratorTypes;
};

/// Returns the indexing maps and iterator types for a data-tiled 2D
/// convolution with the given strides and dilations.
DataTiledConvIterationSpace
getDataTiledConvIterationSpace(MLIRContext *ctx, ArrayRef<int64_t> strides,
                               ArrayRef<int64_t> dilations);

/// Returns true if `op` is a 9D data-tiled convolution generic emitted by
/// encoding materialization (matched by iterator types, operand ranks,
/// indexing maps, and multiply-accumulate body).
bool isDataTiledConvGeneric(Operation *op);

} // namespace mlir::iree_compiler::IREE::Codegen

#endif // IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_
