// Copyright 2025 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_UTILS_UTILS_H_
#define IREE_COMPILER_DIALECT_ENCODING_UTILS_UTILS_H_

#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"

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

constexpr char kDataTilingHint[] = "iree.opt.data_tiling";

/// Returns true if the operation has data-tiling hint attribute.
inline bool hasDataTilingHint(Operation *op) {
  return op->getAttr(kDataTilingHint) ? true : false;
}

/// Adds an unit attribute with `kDataTilingHint` key to the operation.
inline void setDataTilingHint(Operation *op) {
  op->setAttr(kDataTilingHint, UnitAttr::get(op->getContext()));
}

/// Removes the attribute with `kDataTilingHint` key from the operation, if it
/// exists.
inline void removeDataTilingHint(Operation *op) {
  (void)op->removeAttr(kDataTilingHint);
}

/// Returns the encoding attribute from the type if there is an encoding that
/// implements SerializableAttr. Otherwise, returns null.
SerializableAttr getSerializableAttr(RankedTensorType type);

/// Returns the encoding attribute from the type if there is an encoding.
/// Otherwise, returns null.
EncodingAttr getEncodingAttr(RankedTensorType type);

/// Returns true if the type contains packed_storage attribute.
bool hasPackedStorageAttr(RankedTensorType type);

/// Returns the ContractionDimensions for the encoding user_indexing_maps.
FailureOr<linalg::ContractionDimensions>
getEncodingContractionDims(EncodingAttr encoding);

/// Returns the ScaledContractionDimensions for the encoding user_indexing_maps.
FailureOr<IREE::LinalgExt::ScaledContractionDimensions>
getEncodingScaledContractionDims(EncodingAttr encoding);

/// The sizes for contraction-like ops with one each of Batch, M, N, K, Kb dims.
struct BxMxNxKxKb {
  int64_t batch = 1;
  int64_t M = 1;
  int64_t N = 1;
  int64_t K = 1;
  int64_t Kb = 1;
};

/// Returns the contraction-like sizes (batch, M, N, K, Kb) for a given
/// encoding. Supports both regular contractions and scaled contractions. For
/// regular contractions, Kb is set to 1.
FailureOr<BxMxNxKxKb> getEncodingContractionLikeSizes(EncodingAttr encoding);

/// Returns the narrow dim in a given `encoding`, ceiled to a power of two. This
/// works by inspecting the `iteration_sizes` array attribute in the `encoding`.
/// If the `iteration_sizes` of one dimension (M or N) is smaller than the
/// other, then that's the narrow dimension, because the only way it would have
/// been set to be smaller in the first place, is if we previously flagged that
/// dimension as narrow. If the `iteration_sizes` of the M and N dimensions
/// agree, then neither is a narrow dimension and this returns a
/// default-constructed falsish value.
MatmulNarrowDim getPo2MatmulNarrowDim(EncodingAttr encoding);

/// Returns true if `encoding` represents a narrow-N matmul RESULT, e.g. the
/// result of a matvec.
bool isNarrowNResult(EncodingAttr encoding);

/// Rematerialize encoding dims at the given insertion point. This is useful
/// when propagating encodings through operations where the original encoding
/// dims don't dominate the new insertion point.
///
/// Returns the rematerialized dims if successful, or failure if some dims
/// cannot be rematerialized. On failure, no operations are created.
///
/// Recursively rematerializes operands as needed. Supports:
/// - Values that already dominate the insertion point (used as-is).
/// - tensor.dim ops on propagationSource (recreated using newSource).
/// - Pure operations (cloned with rematerialized operands).
FailureOr<SmallVector<Value>>
rematerializeEncodingDims(RewriterBase &builder, Operation *insertionPoint,
                          ValueRange encodingDims, Value propagationSource,
                          Value newSource);

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

#endif // IREE_COMPILER_DIALECT_ENCODING_IR_ENCODINGTYPES_H_
