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

//===---------------------------------------------------------------------===//
// Pass to materialize the encoding of tensor based on target information.
//===---------------------------------------------------------------------===//

#include "iree/compiler/Codegen/Common/EncodingUtils.h"
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
#include "iree/compiler/Codegen/Utils/EncodingUtils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"

#define DEBUG_TYPE "iree-codegen-materialize-encoding"

namespace mlir::iree_compiler {

using IREE::Codegen::MaterializeEncodingInfo;
using IREE::Codegen::TileSwizzle;

//===---------------------------------------------------------------------===//
// Methods to convert `set_encoding` and `unset_encoding` operations
// to `pack` and `unpack` operations respectively.
//===---------------------------------------------------------------------===//

FailureOr<Value> lowerSetEncodingOpToPackOp(
    RewriterBase &rewriter, IREE::Encoding::SetEncodingOp encodingOp,
    Value source, const MaterializeEncodingTypeConverter &typeConverter) {
  RankedTensorType resultType = encodingOp.getResultType();
  MaterializeEncodingInfo encodingInfo =
      typeConverter.getEncodingInfo(resultType);

  // Shortcut to avoid creating new operations.
  if (IREE::Codegen::isIdentityLayout(encodingInfo)) {
    return source;
  }

  // Create `tensor.empty` operation for the result of the pack operation.
  Location loc = encodingOp.getLoc();
  FailureOr<SmallVector<OpFoldResult>> innerTileSizesOfr =
      typeConverter.getInnerTileSizesOfr(rewriter, loc, resultType,
                                         encodingInfo);
  if (failed(innerTileSizesOfr)) {
    return rewriter.notifyMatchFailure(
        encodingOp, "failed to generate runtime tile size query");
  }
  Value paddingValue = arith::ConstantOp::create(
      rewriter, loc, rewriter.getZeroAttr(resultType.getElementType()));
  SmallVector<OpFoldResult> sourceDims =
      tensor::getMixedSizes(rewriter, loc, source);
  SmallVector<OpFoldResult> resultDims = linalg::PackOp::getResultShape(
      rewriter, loc, sourceDims, *innerTileSizesOfr, encodingInfo.innerDimsPos,
      encodingInfo.outerDimsPerm);
  auto emptyOp = tensor::EmptyOp::create(rewriter, loc, resultDims,
                                         resultType.getElementType());
  return linalg::PackOp::create(rewriter, loc, source, emptyOp,
                                encodingInfo.innerDimsPos, *innerTileSizesOfr,
                                paddingValue, encodingInfo.outerDimsPerm)
      .getResult();
}

FailureOr<Value> lowerUnsetEncodingToUnpackOp(
    RewriterBase &rewriter, IREE::Encoding::UnsetEncodingOp encodingOp,
    Value packedValue, const MaterializeEncodingTypeConverter &typeConverter) {
  RankedTensorType sourceType = encodingOp.getSourceType();
  MaterializeEncodingInfo encodingInfo =
      typeConverter.getEncodingInfo(sourceType);

  // Shortcut to avoid creating new operations.
  if (IREE::Codegen::isIdentityLayout(encodingInfo)) {
    return packedValue;
  }

  // Create an `tensor.empty` for the result of the unpack operation.
  Location loc = encodingOp.getLoc();
  SmallVector<OpFoldResult> resultDims =
      getMixedValues(encodingOp.getResultType().getShape(),
                     encodingOp.getResultDims(), rewriter);
  auto emptyOp = tensor::EmptyOp::create(rewriter, loc, resultDims,
                                         sourceType.getElementType());
  FailureOr<SmallVector<OpFoldResult>> innerTileSizesOfr =
      typeConverter.getInnerTileSizesOfr(rewriter, loc, sourceType,
                                         encodingInfo);
  if (failed(innerTileSizesOfr)) {
    return rewriter.notifyMatchFailure(
        encodingOp, "failed to generate runtime tile size query");
  }
  return linalg::UnPackOp::create(rewriter, loc, packedValue, emptyOp,
                                  encodingInfo.innerDimsPos, *innerTileSizesOfr,
                                  encodingInfo.outerDimsPerm)
      .getResult();
}

/// Utility method to convert `tensor.empty` with encoding to a `tensor.empty`
/// of the materialized type.
static FailureOr<Operation *>
lowerOpWithEncoding(RewriterBase &rewriter, tensor::EmptyOp emptyOp,
                    ValueRange convertedOperands,
                    const MaterializeEncodingTypeConverter &typeConverter) {
  auto emptyType = cast<RankedTensorType>(emptyOp->getResultTypes()[0]);
  MaterializeEncodingInfo encodingInfo =
      typeConverter.getEncodingInfo(emptyType);
  Location loc = emptyOp.getLoc();
  if (IREE::Codegen::isIdentityLayout(encodingInfo)) {
    return tensor::EmptyOp::create(rewriter, loc, emptyOp.getMixedSizes(),
                                   emptyType.getElementType())
        .getOperation();
  }

  FailureOr<SmallVector<OpFoldResult>> innerTileSizesOfr =
      typeConverter.getInnerTileSizesOfr(rewriter, loc, emptyType,
                                         encodingInfo);
  if (failed(innerTileSizesOfr)) {
    return rewriter.notifyMatchFailure(
        emptyOp, "failed to generate runtime tile size query");
  }

  SmallVector<OpFoldResult> sourceDims = emptyOp.getMixedSizes();
  (void)foldDynamicIndexList(sourceDims);
  SmallVector<OpFoldResult> newShape = linalg::PackOp::getResultShape(
      rewriter, loc, sourceDims, *innerTileSizesOfr, encodingInfo.innerDimsPos,
      encodingInfo.outerDimsPerm);
  newShape = getSwizzledShape(newShape, encodingInfo);
  Operation *newEmptyOp = tensor::EmptyOp::create(rewriter, loc, newShape,
                                                  emptyType.getElementType());
  return newEmptyOp;
}

namespace {
/// Pattern to materialize the encoding for `hal.interface.binding.subspan`
/// operations.
struct MaterializeInterfaceBindingEncoding
    : OpConversionPattern<IREE::HAL::InterfaceBindingSubspanOp> {
  using OpConversionPattern<
      IREE::HAL::InterfaceBindingSubspanOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(IREE::HAL::InterfaceBindingSubspanOp subspanOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto origResultType = dyn_cast<IREE::TensorExt::DispatchTensorType>(
        subspanOp.getResult().getType());
    if (!origResultType) {
      return rewriter.notifyMatchFailure(
          subspanOp,
          "expected result type to be !iree_tensor_ext.dispatch.tensor");
    }
    auto origBoundTensorType =
        dyn_cast<RankedTensorType>(origResultType.getBoundType());
    if (!origBoundTensorType) {
      return rewriter.notifyMatchFailure(
          subspanOp, "bound type is not a RankedTensorType");
    }

    auto typeConverter = getTypeConverter<MaterializeEncodingTypeConverter>();
    auto convertedResultType =
        typeConverter->convertType<IREE::TensorExt::DispatchTensorType>(
            origResultType);
    if (!convertedResultType) {
      return rewriter.notifyMatchFailure(subspanOp,
                                         "expected converted result type to be "
                                         "!iree_tensor_ext.dispatch.tensor");
    }
    if (origResultType == convertedResultType) {
      return rewriter.notifyMatchFailure(
          subspanOp, "DispatchTensorType type already valid");
    }

    // Get the dynamic dims of the target.
    // TODO(hanchung): We only have getOffsetsSizesStrides interface method that
    // handles all three together. It would be cleaner to have a separate method
    // to get dynamic sizes only.
    Location loc = subspanOp.getLoc();
    ValueRange origDynamicDims = subspanOp.getDynamicDims();
    SmallVector<OpFoldResult> origSizes = getMixedValues(
        origBoundTensorType.getShape(), origDynamicDims, rewriter);
    SmallVector<OpFoldResult> origOffsets(origDynamicDims.size(),
                                          rewriter.getIndexAttr(0));
    SmallVector<OpFoldResult> origStrides(origDynamicDims.size(),
                                          rewriter.getIndexAttr(1));
    SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
    if (failed(typeConverter->getOffsetsSizesStrides(
            rewriter, loc, origResultType, origDynamicDims, origOffsets,
            origSizes, origStrides, newOffsets, newSizes, newStrides))) {
      return failure();
    }

    SmallVector<int64_t> newStaticDims;
    SmallVector<Value> newDynamicDims;
    dispatchIndexOpFoldResults(newSizes, newDynamicDims, newStaticDims);
    rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceBindingSubspanOp>(
        subspanOp, convertedResultType, subspanOp.getLayout(),
        subspanOp.getBinding(), subspanOp.getByteOffset(), newDynamicDims,
        subspanOp.getAlignmentAttr(), subspanOp.getDescriptorFlagsAttr());
    return success();
  }
};

/// Pattern to convert `iree_tensor_ext.dispatch.tensor.load` operation when
/// materializing the encoding.
struct MaterializeTensorExtDispatchTensorLoadOp
    : OpConversionPattern<IREE::TensorExt::DispatchTensorLoadOp> {
  using OpConversionPattern<
      IREE::TensorExt::DispatchTensorLoadOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(IREE::TensorExt::DispatchTensorLoadOp loadOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto sourceType = loadOp.getSourceType();
    auto boundTensorType = cast<RankedTensorType>(sourceType.getBoundType());
    auto typeConverter = getTypeConverter<MaterializeEncodingTypeConverter>();
    if (typeConverter->convertType(boundTensorType) == boundTensorType) {
      return rewriter.notifyMatchFailure(loadOp, "bound type already valid");
    }

    SmallVector<OpFoldResult> newOffsets, newMixedSizes, newStrides;
    if (failed(typeConverter->getOffsetsSizesStrides(
            rewriter, loadOp.getLoc(), sourceType, loadOp.getSourceDims(),
            loadOp.getMixedOffsets(), loadOp.getMixedSizes(),
            loadOp.getMixedStrides(), newOffsets, newMixedSizes, newStrides))) {
      return failure();
    }
    SmallVector<int64_t> newStaticDims;
    SmallVector<Value> newDynamicDims;
    dispatchIndexOpFoldResults(newMixedSizes, newDynamicDims, newStaticDims);
    rewriter.replaceOpWithNewOp<IREE::TensorExt::DispatchTensorLoadOp>(
        loadOp, adaptor.getSource(), newDynamicDims, newOffsets, newMixedSizes,
        newStrides);
    return success();
  }
};

/// Pattern to convert `iree_tensor_ext.dispatch.tensor.store` operation when
/// materializing the encoding.
struct MaterializeTensorExtDispatchTensorStoreOp
    : OpConversionPattern<IREE::TensorExt::DispatchTensorStoreOp> {
  using OpConversionPattern<
      IREE::TensorExt::DispatchTensorStoreOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(IREE::TensorExt::DispatchTensorStoreOp storeOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto targetType = storeOp.getTargetType();
    auto boundTensorType = cast<RankedTensorType>(targetType.getBoundType());
    auto typeConverter = getTypeConverter<MaterializeEncodingTypeConverter>();
    if (typeConverter->convertType(boundTensorType) == boundTensorType) {
      return rewriter.notifyMatchFailure(storeOp, "bound type already valid");
    }

    SmallVector<OpFoldResult> newOffsets, newMixedSizes, newStrides;
    if (failed(typeConverter->getOffsetsSizesStrides(
            rewriter, storeOp.getLoc(), targetType, storeOp.getTargetDims(),
            storeOp.getMixedOffsets(), storeOp.getMixedSizes(),
            storeOp.getMixedStrides(), newOffsets, newMixedSizes,
            newStrides))) {
      return failure();
    }
    SmallVector<int64_t> newStaticDims;
    SmallVector<Value> newDynamicDims;
    dispatchIndexOpFoldResults(newMixedSizes, newDynamicDims, newStaticDims);
    rewriter.replaceOpWithNewOp<IREE::TensorExt::DispatchTensorStoreOp>(
        storeOp, adaptor.getValue(), adaptor.getTarget(), newDynamicDims,
        newOffsets, newMixedSizes, newStrides);
    return success();
  }
};

//===---------------------------------------------------------------------===//
// Patterns for layout transfers. They decompose load/store ops into
// set_encoding/unset_encoding + load/store, if the converted types mismatch.
//===---------------------------------------------------------------------===//

/// Returns the value that brings `src` to `destType` by inserting the necessary
/// encoding ops.
static Value generateEncodingTransferOps(RewriterBase &rewriter, Value src,
                                         ArrayRef<Value> dynamicDims,
                                         RankedTensorType destType) {
  auto srcType = cast<RankedTensorType>(src.getType());
  if (srcType == destType) {
    return src;
  }
  Value value = src;
  if (srcType.getEncoding()) {
    value = IREE::Encoding::UnsetEncodingOp::create(
        rewriter, src.getLoc(), srcType.dropEncoding(), value, dynamicDims,
        /*encoding_dims=*/ValueRange{});
  }
  if (destType.getEncoding()) {
    value = IREE::Encoding::SetEncodingOp::create(
        rewriter, src.getLoc(), destType, value,
        /*encoding_dims=*/ValueRange{});
  }
  return value;
}

/// Pattern to convert `iree_tensor_ext.dispatch.tensor.load` operation when
/// materializing the encoding.
struct DecomposeMismatchEncodingTensorLoadOp
    : OpRewritePattern<IREE::TensorExt::DispatchTensorLoadOp> {
  using Base::Base;

  DecomposeMismatchEncodingTensorLoadOp(
      MaterializeEncodingTypeConverter &converter, MLIRContext *ctx,
      PatternBenefit benefit = 0)
      : OpRewritePattern<IREE::TensorExt::DispatchTensorLoadOp>(ctx, benefit),
        typeConverter(converter) {}

  LogicalResult matchAndRewrite(IREE::TensorExt::DispatchTensorLoadOp loadOp,
                                PatternRewriter &rewriter) const override {
    if (!loadOp.isLoadOfWholeSource()) {
      return rewriter.notifyMatchFailure(loadOp, "unhandled partial loads");
    }

    IREE::TensorExt::DispatchTensorType srcType = loadOp.getSourceType();
    auto boundTensorType = dyn_cast<RankedTensorType>(srcType.getBoundType());
    if (!boundTensorType) {
      return rewriter.notifyMatchFailure(
          loadOp, "source bound type is not a RankedTensorType");
    }

    // Only decompose if there's an encoding involved. If neither the source
    // nor the destination has an encoding, this pattern should not match.
    // This can happen when isLoadOfWholeSource() returns true but the load
    // reshapes the tensor (e.g., loading a 4D tensor from a 5D source).
    RankedTensorType destType = loadOp.getResult().getType();
    if (!boundTensorType.getEncoding() && !destType.getEncoding()) {
      return rewriter.notifyMatchFailure(
          loadOp, "no encoding involved in source or destination");
    }

    // We have to check the bound type from converted DispatchTensorType because
    // it is what we'll see in encoding materialization. E.g.,
    // GPUPaddingResolver converts RankedTensorType into the same type, but it
    // creates different IREE::TensorExt::DispatchTensorType that may have
    // larger tensor shape for bound type.
    auto convertedSrcType =
        typeConverter.convertType<IREE::TensorExt::DispatchTensorType>(srcType);
    if (typeConverter.convertType(convertedSrcType.getBoundType()) ==
        typeConverter.convertType(destType)) {
      return rewriter.notifyMatchFailure(
          loadOp, "the source type and the result type match after conversion");
    }

    LDBG() << "Performance warning: decomposing mismatched encoding load op: "
           << loadOp;
    Location loc = loadOp.getLoc();
    Value result = IREE::TensorExt::DispatchTensorLoadOp::create(
        rewriter, loc, boundTensorType, loadOp.getSource(),
        loadOp.getSourceDims(), loadOp.getMixedOffsets(),
        loadOp.getMixedSizes(), loadOp.getMixedStrides());
    SmallVector<Value> dynamicDims = llvm::to_vector(loadOp.getSizes());
    result =
        generateEncodingTransferOps(rewriter, result, dynamicDims, destType);
    rewriter.replaceOp(loadOp, result);
    return success();
  }

private:
  MaterializeEncodingTypeConverter &typeConverter;
};

/// Pattern to convert `iree_tensor_ext.dispatch.tensor.store` operation when
/// materializing the encoding.
struct DecomposeMismatchEncodingTensorStoreOp
    : OpRewritePattern<IREE::TensorExt::DispatchTensorStoreOp> {
  using Base::Base;

  DecomposeMismatchEncodingTensorStoreOp(
      MaterializeEncodingTypeConverter &converter, MLIRContext *ctx,
      PatternBenefit benefit = 0)
      : OpRewritePattern<IREE::TensorExt::DispatchTensorStoreOp>(ctx, benefit),
        typeConverter(converter) {}

  LogicalResult matchAndRewrite(IREE::TensorExt::DispatchTensorStoreOp storeOp,
                                PatternRewriter &rewriter) const override {
    if (!storeOp.isStoreToWholeTarget()) {
      return rewriter.notifyMatchFailure(storeOp, "unhandled partial stores");
    }

    IREE::TensorExt::DispatchTensorType targetType = storeOp.getTargetType();
    auto boundTensorType =
        dyn_cast<RankedTensorType>(targetType.getBoundType());
    if (!boundTensorType) {
      return rewriter.notifyMatchFailure(
          storeOp, "target bound type is not a RankedTensorType");
    }

    // Only decompose if there's an encoding involved. If neither the value
    // nor the target has an encoding, this pattern should not match.
    // This can happen when isStoreToWholeTarget() returns true but the store
    // reshapes the tensor (e.g., storing a 4D tensor to a 5D target).
    RankedTensorType valueType = storeOp.getValue().getType();
    if (!boundTensorType.getEncoding() && !valueType.getEncoding()) {
      return rewriter.notifyMatchFailure(
          storeOp, "no encoding involved in value or target");
    }

    // Similar to DecomposeMismatchEncodingTensorLoadOp, we have to check with
    // the bound type from converted DispatchTensorType.
    auto convertedTargetType =
        typeConverter.convertType<IREE::TensorExt::DispatchTensorType>(
            targetType);
    if (typeConverter.convertType(convertedTargetType.getBoundType()) ==
        typeConverter.convertType(valueType)) {
      return rewriter.notifyMatchFailure(
          storeOp, "the value type and the target type match");
    }

    LDBG() << "Performance warning: decomposing mismatched encoding store op: "
           << storeOp;
    Location loc = storeOp.getLoc();
    Value valueToStore = storeOp.getValue();
    SmallVector<Value> dynamicDims = llvm::to_vector(storeOp.getSizes());
    valueToStore = generateEncodingTransferOps(rewriter, valueToStore,
                                               dynamicDims, boundTensorType);
    IREE::TensorExt::DispatchTensorStoreOp::create(
        rewriter, loc, valueToStore, storeOp.getTarget(),
        storeOp.getTargetDims(), storeOp.getMixedOffsets(),
        storeOp.getMixedSizes(), storeOp.getMixedStrides());
    rewriter.eraseOp(storeOp);
    return success();
  }

private:
  MaterializeEncodingTypeConverter &typeConverter;
};

//===---------------------------------------------------------------------===//
// Patterns to lower ops with encodings. These are written as
// dialect conversion patterns for now. These are just drivers around
// the core conversion utilities.
//===---------------------------------------------------------------------===//

/// Generic pattern to convert an operation.
template <typename OpTy>
struct MaterializeOperation : OpConversionPattern<OpTy> {
  using OpConversionPattern<OpTy>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto converter =
        this->template getTypeConverter<MaterializeEncodingTypeConverter>();
    FailureOr<Operation *> convertedOp =
        lowerOpWithEncoding(rewriter, op, adaptor.getOperands(), *converter);
    if (failed(convertedOp)) {
      return failure();
    }

    rewriter.replaceOp(op, convertedOp.value());
    return success();
  }
};

struct MaterializeOptimizationBarrierOp
    : OpConversionPattern<IREE::Util::OptimizationBarrierOp> {
  using OpConversionPattern<
      IREE::Util::OptimizationBarrierOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(IREE::Util::OptimizationBarrierOp op,
                  IREE::Util::OptimizationBarrierOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (llvm::none_of(op.getOperandTypes(), [](Type type) -> bool {
          auto tensorType = dyn_cast<RankedTensorType>(type);
          return tensorType && tensorType.getEncoding();
        })) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<IREE::Util::OptimizationBarrierOp>(
        op, adaptor.getOperands());
    return success();
  }
};

static SmallVector<ReassociationIndices>
getReassociationIndices(int outerDims,
                        const TileSwizzle::ExpandShapeType &expandShape) {
  SmallVector<ReassociationIndices> result;
  int expandedIdx = 0;
  for (int i = 0; i < outerDims; ++i) {
    result.push_back({expandedIdx++});
  }
  for (auto expandShapeDim : expandShape) {
    result.push_back({});
    for (int i = 0, e = expandShapeDim.size(); i < e; ++i) {
      result.back().push_back(expandedIdx++);
    }
  }
  return result;
}

/// Returns true when the swizzle's `permutation` only reorders unit dims
/// relative to non-unit dims, i.e. it preserves the relative order of all
/// non-unit dims. In that case the unpermuted and permuted inner-tile shapes
/// have the same linearized memory layout, so the `linalg.transpose` that the
/// pack/unpack lowering would otherwise emit moves no data.
static bool isSwizzlePermutationLayoutNoOp(const TileSwizzle &swizzle) {
  SmallVector<int64_t> expandedTileShape =
      IREE::Codegen::getExpandedTileShape(swizzle.expandShape());
  // `swizzle.permutation()[i]` is the unpermuted `expandedTileShape` index of
  // the dim landing at permuted position `i`. Unit dims span no memory, so
  // only the non-unit dims matter: the layout is preserved iff they keep
  // their source order, i.e. `srcDim` increases across just the non-unit dims.
  int64_t prevNonUnitSrc = -1;
  for (int64_t srcDim : swizzle.permutation()) {
    if (expandedTileShape[srcDim] == 1) {
      continue;
    }
    if (srcDim < prevNonUnitSrc) {
      return false;
    }
    prevNonUnitSrc = srcDim;
  }
  return true;
}

/// Tries to build the reassociation for the single reshape that replaces the
/// swizzle's expand_shape + transpose when the transpose is a layout no-op:
/// from the packed shape (outer + `innerTileSizes`) to the permuted-inner
/// shape (outer + `applyPermutation(expandedTileShape, permutation)`).
///
/// The permuted-inner shape keeps the swizzle's unit dims: it is the
/// materialized encoding type fixed by `getEncodingInfo`, which the reshape
/// must match exactly for dialect conversion. Dropping unit dims is not a
/// choice this lowering can make locally.
///
/// Intended to be called only when `isSwizzlePermutationLayoutNoOp(swizzle)`,
/// but that predicate is not sufficient: it guarantees the linearized layout
/// is preserved, but a layout-no-op transpose isn't always expressible as a
/// single reshape. Specifically, a later group's non-unit dim can land
/// positionally before an earlier (all-unit) group's content, leaving no
/// valid ascending-order reassociation. Returns `failure()` in that case so
/// the caller falls back to the general expand + transpose lowering.
static FailureOr<SmallVector<ReassociationIndices>>
getSwizzleFoldedReassociation(int outerDims, ArrayRef<int64_t> innerTileSizes,
                              const TileSwizzle &swizzle) {
  SmallVector<int64_t> permutedInnerShape =
      IREE::Codegen::getExpandedTileShape(swizzle.expandShape());
  applyPermutationToVector(permutedInnerShape, swizzle.permutation());

  SmallVector<ReassociationIndices> result;
  for (int i = 0; i < outerDims; ++i) {
    result.push_back({i});
  }
  int targetIdx = 0;
  for (int64_t innerSize : innerTileSizes) {
    ReassociationIndices group;
    int64_t accum = 1;
    // Grow the slice until its product reaches `innerSize`. `do/while` (not
    // `while`) consumes at least one dim, so `innerSize == 1` slices aren't
    // empty. Fail if a single dim overshoots `innerSize` (a later group's
    // non-unit dim positioned before this group's content) or we run out of
    // dims.
    do {
      if (targetIdx >= static_cast<int>(permutedInnerShape.size())) {
        return failure();
      }
      group.push_back(outerDims + targetIdx);
      accum *= permutedInnerShape[targetIdx];
      ++targetIdx;
    } while (accum < innerSize);
    if (accum != innerSize) {
      return failure();
    }
    result.push_back(std::move(group));
  }
  // Trailing unit dims (after the last non-unit dim) go to the last slice.
  // Absorbing them per-slice instead would let an earlier slice eat unit dims
  // that a later, all-unit packed dim still needs.
  while (targetIdx < static_cast<int>(permutedInnerShape.size())) {
    if (permutedInnerShape[targetIdx] != 1) {
      return failure();
    }
    result.back().push_back(outerDims + targetIdx);
    ++targetIdx;
  }
  return result;
}

/// Convert iree_linalg_ext.set_encoding op to pack + tile swizzling ops. We use
/// expand_shape + linalg.transpose to represent a tile swizzling op.
struct SetEncodingOpLoweringConversion
    : OpConversionPattern<IREE::Encoding::SetEncodingOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(IREE::Encoding::SetEncodingOp encodingOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (encodingOp.getSource().getType().getRank() == 0) {
      rewriter.replaceOp(encodingOp, adaptor.getSource());
      return success();
    }
    auto converter = getTypeConverter<MaterializeEncodingTypeConverter>();
    auto packedValue = lowerSetEncodingOpToPackOp(
        rewriter, encodingOp, adaptor.getSource(), *converter);
    if (failed(packedValue)) {
      rewriter.replaceOp(encodingOp, adaptor.getSource());
      return success();
    }

    MaterializeEncodingInfo encodingInfo =
        converter->getEncodingInfo(encodingOp.getResultType());
    if (!encodingInfo.swizzle) {
      rewriter.replaceOp(encodingOp, packedValue.value());
      return success();
    }

    Location loc = encodingOp.getLoc();
    int origRank = encodingOp.getSourceType().getRank();

    // Final shape after the swizzle is applied: outer dims + permuted
    // expand-shape inner tile.
    SmallVector<int64_t> permutedInnerShape =
        getExpandedTileShape(encodingInfo.swizzle->expandShape());
    applyPermutationToVector(permutedInnerShape,
                             encodingInfo.swizzle->permutation());
    SmallVector<int64_t> finalShape(cast<ShapedType>(packedValue->getType())
                                        .getShape()
                                        .take_front(origRank));
    finalShape.append(permutedInnerShape);
    RankedTensorType finalType = encodingOp.getSourceType().clone(finalShape);

    // Fast path: when the permutation is a layout no-op and the fold is
    // expressible as a single reshape, replace expand_shape + transpose with
    // one expand_shape straight to the permuted shape. This drops the
    // intermediate buffer that, under dynamic shapes, gets hoisted to a stack
    // alloca with a bogus index-umax-derived static size.
    if (isSwizzlePermutationLayoutNoOp(*encodingInfo.swizzle)) {
      if (auto reassoc = getSwizzleFoldedReassociation(
              origRank, encodingInfo.innerTileSizes, *encodingInfo.swizzle);
          succeeded(reassoc)) {
        auto expandShapeOp = tensor::ExpandShapeOp::create(
            rewriter, loc, finalType, packedValue.value(), *reassoc);
        rewriter.replaceOp(encodingOp, expandShapeOp.getResult());
        return success();
      }
    }

    // General case: expand_shape to the unpermuted inner tile, then
    // linalg.transpose to apply the swizzle permutation.
    SmallVector<int64_t> expandShapeShape(
        cast<ShapedType>(packedValue->getType())
            .getShape()
            .take_front(origRank));
    expandShapeShape.append(
        getExpandedTileShape(encodingInfo.swizzle->expandShape()));
    RankedTensorType expandShapeType =
        encodingOp.getSourceType().clone(expandShapeShape);

    SmallVector<ReassociationIndices> reassociation =
        getReassociationIndices(origRank, encodingInfo.swizzle->expandShape());
    auto expandShapeOp = tensor::ExpandShapeOp::create(
        rewriter, loc, expandShapeType, packedValue.value(), reassociation);

    SmallVector<int64_t> transposePerm =
        llvm::to_vector(llvm::seq<int64_t>(0, origRank));
    for (auto perm : encodingInfo.swizzle->permutation()) {
      transposePerm.push_back(origRank + perm);
    }
    SmallVector<OpFoldResult> transposeResultDims =
        tensor::getMixedSizes(rewriter, loc, expandShapeOp.getResult());
    applyPermutationToVector(transposeResultDims, transposePerm);

    auto emptyTensor =
        tensor::EmptyOp::create(rewriter, loc, transposeResultDims,
                                encodingOp.getSourceType().getElementType());
    auto transposeOp = linalg::TransposeOp::create(rewriter, loc, expandShapeOp,
                                                   emptyTensor, transposePerm);
    rewriter.replaceOp(encodingOp, transposeOp->getResult(0));

    return success();
  }
};

struct UnsetEncodingOpLoweringConversion
    : OpConversionPattern<IREE::Encoding::UnsetEncodingOp> {
  using OpConversionPattern<
      IREE::Encoding::UnsetEncodingOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(IREE::Encoding::UnsetEncodingOp unsetEncodingOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto converter = getTypeConverter<MaterializeEncodingTypeConverter>();
    MaterializeEncodingInfo encodingInfo =
        converter->getEncodingInfo(unsetEncodingOp.getSource().getType());
    if (IREE::Codegen::isIdentityLayout(encodingInfo)) {
      rewriter.replaceOp(unsetEncodingOp, adaptor.getSource());
      return success();
    }

    Location loc = unsetEncodingOp.getLoc();
    Value unpackSrc = adaptor.getSource();
    if (encodingInfo.swizzle) {
      int targetRank = unsetEncodingOp.getResultType().getRank();
      auto srcConvertedType =
          cast<RankedTensorType>(adaptor.getSource().getType());
      SmallVector<int64_t> unpackSrcShape(
          srcConvertedType.getShape().take_front(targetRank));
      unpackSrcShape.append(encodingInfo.innerTileSizes.begin(),
                            encodingInfo.innerTileSizes.end());
      RankedTensorType unpackSrcType =
          unsetEncodingOp.getResultType().clone(unpackSrcShape);

      // Fast path mirroring the SetEncoding side: fold the inverse transpose +
      // collapse_shape into a single collapse_shape when the permutation is a
      // layout no-op and the fold is a single reshape.
      FailureOr<SmallVector<ReassociationIndices>> foldedReassoc = failure();
      if (isSwizzlePermutationLayoutNoOp(*encodingInfo.swizzle)) {
        foldedReassoc = getSwizzleFoldedReassociation(
            targetRank, encodingInfo.innerTileSizes, *encodingInfo.swizzle);
      }
      if (succeeded(foldedReassoc)) {
        unpackSrc = tensor::CollapseShapeOp::create(
            rewriter, loc, unpackSrcType, adaptor.getSource(), *foldedReassoc);
      } else {
        SmallVector<OpFoldResult> emptyShape =
            tensor::getMixedSizes(rewriter, loc, adaptor.getSource());
        emptyShape.resize(targetRank);
        for (auto i :
             getExpandedTileShape(encodingInfo.swizzle->expandShape())) {
          emptyShape.push_back(rewriter.getIndexAttr(i));
        }
        auto emptyTensor = tensor::EmptyOp::create(
            rewriter, loc, emptyShape,
            unsetEncodingOp.getSourceType().getElementType());

        SmallVector<int64_t> transposePerm =
            llvm::to_vector(llvm::seq<int64_t>(0, targetRank));
        for (auto perm : encodingInfo.swizzle->permutation()) {
          transposePerm.push_back(targetRank + perm);
        }
        auto invertedTransposePerm = invertPermutationVector(transposePerm);
        auto transposeOp =
            linalg::TransposeOp::create(rewriter, loc, adaptor.getSource(),
                                        emptyTensor, invertedTransposePerm);

        SmallVector<ReassociationIndices> reassociation =
            getReassociationIndices(targetRank,
                                    encodingInfo.swizzle->expandShape());
        unpackSrc = tensor::CollapseShapeOp::create(
            rewriter, loc, unpackSrcType, transposeOp->getResult(0),
            reassociation);
      }
    }

    auto unpackedValue = lowerUnsetEncodingToUnpackOp(rewriter, unsetEncodingOp,
                                                      unpackSrc, *converter);
    if (failed(unpackedValue)) {
      rewriter.replaceOp(unsetEncodingOp, adaptor.getSource());
      return success();
    }
    rewriter.replaceOp(unsetEncodingOp, unpackedValue.value());
    return success();
  }
};

/// Pattern to rewrite linalg::LinalgOp by materializing its encoding using the
/// provided LayoutMaterializerAttr.
class MaterializeLinalgOp
    : public OpInterfaceConversionPattern<linalg::LinalgOp> {
public:
  MaterializeLinalgOp(const MaterializeEncodingTypeConverter &typeConverter,
                      MLIRContext *context, PatternBenefit benefit = 1)
      : OpInterfaceConversionPattern<linalg::LinalgOp>(typeConverter, context,
                                                       benefit) {}

  LogicalResult
  matchAndRewrite(linalg::LinalgOp op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    auto converter = getTypeConverter<MaterializeEncodingTypeConverter>();
    IREE::Encoding::LayoutMaterializerAttr layoutAttr =
        converter->getLayoutAttr();
    SmallVector<Type> convertedResTypes;
    for (auto init : op.getDpsInits()) {
      convertedResTypes.push_back(converter->convertType(init.getType()));
    }
    Operation *newOp =
        layoutAttr.lowerOp(rewriter, op, convertedResTypes, operands);
    if (!newOp) {
      return rewriter.notifyMatchFailure(
          op, "unsupported operation for encoding materialization");
    }
    rewriter.replaceOp(op, newOp->getResults());
    return success();
  }
};

static bool isRankedTensorTypeWithEncoding(Type type) {
  auto rankedTensorType = dyn_cast<RankedTensorType>(type);
  if (!rankedTensorType) {
    return false;
  }
  return rankedTensorType.getEncoding() ? true : false;
}

struct MaterializeFuncReturnOp final : OpConversionPattern<func::ReturnOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<func::ReturnOp>(op, adaptor.getOperands());
    return success();
  }
};

} // namespace

void populateDecomposeMismatchedLayoutLoadStoreOpsPatterns(
    RewritePatternSet &patterns,
    MaterializeEncodingTypeConverter &typeConverter) {
  patterns.insert<DecomposeMismatchEncodingTensorLoadOp,
                  DecomposeMismatchEncodingTensorStoreOp>(
      typeConverter, patterns.getContext());
}

void populateMaterializeEncodingPatterns(
    RewritePatternSet &patterns, MaterializeEncodingConversionTarget &target,
    MaterializeEncodingTypeConverter &typeConverter) {
  MLIRContext *context = patterns.getContext();
  target.addDynamicallyLegalOp<IREE::HAL::InterfaceBindingSubspanOp>(
      [&typeConverter](IREE::HAL::InterfaceBindingSubspanOp subspanOp) {
        auto resultType = dyn_cast<IREE::TensorExt::DispatchTensorType>(
            subspanOp.getResult().getType());
        // For types that are not `TensorExt::DispatchTensorType` mark as legal.
        if (!resultType) {
          return true;
        }
        return resultType == typeConverter.convertType(resultType);
      });
  target.addIllegalOp<IREE::Encoding::SetEncodingOp,
                      IREE::Encoding::UnsetEncodingOp>();
  target.addDynamicallyLegalOp<IREE::TensorExt::DispatchTensorStoreOp>(
      [&typeConverter](IREE::TensorExt::DispatchTensorStoreOp storeOp) {
        auto resultType = dyn_cast<IREE::TensorExt::DispatchTensorType>(
            storeOp.getTargetType());
        // For types that are not `TensorExt::DispatchTensorType` mark as legal.
        if (!resultType) {
          return true;
        }
        return resultType == typeConverter.convertType(resultType);
      });
  target.addDynamicallyLegalOp<IREE::TensorExt::DispatchTensorLoadOp>(
      [&typeConverter](IREE::TensorExt::DispatchTensorLoadOp loadOp) {
        auto resultType = dyn_cast<IREE::TensorExt::DispatchTensorType>(
            loadOp.getSourceType());
        // For types that are not `TensorExt::DispatchTensorType` mark as legal.
        if (!resultType) {
          return true;
        }
        return resultType == typeConverter.convertType(resultType);
      });
  target.addDynamicallyLegalOp<func::ReturnOp>([](func::ReturnOp returnOp) {
    return llvm::none_of(returnOp.getOperandTypes(),
                         isRankedTensorTypeWithEncoding);
  });

  patterns.insert<MaterializeLinalgOp, SetEncodingOpLoweringConversion,
                  UnsetEncodingOpLoweringConversion,
                  MaterializeOperation<tensor::EmptyOp>,
                  MaterializeOptimizationBarrierOp,
                  MaterializeTensorExtDispatchTensorLoadOp,
                  MaterializeTensorExtDispatchTensorStoreOp,
                  MaterializeInterfaceBindingEncoding, MaterializeFuncReturnOp>(
      typeConverter, context);
};

} // namespace mlir::iree_compiler
