// 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

#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"

#include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
#include "iree/compiler/Utils/EncodingUtils.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LLVM.h"

#include <cassert>

#define DEBUG_TYPE "iree-encoding-attrs"

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

/// Returns the bit-width of the scalar type. If the type is complex, it
/// returns the type of individual elements * 2 (1 for real and 1 for
/// complex).
static unsigned getTypeBitWidth(Type type) {
  if (auto complexType = dyn_cast<ComplexType>(type)) {
    return 2 * complexType.getElementType().getIntOrFloatBitWidth();
  }
  return type.getIntOrFloatBitWidth();
}

/// Returns the number of bytes an element of the given type occupies in
/// memory. This is in the default dense conversion to machine words where
/// sizes must be powers of two aligned to bytes.
///
/// Examples:
///   getRoundedElementByteWidth(i1) = 1
///   getRoundedElementByteWidth(i23) = 4
///   getRoundedElementByteWidth(i32) = 4
///   getRoundedElementByteWidth(bf16) = 2
///   getRoundedElementByteWidth(i33) = 8
///   getRoundedElementByteWidth(complex<f32>) = 8
static int32_t getRoundedElementByteWidth(Type type) {
  unsigned bitsUnaligned = getTypeBitWidth(type);
  assert(bitsUnaligned > 0 && "0-width types unsupported");
  // Round up to 8-bit aligned bytes.
  unsigned byteAligned = (bitsUnaligned + 8 - 1) / 8;
  // Round up to the next power of two (unless already a power of two).
  return llvm::PowerOf2Ceil(byteAligned);
}

//===---------------------------------------------------------------------===//
// iree_encoding.layout
//===---------------------------------------------------------------------===//

LayoutAttr
LayoutAttr::getChecked(llvm::function_ref<InFlightDiagnostic()> emitError,
                       MLIRContext *context, ArrayAttr layoutsAttr) {
  if (failed(LayoutAttr::verify(emitError, layoutsAttr))) {
    return LayoutAttr();
  }
  return LayoutAttr::get(context, layoutsAttr);
}

LayoutAttr LayoutAttr::get(MLIRContext *context, ArrayAttr layoutsAttr) {
  auto emitError = mlir::detail::getDefaultDiagnosticEmitFn(context);
  if (failed(LayoutAttr::verify(emitError, layoutsAttr))) {
    return LayoutAttr();
  }
  return Base::get(context, layoutsAttr);
}

LogicalResult
LayoutAttr::verify(function_ref<mlir::InFlightDiagnostic()> emitError,
                   ArrayAttr layoutsAttr) {
  ArrayRef<Attribute> layouts = layoutsAttr.getValue();
  if (layouts.empty()) {
    return emitError() << "expected non-empty layouts";
  }
  if (!llvm::all_of(layouts, llvm::IsaPred<SerializableAttr>)) {
    return emitError() << "expected all the layout attributes to implement "
                          "SerializableAttr";
  }
  return success();
}

bool LayoutAttr::isSerialized() const {
  ArrayRef<Attribute> layouts = getLayouts().getValue();
  return llvm::all_of(layouts, [](Attribute layout) {
    return cast<SerializableAttr>(layout).isSerialized();
  });
}

bool LayoutAttr::isIdentityLayout() const {
  auto layouts = getLayouts().getAsRange<SerializableAttr>();
  return llvm::all_of(layouts,
                      [](auto attr) { return attr.isIdentityLayout(); });
}

Value LayoutAttr::calculateStorageSizeInBytes(Location loc, OpBuilder &builder,
                                              RankedTensorType type,
                                              ValueRange dynamicDims) const {
  ArrayAttr layoutsAttr = getLayouts();
  Value res;
  for (auto attr : layoutsAttr.getAsRange<SerializableAttr>()) {
    Value requestedSize =
        attr.calculateStorageSizeInBytes(loc, builder, type, dynamicDims);
    if (!res) {
      res = requestedSize;
      continue;
    }
    res = arith::MaxUIOp::create(builder, loc, res, requestedSize);
  }
  return res;
}

LogicalResult LayoutAttr::verifyEncoding(
    llvm::ArrayRef<int64_t> shape, mlir::Type elementType,
    llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
  for (Attribute attr : getLayouts()) {
    auto verifiableAttr = dyn_cast<VerifiableTensorEncoding>(attr);
    if (verifiableAttr &&
        failed(verifiableAttr.verifyEncoding(shape, elementType, emitError))) {
      return failure();
    }
  }
  return success();
}

//===---------------------------------------------------------------------===//
// iree_encoding.packed_storage
//===---------------------------------------------------------------------===//

bool PackedStorageAttr::isSerialized() const { return true; }

bool PackedStorageAttr::isCompatibleWith(
    IREE::Encoding::SerializableAttr other) const {
  // The packed_storage encodings is compatible with empty encodings, because it
  // doesn't impact tensor shape.
  if (!other) {
    return true;
  }
  // Otherwise, packed_storage is only compatible with itself.
  return isa<IREE::Encoding::PackedStorageAttr>(other);
}

Value PackedStorageAttr::calculateStorageSizeInBytes(
    Location loc, OpBuilder &builder, RankedTensorType type,
    ValueRange dynamicDims) const {
  unsigned elementBits = getTypeBitWidth(type.getElementType());
  assert(elementBits <= 8 && "packed_storage only allowed for sub-byte types");
  assert(llvm::isPowerOf2_32(elementBits) &&
         "packed_storage only allowed for power-of-two types");

  // Calculate static dimensions if there are any.
  int64_t staticCount = 1;
  for (unsigned i = 0; i < type.getRank(); ++i) {
    if (!type.isDynamicDim(i)) {
      staticCount *= type.getDimSize(i);
    }
  }

  // Emit computation of dynamic dimensions.
  auto value =
      arith::ConstantIndexOp::create(builder, loc, staticCount).getResult();
  for (Value dim : dynamicDims) {
    value = builder.createOrFold<arith::MulIOp>(
        loc, value, dim, arith::IntegerOverflowFlags::nsw);
  }

  // For sub-byte element types we need to divide by the number of elements that
  // can be packed into one byte.
  unsigned elementsPerByte = 8 / elementBits;
  auto divisor = arith::ConstantIndexOp::create(builder, loc, elementsPerByte);
  value = builder.createOrFold<arith::CeilDivUIOp>(loc, value, divisor);
  return value;
}

LogicalResult PackedStorageAttr::verifyEncoding(
    llvm::ArrayRef<int64_t> shape, mlir::Type elementType,
    llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
  unsigned elementBitWidth = getTypeBitWidth(elementType);
  if (elementBitWidth > 7) {
    return emitError() << "bit-width of the element type is " << elementBitWidth
                       << " but packed_storage is currently only supported for "
                       << "sub-byte types";
  }
  if (!llvm::isPowerOf2_32(elementBitWidth)) {
    return emitError()
           << "bit-width of the element type is " << elementBitWidth
           << " but packed_storage currently only supports types with "
           << "power-of-two bitwidth";
  }
  return success();
}

Attribute PackedStorageAttr::getCollapsedEncoding(ArrayRef<int64_t>,
                                                  ArrayRef<int64_t>) const {
  return *this;
}

//===---------------------------------------------------------------------===//
// iree_encoding.encoding
//===---------------------------------------------------------------------===//

/// Returns a composed affine map from the provided attribute. `attr` can be
/// either an `AffineMapAttr` or an `ArrayAttr` containing `AffineMapAttr`. In
/// case of an empty `attr`, an empty affine map is returned. In case of
/// unrecognized attribute types, a failure is returned.
static FailureOr<AffineMap> getComposedAffineMap(Attribute attr) {
  if (!attr) {
    return AffineMap();
  }
  if (auto mapAttr = dyn_cast<AffineMapAttr>(attr)) {
    return mapAttr.getAffineMap();
  }
  if (auto mapsAttr = dyn_cast<ArrayAttr>(attr)) {
    if (mapsAttr.empty()) {
      return AffineMap();
    }
    // All entries should have type `AffineMapAttr`.
    if (!llvm::all_of(mapsAttr, llvm::IsaPred<AffineMapAttr>)) {
      return failure();
    }
    AffineMap map =
        cast<AffineMapAttr>(mapsAttr[mapsAttr.size() - 1]).getAffineMap();
    for (ssize_t i = mapsAttr.size() - 2; i >= 0; i--) {
      map = map.compose(cast<AffineMapAttr>(mapsAttr[i]).getAffineMap());
    }
    return map;
  }
  // Return failure in case of an unrecognized attribute type.
  return failure();
}

EncodingAttr EncodingAttr::get(MLIRContext *ctx, int64_t operandIndex,
                               EncodingOpType opType, ArrayRef<Type> elemTypes,
                               Type originalElementType,
                               ArrayRef<AffineMap> maps,
                               ArrayRef<int64_t> iterationSizes) {
  Builder b(ctx);
  auto opTypeAttr = EncodingOpTypeAttr::get(ctx, opType);
  auto originalElementTypeAttr =
      originalElementType ? TypeAttr::get(originalElementType) : TypeAttr();
  auto mapsAttr = maps.empty() ? ArrayAttr() : b.getAffineMapArrayAttr(maps);
  auto iterationSizesAttr =
      iterationSizes.empty() ? ArrayAttr() : b.getI64ArrayAttr(iterationSizes);
  return get(ctx, b.getIndexAttr(operandIndex), opTypeAttr,
             b.getTypeArrayAttr(elemTypes), originalElementTypeAttr, mapsAttr,
             iterationSizesAttr);
}

LogicalResult EncodingAttr::verify(
    function_ref<mlir::InFlightDiagnostic()> emitError,
    IntegerAttr operandIndexAttr, EncodingOpTypeAttr opTypeAttr,
    ArrayAttr elementTypesAttr, TypeAttr originalElementTypeAttr,
    ArrayAttr userIndexingMapsAttr, ArrayAttr iterationSizesAttr) {
  AffineMap indexingMap;
  if (userIndexingMapsAttr) {
    unsigned operandIndex = operandIndexAttr.getValue().getZExtValue();
    if (operandIndex >= userIndexingMapsAttr.size()) {
      return emitError()
             << "`operandIndex` exceeds the size of `user_indexing_maps`";
    }
    for (auto &&[idx, attr] : llvm::enumerate(userIndexingMapsAttr)) {
      FailureOr<AffineMap> composedMap = getComposedAffineMap(attr);
      if (failed(composedMap)) {
        return emitError() << "found a non-composable attribute in "
                              "`user_indexing_maps` at index: "
                           << idx;
      }
      if (idx == operandIndex) {
        // Keep track of the indexing map for later verification use.
        indexingMap = composedMap.value();
      }
    }
  }
  if (iterationSizesAttr) {
    if (!indexingMap) {
      return emitError() << "found `iteration_sizes` without any corresponding "
                            "`user_indexing_maps`";
    }
    if (iterationSizesAttr.size() != indexingMap.getNumDims()) {
      return emitError() << "found an encoding with "
                         << iterationSizesAttr.size()
                         << " iteration sizes, but expected "
                         << indexingMap.getNumDims()
                         << " based on the user indexing maps";
    }
  }
  return success();
}

AffineMap EncodingAttr::getMapForOperandIndex() const {
  unsigned index = getOperandIndex().getValue().getZExtValue();
  ArrayAttr userIndexingMaps = getUserIndexingMaps();
  if (!userIndexingMaps) {
    return AffineMap();
  }
  FailureOr<AffineMap> map = getComposedAffineMap(userIndexingMaps[index]);
  assert(!failed(map) &&
         "Expected a composable map. The verifier should ensure that all "
         "`user_indexing_maps` are composable.");
  return map.value();
}

SmallVector<AffineMap> EncodingAttr::getRootMaps() const {
  return llvm::map_to_vector(
      getUserIndexingMaps(), [](Attribute m) -> AffineMap {
        if (auto mapAttr = dyn_cast<AffineMapAttr>(m)) {
          return cast<AffineMapAttr>(m).getAffineMap();
        }
        if (auto mapsAttr = dyn_cast<ArrayAttr>(m)) {
          if (mapsAttr.empty()) {
            return AffineMap();
          }
          return cast<AffineMapAttr>(mapsAttr[0]).getAffineMap();
        }
        return AffineMap();
      });
}

AffineMap EncodingAttr::getLastMapForOperandIndex() const {
  unsigned index = getOperandIndex().getValue().getZExtValue();
  ArrayAttr userIndexingMaps = getUserIndexingMaps();
  if (!userIndexingMaps) {
    return AffineMap();
  }
  Attribute indexingMap = userIndexingMaps[index];
  if (auto mapAttr = dyn_cast<AffineMapAttr>(indexingMap)) {
    return mapAttr.getAffineMap();
  }
  if (auto mapsAttr = dyn_cast<ArrayAttr>(indexingMap)) {
    if (mapsAttr.empty()) {
      return AffineMap();
    }
    return cast<AffineMapAttr>(mapsAttr[mapsAttr.size() - 1]).getAffineMap();
  }
  return AffineMap();
}

std::optional<unsigned>
EncodingAttr::mapDimToOperandIndex(int64_t dimPos) const {
  return getMapForOperandIndex().getResultPosition(
      getAffineDimExpr(dimPos, getContext()));
}

SmallVector<int64_t> EncodingAttr::getIterationSizesArray() const {
  ArrayAttr iterationSizes = getIterationSizes();
  if (!iterationSizes) {
    return {};
  }
  return llvm::map_to_vector(iterationSizes, [](Attribute attr) {
    return cast<IntegerAttr>(attr).getInt();
  });
}

SmallVector<Type> EncodingAttr::getElementTypesArray() const {
  return llvm::map_to_vector(getElementTypes().getValue(), [](Attribute a) {
    return cast<TypeAttr>(a).getValue();
  });
}

EncodingAttr
EncodingAttr::cloneWithNewOperandIndexingMap(AffineMap newIndexingMap) {
  if (!newIndexingMap) {
    return *this;
  }
  ArrayAttr userIndexingMaps = getUserIndexingMaps();
  SmallVector<Attribute> newMaps(userIndexingMaps.begin(),
                                 userIndexingMaps.end());
  unsigned operandIndex = getOperandIndex().getValue().getZExtValue();
  SmallVector<Attribute> maps;
  if (auto mapForIndex = dyn_cast<AffineMapAttr>(newMaps[operandIndex])) {
    maps.push_back(AffineMapAttr::get(mapForIndex.getAffineMap()));
  } else if (auto mapForIndex = dyn_cast<ArrayAttr>(newMaps[operandIndex])) {
    maps.assign(mapForIndex.begin(), mapForIndex.end());
  }
  maps.push_back(AffineMapAttr::get(newIndexingMap));
  newMaps[operandIndex] = ArrayAttr::get(getContext(), maps);
  return get(getContext(), getOperandIndex(), getOpType(), getElementTypes(),
             getOriginalElementType(), ArrayAttr::get(getContext(), newMaps),
             getIterationSizes());
}

bool EncodingAttr::isSerialized() const { return false; }

std::optional<int64_t> EncodingAttr::getNumDynamicEncodingDims() const {
  ArrayAttr iterationSizes = getIterationSizes();
  if (!iterationSizes) {
    return 0;
  }
  return llvm::count_if(iterationSizes, [](Attribute attr) {
    return ShapedType::isDynamic(cast<IntegerAttr>(attr).getInt());
  });
}

Attribute EncodingAttr::cloneWithLayouts(ArrayRef<Attribute> layouts) const {
  MLIRContext *ctx = getContext();
  return LayoutAttr::get(ctx, ArrayAttr::get(ctx, layouts));
}

/// Helper to compute new shape when bitcasting a tensor to a different element
/// type. Returns failure if the shape is empty or bits don't divide evenly.
/// On success, returns the new shape with the last dimension adjusted.
static FailureOr<SmallVector<int64_t>>
computeBitcastShape(ArrayRef<int64_t> shape, Type originalElementType,
                    Type newElementType) {
  if (shape.empty()) {
    return failure();
  }

  unsigned oldBitWidth = getTypeBitWidth(originalElementType);
  unsigned newBitWidth = getTypeBitWidth(newElementType);

  // If bit widths are the same, shape doesn't change.
  if (oldBitWidth == newBitWidth) {
    return SmallVector<int64_t>(shape);
  }

  int64_t lastDim = shape.back();
  if (ShapedType::isDynamic(lastDim)) {
    return failure();
  }

  // Compute the new last dimension based on the ratio of bit widths.
  // totalBits = lastDim * oldBitWidth must equal newLastDim * newBitWidth.
  int64_t totalBits = lastDim * oldBitWidth;
  if (totalBits % newBitWidth != 0) {
    return failure();
  }
  int64_t newLastDim = totalBits / newBitWidth;

  SmallVector<int64_t> newShape(shape.drop_back());
  newShape.push_back(newLastDim);
  return newShape;
}

FailureOr<BitcastEncodingInfo>
EncodingAttr::convertForBitcast(ArrayRef<int64_t> shape,
                                Type originalElementType,
                                Type newElementType) const {
  FailureOr<SmallVector<int64_t>> newShape =
      computeBitcastShape(shape, originalElementType, newElementType);
  if (failed(newShape)) {
    return failure();
  }

  // Preserve the original element type if not already set, or keep the
  // existing one if this is a chain of bitcasts.
  Type origElemType = getOriginalElementType()
                          ? getOriginalElementType().getValue()
                          : originalElementType;
  EncodingAttr newEncoding = EncodingAttr::get(
      getContext(), getOperandIndex().getValue().getZExtValue(),
      getOpType().getValue(), getElementTypesArray(), origElemType,
      getRootMaps(), getIterationSizesArray());
  return BitcastEncodingInfo{*newShape, newEncoding};
}

std::optional<SmallVector<int32_t>> EncodingAttr::getReductionDims() const {
  if (!getUserIndexingMaps()) {
    return std::nullopt;
  }
  SmallVector<AffineMap> indexingMaps = getRootMaps();
  FailureOr<linalg::ContractionDimensions> contractionDims =
      linalg::inferContractionDims(indexingMaps);
  if (failed(contractionDims)) {
    return std::nullopt;
  }
  SmallVector<int32_t> result;
  for (unsigned k : contractionDims->k) {
    if (std::optional<unsigned> dimIdx = mapDimToOperandIndex(k)) {
      result.push_back(dimIdx.value());
    }
  }
  return result;
}

//===---------------------------------------------------------------------===//
// iree_encoding.padding
//===---------------------------------------------------------------------===//

/// Custom printer/parser methods to handle dynamic shapes.
ParseResult parsePadding(AsmParser &parser, DenseI64ArrayAttr &padding) {
  FailureOr<SmallVector<int64_t>> integerVals =
      parseDynamicI64IntegerList(parser);
  if (failed(integerVals)) {
    return failure();
  }
  padding = DenseI64ArrayAttr::get(parser.getContext(), integerVals.value());
  return success();
}
void printPadding(AsmPrinter &printer, DenseI64ArrayAttr padding) {
  return printDynamicI64IntegerList(printer, padding.asArrayRef());
}

PaddingAttr PaddingAttr::get(MLIRContext *ctx, ArrayRef<int64_t> padding) {
  return get(ctx, DenseI64ArrayAttr::get(ctx, padding));
}

PaddingAttr PaddingAttr::getIdentityAttr(MLIRContext *ctx, int rank) {
  SmallVector<int64_t> zeros(rank, 0);
  return get(ctx, zeros);
}

Attribute PaddingAttr::cloneWithLayouts(ArrayRef<Attribute> layouts) const {
  MLIRContext *ctx = getContext();
  return LayoutAttr::get(ctx, ArrayAttr::get(ctx, layouts));
}

bool PaddingAttr::isSerialized() const {
  return ShapedType::isStaticShape(getPadding().asArrayRef());
}

bool PaddingAttr::isIdentityLayout() const {
  ArrayRef<int64_t> padding = getPadding().asArrayRef();
  return llvm::all_of(padding, [](int64_t val) { return val == 0; });
}

Value PaddingAttr::calculateStorageSizeInBytes(Location loc, OpBuilder &builder,
                                               RankedTensorType type,
                                               ValueRange dynamicDims) const {
  ArrayRef<int64_t> padding = getPadding().asArrayRef();
  assert(padding.size() == type.getRank() && "Invalid padding");
  LLVM_DEBUG(if (llvm::any_of(padding, [](int64_t x) { return x != 0; })) {
    llvm::dbgs() << "Non-zero padding: " << type << "\n";
  });

  const int64_t elementSize = getRoundedElementByteWidth(type.getElementType());
  int64_t staticProduct = elementSize;
  Value dynamicProduct = arith::ConstantIndexOp::create(builder, loc, 1);

  size_t dynamicDimIdx = 0;
  for (auto [dimSize, padValue] : llvm::zip_equal(type.getShape(), padding)) {
    if (ShapedType::isStatic(dimSize)) {
      staticProduct *= (dimSize + padValue);
      continue;
    }

    Value dynamicDimSize = dynamicDims[dynamicDimIdx];
    ++dynamicDimIdx;

    if (padValue != 0) {
      dynamicDimSize = arith::AddIOp::create(
          builder, loc, dynamicDimSize,
          arith::ConstantIndexOp::create(builder, loc, padValue),
          arith::IntegerOverflowFlags::nsw);
    }
    dynamicProduct = builder.createOrFold<arith::MulIOp>(
        loc, dynamicProduct, dynamicDimSize, arith::IntegerOverflowFlags::nsw);
  }

  return builder.createOrFold<arith::MulIOp>(
      loc, arith::ConstantIndexOp::create(builder, loc, staticProduct),
      dynamicProduct, arith::IntegerOverflowFlags::nsw);
}

LogicalResult PaddingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                                  DenseI64ArrayAttr padding) {
  // You can only verify that the value is non-negative or dynamic.
  if (!llvm::all_of(padding.asArrayRef(), [](int64_t val) {
        return val == ShapedType::kDynamic || val >= 0;
      })) {
    return emitError() << "expected all padding values need to be "
                          "non-negative or dynamic";
  }
  return success();
}

//===---------------------------------------------------------------------===//
// iree_encoding.identity
//===---------------------------------------------------------------------===//

Value IdentityAttr::calculateStorageSizeInBytes(Location loc,
                                                OpBuilder &builder,
                                                RankedTensorType type,
                                                ValueRange dynamicDims) const {
  const int64_t elementSize = getRoundedElementByteWidth(type.getElementType());
  int64_t staticProduct = elementSize;
  Value dynamicProduct = arith::ConstantIndexOp::create(builder, loc, 1);

  size_t dynamicDimIdx = 0;
  for (int64_t dimSize : type.getShape()) {
    if (ShapedType::isStatic(dimSize)) {
      staticProduct *= dimSize;
      continue;
    }

    Value dynamicDimSize = dynamicDims[dynamicDimIdx];
    ++dynamicDimIdx;
    dynamicProduct = builder.createOrFold<arith::MulIOp>(
        loc, dynamicProduct, dynamicDimSize, arith::IntegerOverflowFlags::nsw);
  }
  return builder.createOrFold<arith::MulIOp>(
      loc, arith::ConstantIndexOp::create(builder, loc, staticProduct),
      dynamicProduct, arith::IntegerOverflowFlags::nsw);
}

bool IdentityAttr::isIdentityLayout() const { return true; }

bool IdentityAttr::isSerialized() const { return true; }

//===---------------------------------------------------------------------===//
// iree_encoding.identity_resolver
//===---------------------------------------------------------------------===//

Attribute
IdentityResolverAttr::cloneWithSimplifiedConfig(DictionaryAttr) const {
  return *this;
}

Attribute IdentityResolverAttr::getLayout(RankedTensorType type) const {
  return Encoding::IdentityAttr::get(getContext());
}

Attribute
IdentityResolverAttr::getUnifiedEncoding(ArrayRef<Attribute> encodings) const {
  MLIRContext *ctx = getContext();
  return Encoding::IdentityAttr::get(ctx);
}

Type IdentityResolverAttr::convertType(Type type) const {
  using IREE::TensorExt::DispatchTensorType;
  return TypeSwitch<Type, Type>(type)
      .Case([&](RankedTensorType rankedTensorType) {
        return rankedTensorType.dropEncoding();
      })
      .Case([&](DispatchTensorType dispatchTensorType) {
        auto boundType =
            dyn_cast<RankedTensorType>(dispatchTensorType.getBoundType());
        if (!boundType || !boundType.getEncoding()) {
          return dispatchTensorType;
        }
        Type convertedBoundType = this->convertType(boundType);
        return DispatchTensorType::get(dispatchTensorType.getAccess(),
                                       convertedBoundType);
      })
      .Default([&](auto concreteType) { return concreteType; });
}

LogicalResult IdentityResolverAttr::getOffsetsSizesStrides(
    OpBuilder &builder, Location loc, IREE::TensorExt::DispatchTensorType type,
    ValueRange dynamicDims, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides,
    SmallVectorImpl<OpFoldResult> &newOffsets,
    SmallVectorImpl<OpFoldResult> &newSizes,
    SmallVectorImpl<OpFoldResult> &newStrides) const {
  newSizes.assign(sizes.begin(), sizes.end());
  newOffsets.assign(offsets.begin(), offsets.end());
  newStrides.assign(strides.begin(), strides.end());
  return success();
}

Operation *IdentityResolverAttr::lowerOp(OpBuilder &b, Operation *op,
                                         TypeRange convertedResTypes,
                                         ValueRange convertedOperands) const {
  return clone(b, op, convertedResTypes, convertedOperands);
}

//===---------------------------------------------------------------------===//
// iree_encoding.unsupported_resolver
//===---------------------------------------------------------------------===//

Attribute
UnsupportedResolverAttr::cloneWithSimplifiedConfig(DictionaryAttr) const {
  return *this;
}

Attribute UnsupportedResolverAttr::getLayout(RankedTensorType) const {
  return nullptr;
}

//===---------------------------------------------------------------------===//
// Encoding attributes that are mainly for testing purpose.
//===---------------------------------------------------------------------===//

bool TestingAttr::isSerialized() const { return getLayouts() ? true : false; }

Attribute TestingAttr::cloneWithLayouts(ArrayRef<Attribute> layouts) const {
  MLIRContext *ctx = getContext();
  return TestingAttr::get(ctx, ArrayAttr::get(ctx, layouts),
                          getOriginalElementType());
}

FailureOr<BitcastEncodingInfo>
TestingAttr::convertForBitcast(ArrayRef<int64_t> shape,
                               Type originalElementType,
                               Type newElementType) const {
  FailureOr<SmallVector<int64_t>> newShape =
      computeBitcastShape(shape, originalElementType, newElementType);
  if (failed(newShape)) {
    return failure();
  }

  // Preserve the original element type if not already set, or keep the
  // existing one if this is a chain of bitcasts.
  Type origElemType =
      getOriginalElementType() ? getOriginalElementType() : originalElementType;
  auto newAttr = TestingAttr::get(getContext(), getLayouts(), origElemType);
  return BitcastEncodingInfo{*newShape, newAttr};
}

Attribute
SpecializationResolverAttr::cloneWithSimplifiedConfig(DictionaryAttr) const {
  return *this;
}

Attribute SpecializationResolverAttr::getLayout(RankedTensorType type) const {
  MLIRContext *ctx = getContext();
  return SpecializedAttr::get(ctx, getSeed(),
                              TypeAttr::get(type.dropEncoding()));
}

Attribute SpecializationResolverAttr::getUnifiedEncoding(
    ArrayRef<Attribute> encodings) const {
  MLIRContext *ctx = getContext();
  return SpecializedAttr::get(ctx, getSeed(), /*type=*/nullptr);
}

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

using namespace mlir::iree_compiler::IREE::Encoding;

#define GET_ATTRDEF_CLASSES
#include "iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp.inc"

void IREEEncodingDialect::registerAttributes() {
  addAttributes<
#define GET_ATTRDEF_LIST
#include "iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp.inc"
      >();
}
