// Copyright 2024 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_DIALECT_ENCODING_IR_ENCODINGTYPES_H_
#define IREE_COMPILER_DIALECT_ENCODING_IR_ENCODINGTYPES_H_

#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/TensorEncoding.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LogicalResult.h"

// clang-format off: must be included after all LLVM/MLIR headers.
#include "iree/compiler/Dialect/Encoding/IR/EncodingEnums.h.inc" // IWYU pragma: export
// clang-format on

namespace mlir::iree_compiler::IREE::Encoding {

static constexpr char kEncodingResolverAttrName[] = "iree.encoding.resolver";

/// Assign a name to operand indices for clarity
const int64_t MATMUL_LHS = 0;
const int64_t MATMUL_RHS = 1;
const int64_t MATMUL_RESULT = 2;
/// Scaled matmul
const int64_t SCALED_MATMUL_LHS = 0;
const int64_t SCALED_MATMUL_RHS = 1;
const int64_t SCALED_MATMUL_LHS_SCALES = 2;
const int64_t SCALED_MATMUL_RHS_SCALES = 3;
const int64_t SCALED_MATMUL_RESULT = 4;
/// Convolutions
const int64_t CONV_IN = 0;
const int64_t CONV_FILTER = 1;
const int64_t CONV_OUT = 2;

/// Convert operand index to strings for printing
std::string stringifyOperandIndex(EncodingOpType opType,
                                  IntegerAttr operandIndex);

/// Designates a dimension in a matmul (either the M or the N dimension) as
/// being "narrow", i.e. small enough that we bother lowering the amount of
/// padding along that dimension compared to how padding we apply to
/// sufficiently large dimensions.
struct MatmulNarrowDim {
  // Enumerates dimensions of a matmul that may be labelled as narrow.
  enum class Dim {
    None,
    M,
    N,
  };
  Dim dim = Dim::None; // Which dimension is designated by *this.
  int64_t size = 0;    // Size of the designated dimension, or kDynamic.

  explicit operator bool() const { return dim != Dim::None; }
  bool isM() const { return dim == Dim::M; }
  bool isN() const { return dim == Dim::N; }
};

/// Returns the narrow dim in a given `linalgOp`, with respect to the given
/// `narrowThreshold` below which a dimension is eligible to be considered
/// narrow. If both M and N are narrow, M is returned. If neither M nor N are
/// narrow, this returns a default-constructed falsish value.
MatmulNarrowDim getMatmulNarrowDim(linalg::LinalgOp linalgOp,
                                   int narrowThreshold);

// The structs defined here because they are used by encoding_interfaces.td.

/// Bundles an encoding attribute with its associated dynamic information.
/// This is returned by interface methods that query encoding properties
/// from operations. Dynamic values can include runtime information needed
/// by the encoding (e.g., Batch, M, N, K dimensions for matmul encodings).
struct EncodingProperties {
  /// The encoding attribute for the operand/result.
  Attribute encoding;
  /// Optional dynamic values needed by the encoding. This field may be empty
  /// for operations that don't require dynamic dimension tracking.
  ///
  /// For operations implementing IndexingMapOpInterface (e.g., linalg ops),
  /// these correspond to the dynamic entries in iteration_sizes, ordered by
  /// loop dimension index. For example:
  /// - matmul with iteration_sizes=[?, ?, ?] -> [M, N, K]
  /// - batch_matmul with iteration_sizes=[?, ?, ?, ?] -> [Batch, M, N, K]
  /// - matmul with iteration_sizes=[?, 128, ?] -> [M, K] (N is static)
  ///
  /// For operations that do not implement IndexingMapOpInterface, the meaning
  /// and ordering of these values is defined by downstream users.
  SmallVector<Value> dynamicValues;
};

/// Bundles encoding properties for all operands and DPS init operands of an
/// operation. This is returned by SerializableAttr::getEncodingProperties()
/// when deriving encodings from an operation (e.g., linalg.matmul).
struct OpEncodingProperties {
  SmallVector<EncodingProperties> operands;
  SmallVector<EncodingProperties> inits;
};

struct PropagationEncoding {
  SmallVector<Attribute> operandEncodings;
  SmallVector<Attribute> resultEncodings;
  /// Optional dynamic encoding dimension values that should be carried through
  /// during propagation. See EncodingProperties::dynamicValues for details on
  /// the ordering and semantics of these values.
  SmallVector<Value> encodingDims;
};

struct PropagationResult {
  // A list of operations that are created by the propagation. They are returned
  // to the caller for further transformation.
  SmallVector<Operation *> encodedOps;

  // A list of new set_encoding/unset_encoding ops that are generated by
  // the propagation.
  SmallVector<Operation *> generatedEncodingOps;

  // The new results created after propagating an encoding through an operation.
  // It is returned to the caller for further transformation or replacement.
  SmallVector<Value> replacements;
};

/// Information returned by convertForBitcast interface method.
/// Contains the new shape and encoding after bitcasting to a new element type.
struct BitcastEncodingInfo {
  /// The new shape after bitcasting.
  SmallVector<int64_t> newShape;
  /// The new encoding attribute after bitcasting.
  Attribute encoding;
};

} // namespace mlir::iree_compiler::IREE::Encoding

// clang-format off
#include "iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.h.inc" // IWYU pragma: export
#define GET_ATTRDEF_CLASSES
#include "iree/compiler/Dialect/Encoding/IR/EncodingAttrs.h.inc" // IWYU pragma: export
#undef GET_ATTRDEF_CLASSES
#define GET_TYPEDEF_CLASSES
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h.inc" // IWYU pragma: export
#undef GET_TYPEDEF_CLASSES
// The EncodingTypeInterfaces.h.inc needs to be included after
// EncodingTypes.h.inc because an interface method could have EncodingAttr
// types.
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypeInterfaces.h.inc" // IWYU pragma: export
#include "iree/compiler/Dialect/Encoding/IR/EncodingOpInterfaces.h.inc" // IWYU pragma: export
// clang-format on

#endif // IREE_COMPILER_DIALECT_ENCODING_IR_ENCODINGTYPES_H_
