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

#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
#include "iree/compiler/Preprocessing/Common/Passes.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Arith/IR/Arith.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/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-preprocessing-convert-conv-filter-to-channels-last"

namespace mlir::iree_compiler::Preprocessing {

#define GEN_PASS_DEF_CONVERTCONVFILTERTOCHANNELSLASTPASS
#include "iree/compiler/Preprocessing/Common/Passes.h.inc"

static AffineMap applyPermutationToResults(AffineMap map,
                                           ArrayRef<int64_t> perm) {
  unsigned numDims = map.getNumDims();
  ArrayRef<AffineExpr> mapResults = map.getResults();
  SmallVector<AffineExpr> exprs;
  for (int i = 0, e = perm.size(); i < e; ++i) {
    exprs.push_back(mapResults[perm[i]]);
  }
  return AffineMap::get(numDims, map.getNumSymbols(), exprs, map.getContext());
}

static Value createTransposeOp(RewriterBase &rewriter, Location loc,
                               Value tensor, ArrayRef<int64_t> perm) {
  SmallVector<OpFoldResult> dimSizes =
      tensor::getMixedSizes(rewriter, loc, tensor);
  applyPermutationToVector(dimSizes, perm);

  auto tensorType = cast<RankedTensorType>(tensor.getType());
  auto emptyTensor = tensor::EmptyOp::create(rewriter, loc, dimSizes,
                                             tensorType.getElementType());
  return linalg::TransposeOp::create(rewriter, loc, tensor, emptyTensor, perm)
      .getResult()[0];
}

static LogicalResult
convertConvFilterToTargetLayout(linalg::Conv2DNhwcHwcfOp convOp,
                                RewriterBase &rewriter,
                                SmallVector<int64_t> &perm) {
  Location loc = convOp.getLoc();

  Value input = convOp.getInputs()[0];
  Value filter = convOp.getInputs()[1];
  Value output = convOp.getOutputs()[0];

  AffineMap inputMap = convOp.getIndexingMapsArray()[0];
  AffineMap filterMap = convOp.getIndexingMapsArray()[1];
  AffineMap outputMap = convOp.getIndexingMapsArray()[2];

  AffineMap transposedFilterMap = applyPermutationToResults(filterMap, perm);
  Value transposedFilter = createTransposeOp(rewriter, loc, filter, perm);

  SmallVector<utils::IteratorType> iterators = convOp.getIteratorTypesArray();

  auto genericOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(), ValueRange{input, transposedFilter},
      output, ArrayRef<AffineMap>{inputMap, transposedFilterMap, outputMap},
      iterators);

  // Reuse the same payload as the original convolution op.
  rewriter.inlineRegionBefore(convOp->getRegion(0), genericOp.getRegion(),
                              genericOp.getRegion().begin());

  rewriter.replaceOp(convOp, genericOp->getResults());
  return success();
}

namespace {
struct ConvertHwcfToHwfc : OpRewritePattern<linalg::Conv2DNhwcHwcfOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::Conv2DNhwcHwcfOp convOp,
                                PatternRewriter &rewriter) const override {
    SmallVector<int64_t> perm = {0, 1, 3, 2};
    return convertConvFilterToTargetLayout(convOp, rewriter, perm);
  }
};

struct ConvertHwcfToFhwc : OpRewritePattern<linalg::Conv2DNhwcHwcfOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::Conv2DNhwcHwcfOp convOp,
                                PatternRewriter &rewriter) const override {
    SmallVector<int64_t> perm = {3, 0, 1, 2};
    return convertConvFilterToTargetLayout(convOp, rewriter, perm);
  }
};

/// Transpose the generic form filter layout of `CHWF`, `CFHW`, or `HWCF` to
/// `FHWC`. Also matches convolution-like generics with extra operands (e.g.
/// quantization zero-points).
struct ConvertGenericFilterToFhwc : OpRewritePattern<linalg::GenericOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(linalg::GenericOp op,
                                PatternRewriter &rewriter) const override {
    auto linalgOp = cast<linalg::LinalgOp>(op.getOperation());
    linalg::ConvolutionDimensions convolutionDims;
    StringRef errString = getMatchConvolutionMessage(
        linalg::detail::isConvolutionInterfaceImpl(linalgOp, &convolutionDims));
    if (!errString.empty()) {
      return failure();
    }

    // Require non-empty filter, input and output channel dimensions.
    if (convolutionDims.outputChannel.empty() ||
        convolutionDims.inputChannel.empty() ||
        convolutionDims.filterLoop.empty()) {
      return failure();
    }

    SmallVector<Value> inputs = linalgOp.getDpsInputs();
    Value inputVal = inputs[0];
    Value filterVal = inputs[1];
    Value outputVal = linalgOp.getDpsInits()[0];

    auto indexingMaps = linalgOp.getIndexingMapsArray();
    AffineMap inputMap = indexingMaps[0];
    AffineMap filterMap = indexingMaps[1];
    AffineMap outputMap = indexingMaps.back();

    ArrayRef<int64_t> inputShape =
        cast<ShapedType>(inputVal.getType()).getShape();
    ArrayRef<int64_t> filterShape =
        cast<ShapedType>(filterVal.getType()).getShape();
    ArrayRef<int64_t> outputShape =
        cast<ShapedType>(outputVal.getType()).getShape();

    // TODO(vivian): Once the matmul shape check below is dropped, the
    // dynamic-shape check can also be removed.
    if (ShapedType::isDynamicShape(inputShape) ||
        ShapedType::isDynamicShape(filterShape) ||
        ShapedType::isDynamicShape(outputShape)) {
      return failure();
    }

    auto getDimPositions = [&](ArrayRef<unsigned> dims, const AffineMap &map) {
      SmallVector<int64_t> positions;
      for (auto dim : dims) {
        for (auto [idx, e] : llvm::enumerate(map.getResults())) {
          if (e.isFunctionOfDim(dim)) {
            positions.push_back(idx);
          }
        }
      }
      return positions;
    };

    // Don't transpose when the input is in not batch-first layout (e.g., CHWN).
    SmallVector<int64_t> batchInputPos =
        getDimPositions(convolutionDims.batch, inputMap);
    if (!batchInputPos.empty() && batchInputPos.front() != 0) {
      return failure();
    }

    // Only transpose when the filter is CHWF, CFHW, or HWCF layout.
    SmallVector<int64_t> fFilterPos =
        getDimPositions(convolutionDims.outputChannel, filterMap);
    SmallVector<int64_t> cFilterPos =
        getDimPositions(convolutionDims.inputChannel, filterMap);
    SmallVector<int64_t> kFilterPos =
        getDimPositions(convolutionDims.filterLoop, filterMap);
    int64_t fPos = fFilterPos.back();
    int64_t cPos = cFilterPos.back();

    bool isChwf =
        (cPos < kFilterPos.front()) && (fPos == filterShape.size() - 1);
    bool isCfhw = (cPos < fFilterPos.front()) && (fPos < kFilterPos.front());
    bool isHwcf = (kFilterPos.back() < cPos) && (cPos < fPos);
    if (!isChwf && !isCfhw && !isHwcf) {
      return failure();
    }

    // Don't transpose if it is a matmul and the input shape is small.
    // TODO(vivian): Solve the fusion of transpose op and remove this check.
    SmallVector<int64_t> imagePos =
        getDimPositions(convolutionDims.outputImage, outputMap);
    SmallVector<int64_t> batchPos =
        getDimPositions(convolutionDims.batch, outputMap);
    SmallVector<int64_t> mPos = imagePos;
    mPos.append(batchPos.begin(), batchPos.end());

    auto getProduct = [](ArrayRef<int64_t> shape, ArrayRef<int64_t> pos) {
      return llvm::accumulate(pos, int64_t{1}, [&](int64_t a, int64_t idx) {
        return a * shape[idx];
      });
    };

    int64_t mSize = getProduct(outputShape, mPos);
    int64_t nSize = getProduct(filterShape, fFilterPos);
    int64_t kSize = getProduct(filterShape, cFilterPos);
    int64_t filterProd = getProduct(filterShape, kFilterPos);
    bool smallShape = mSize < 384 || nSize < 384 || kSize < 384;
    if (filterProd == 1 && smallShape) {
      return failure();
    }

    // Build permutation to convert to FHWC layout (or GFHWC for grouped conv).
    llvm::SmallDenseSet<int64_t> processedDims;
    processedDims.insert(cFilterPos.begin(), cFilterPos.end());
    processedDims.insert(fFilterPos.begin(), fFilterPos.end());
    processedDims.insert(kFilterPos.begin(), kFilterPos.end());
    // First add dimensions that are not C, F, H, W, (e.g., G for grouped conv).
    SmallVector<int64_t> perm;
    for (int64_t i = 0; i < filterShape.size(); ++i) {
      if (!processedDims.contains(i)) {
        perm.push_back(i);
      }
    }
    // Then append dimensions in the order of F, H, W, C.
    perm.append(fFilterPos.begin(), fFilterPos.end());
    perm.append(kFilterPos.begin(), kFilterPos.end());
    perm.append(cFilterPos.begin(), cFilterPos.end());

    Location loc = linalgOp.getLoc();

    AffineMap transposedFilterMap = applyPermutationToResults(filterMap, perm);
    Value transposedFilter = createTransposeOp(rewriter, loc, filterVal, perm);

    // Insert compute_barrier.start to avoid propagation of reshape ops and
    // undesirable fusion.
    auto barrierStartOp = IREE::TensorExt::ComputeBarrierStartOp::create(
        rewriter, loc, transposedFilter);

    SmallVector<utils::IteratorType> iterators =
        linalgOp.getIteratorTypesArray();

    SmallVector<Value> newInputs = inputs;
    newInputs[1] = barrierStartOp.getResult();
    SmallVector<AffineMap> newMaps = indexingMaps;
    newMaps[1] = transposedFilterMap;

    auto genericOp =
        linalg::GenericOp::create(rewriter, loc, outputVal.getType(), newInputs,
                                  outputVal, newMaps, iterators);

    // Reuse the same payload as the original convolution op.
    rewriter.inlineRegionBefore(linalgOp->getRegion(0), genericOp.getRegion(),
                                genericOp.getRegion().begin());

    // Reorder the indexing dimensions so that the input channel loops appears
    // after the filter loops.
    unsigned numParallelLoop = genericOp.getNumParallelLoops();
    SmallVector<unsigned> interchange =
        llvm::to_vector(llvm::seq<unsigned>(0, numParallelLoop));
    interchange.append(convolutionDims.filterLoop.begin(),
                       convolutionDims.filterLoop.end());
    interchange.append(convolutionDims.inputChannel.begin(),
                       convolutionDims.inputChannel.end());

    FailureOr<linalg::GenericOp> reorderOp =
        linalg::interchangeGenericOp(rewriter, genericOp, interchange);
    if (failed(reorderOp)) {
      return failure();
    }

    rewriter.replaceOp(linalgOp, reorderOp->getResults());
    return success();
  }
};

class ConvertConvFilterToChannelsLastPass
    : public iree_compiler::Preprocessing::impl::
          ConvertConvFilterToChannelsLastPassBase<
              ConvertConvFilterToChannelsLastPass> {
public:
  using iree_compiler::Preprocessing::impl::
      ConvertConvFilterToChannelsLastPassBase<
          ConvertConvFilterToChannelsLastPass>::
          ConvertConvFilterToChannelsLastPassBase;

  void runOnOperation() override {
    auto op = getOperation();
    MLIRContext *context = &getContext();

    RewritePatternSet patterns(context);
    if (filterLayout == "hwfc") {
      LDBG() << "Converting filter layout to hwfc.";
      patterns.add<ConvertHwcfToHwfc>(context);
    } else if (filterLayout == "fhwc") {
      LDBG() << "Converting filter layout to fhwc.";
      patterns.add<ConvertHwcfToFhwc, ConvertGenericFilterToFhwc>(context);
    } else {
      LDBG() << "convert-filter-to-channels-last pass didn't apply since an "
                "unsupported layout is given. Please use hwfc or fhwc as pass "
                "filter-layout option.";
      return signalPassFailure();
    }

    if (failed(applyPatternsGreedily(op, std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

} // namespace
} // namespace mlir::iree_compiler::Preprocessing
