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

// Implements logic for lowering StableHLO dot product ops to Linalg dialect.
// These patterns are separated out to their own file to save on the compilation
// times, given that we instantiate a large number of class templates here.

#include "compiler/plugins/input/StableHLO/Conversion/LegalizeToLinalgUtils.h"
#include "compiler/plugins/input/StableHLO/Conversion/Rewriters.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "stablehlo/dialect/StablehloOps.h"

namespace mlir::iree_compiler::stablehlo {
namespace {
enum class DotOperationType {
  kVectorDot = 0,
  kMatrixVector,
  kVectorMatrix,
  kMatrixMatrix,
  kUnsupported
};

DotOperationType getDotOperationType(mlir::stablehlo::DotOp dotOp) {
  ArrayRef<int64_t> lhsShape =
      cast<ShapedType>(dotOp.getLhs().getType()).getShape();
  ArrayRef<int64_t> rhsShape =
      cast<ShapedType>(dotOp.getRhs().getType()).getShape();
  auto shapeMatches = [](int64_t a, int64_t b) {
    return ShapedType::isDynamic(a) || ShapedType::isDynamic(b) || a == b;
  };
  if (lhsShape.size() == 1 && rhsShape.size() == 1 &&
      shapeMatches(lhsShape[0], rhsShape[0])) {
    return DotOperationType::kVectorDot;
  }
  if (lhsShape.size() == 2 && rhsShape.size() == 1 &&
      shapeMatches(lhsShape[1], rhsShape[0])) {
    return DotOperationType::kMatrixVector;
  }
  if (lhsShape.size() == 1 && rhsShape.size() == 2 &&
      shapeMatches(lhsShape[0], rhsShape[0])) {
    return DotOperationType::kVectorMatrix;
  }
  if (lhsShape.size() == 2 && rhsShape.size() == 2 &&
      shapeMatches(lhsShape[1], rhsShape[0])) {
    return DotOperationType::kMatrixMatrix;
  }
  return DotOperationType::kUnsupported;
}

SmallVector<Value, 2> getDotOpEmptyTensorDynSizes(OpBuilder &b, Location loc,
                                                  Value lhs, Value rhs,
                                                  DotOperationType type) {
  SmallVector<Value, 2> dynShape;
  switch (type) {
  case DotOperationType::kMatrixMatrix: {
    if (llvm::cast<ShapedType>(lhs.getType()).isDynamicDim(0))
      dynShape.push_back(b.create<tensor::DimOp>(loc, lhs, 0));
    if (llvm::cast<ShapedType>(rhs.getType()).isDynamicDim(1))
      dynShape.push_back(b.create<tensor::DimOp>(loc, rhs, 1));
    break;
  }
  case DotOperationType::kMatrixVector: {
    if (llvm::cast<ShapedType>(lhs.getType()).isDynamicDim(0))
      dynShape.push_back(b.create<tensor::DimOp>(loc, lhs, 0));
    break;
  }
  case DotOperationType::kVectorMatrix: {
    if (llvm::cast<ShapedType>(rhs.getType()).isDynamicDim(1))
      dynShape.push_back(b.create<tensor::DimOp>(loc, rhs, 1));
    break;
  }
  case DotOperationType::kVectorDot:
  case DotOperationType::kUnsupported:
    break;
  }
  return dynShape;
}

template <DotOperationType op_type, typename LinalgOp>
struct DotOpConversion final : OpConversionPattern<mlir::stablehlo::DotOp> {
  using OpConversionPattern<mlir::stablehlo::DotOp>::OpConversionPattern;
  using OpAdaptor = mlir::stablehlo::DotOp::Adaptor;

  LogicalResult
  matchAndRewrite(mlir::stablehlo::DotOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    if (failed(verifyHloOpBufferOrTensorSemantics(op))) {
      return failure();
    }
    if (getDotOperationType(op) != op_type)
      return failure();

    Location loc = op.getLoc();
    // Convert unsigned to signed. This works because signed and unsigned
    // integer matmul is the same operation in two's complement.
    auto outputType =
        cast<ShapedType>(getTypeConverter()->convertType(op.getType()));
    SmallVector<Value, 2> dynShape = getDotOpEmptyTensorDynSizes(
        rewriter, loc, adaptor.getLhs(), adaptor.getRhs(), op_type);
    Value emptyTensor =
        !sparse_tensor::getSparseTensorEncoding(outputType)
            ? getEmptyTensor(rewriter, loc, outputType, dynShape)
            : getEmptySparseTensor(rewriter, loc, outputType, dynShape);
    Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor);
    rewriter.replaceOpWithNewOp<LinalgOp>(
        op, TypeRange{outputType},
        ValueRange{adaptor.getLhs(), adaptor.getRhs()}, ValueRange{zeroTensor},
        linalg::getPrunedAttributeList(op));
    return success();
  }
};

struct DotGeneralBatchMatMulOpConversion final
    : OpConversionPattern<mlir::stablehlo::DotGeneralOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(mlir::stablehlo::DotGeneralOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    if (failed(verifyHloOpBufferOrTensorSemantics(op))) {
      return failure();
    }
    if (llvm::cast<RankedTensorType>(op.getType()).getRank() != 3) {
      return rewriter.notifyMatchFailure(op, "expected a batch matmul");
    }

    mlir::stablehlo::DotDimensionNumbersAttr dimNumbers =
        op.getDotDimensionNumbers();
    ArrayRef<int64_t> lhsBatchingDims = dimNumbers.getLhsBatchingDimensions();
    ArrayRef<int64_t> rhsBatchingDims = dimNumbers.getRhsBatchingDimensions();
    ArrayRef<int64_t> lhsContractingDims =
        dimNumbers.getLhsContractingDimensions();
    ArrayRef<int64_t> rhsContractingDims =
        dimNumbers.getRhsContractingDimensions();
    if (lhsBatchingDims.size() != 1 || lhsBatchingDims[0] != 0) {
      return rewriter.notifyMatchFailure(
          op, "expected lhs batching dimensions exactly {0}");
    }
    if (rhsBatchingDims.size() != 1 || rhsBatchingDims[0] != 0) {
      return rewriter.notifyMatchFailure(
          op, "expected rhs batching dimensions exactly {0}");
    }
    if (lhsContractingDims.size() != 1 || lhsContractingDims[0] != 2) {
      return rewriter.notifyMatchFailure(
          op, "expected lhs contracting dimensions exactly {2}");
    }
    if (rhsContractingDims.size() != 1 || rhsContractingDims[0] != 1) {
      return rewriter.notifyMatchFailure(
          op, "expected rhs contracting dimensions exactly {1}");
    }

    Location loc = op.getLoc();
    // Convert unsigned to signed. This works because signed and unsigned
    // integer matmul is the same operation in two's complement.
    auto outputType =
        cast<ShapedType>(typeConverter->convertType(op.getType()));
    Value emptyTensor =
        getEmptyTensorFor(rewriter, loc, outputType, op, adaptor.getOperands());
    Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor);
    Operation *linalgOp = rewriter.create<linalg::BatchMatmulOp>(
        loc, /*resultTensorTypes=*/TypeRange{outputType},
        /*inputs=*/ValueRange{adaptor.getLhs(), adaptor.getRhs()},
        /*outputBuffers=*/ValueRange{zeroTensor},
        linalg::getPrunedAttributeList(op));

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

struct DotGeneralOpConversion final
    : OpConversionPattern<mlir::stablehlo::DotGeneralOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(mlir::stablehlo::DotGeneralOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    if (failed(verifyHloOpBufferOrTensorSemantics(op))) {
      return failure();
    }

    // Get various dimension iterator information
    mlir::stablehlo::DotDimensionNumbersAttr dimNumbers =
        op.getDotDimensionNumbers();
    ArrayRef<int64_t> lhsBatchingDims = dimNumbers.getLhsBatchingDimensions();
    ArrayRef<int64_t> rhsBatchingDims = dimNumbers.getRhsBatchingDimensions();
    ArrayRef<int64_t> lhsContractingDims =
        dimNumbers.getLhsContractingDimensions();
    ArrayRef<int64_t> rhsContractingDims =
        dimNumbers.getRhsContractingDimensions();

    // Get shape information and initialize output
    assert(lhsContractingDims.size() == rhsContractingDims.size() &&
           "number of contracting dims must be equal");
    size_t numContracting = lhsContractingDims.size();
    // Convert unsigned to signed. This works because signed and unsigned
    // integer matmul is the same operation in two's complement.
    auto outputType =
        cast<ShapedType>(typeConverter->convertType(op.getType()));
    size_t targetRank = outputType.getRank();
    size_t totalLoopCount = numContracting + targetRank;

    int64_t lhsRank =
        llvm::cast<ShapedType>(adaptor.getLhs().getType()).getRank();
    size_t lhsExtraDims =
        lhsRank - lhsBatchingDims.size() - lhsContractingDims.size();
    int64_t rhsRank =
        llvm::cast<ShapedType>(adaptor.getRhs().getType()).getRank();

    Location loc = op.getLoc();
    Value emptyTensor =
        getEmptyTensorFor(rewriter, loc, outputType, op, adaptor.getOperands());
    Value zeroTensor = fillTensorWithZeros(rewriter, loc, emptyTensor);
    SmallVector<AffineMap, 3> indexingMaps;

    auto getMap = [&](int64_t rank, ArrayRef<int64_t> batchingDims,
                      ArrayRef<int64_t> contractingDims, size_t extraDims) {
      llvm::SmallVector<AffineExpr> indices(rank);
      for (const auto &i : llvm::enumerate(batchingDims)) {
        indices[i.value()] = rewriter.getAffineDimExpr(i.index());
      }
      for (const auto &i : llvm::enumerate(contractingDims)) {
        indices[i.value()] = rewriter.getAffineDimExpr(i.index() + targetRank);
      }
      for (int i = 0; i < rank; ++i) {
        if (!indices[i]) {
          indices[i] = rewriter.getAffineDimExpr(extraDims++);
        }
      }
      indexingMaps.push_back(AffineMap::get(/*dimCount=*/totalLoopCount,
                                            /*symbolCount=*/0, indices,
                                            op->getContext()));
    };
    getMap(lhsRank, lhsBatchingDims, lhsContractingDims,
           lhsBatchingDims.size());
    getMap(rhsRank, rhsBatchingDims, rhsContractingDims,
           rhsBatchingDims.size() + lhsExtraDims);

    {
      SmallVector<AffineExpr> dimExprs;
      dimExprs.reserve(targetRank);
      for (unsigned i = 0; i < targetRank; ++i)
        dimExprs.push_back(rewriter.getAffineDimExpr(i));
      indexingMaps.push_back(AffineMap::get(/*dimCount=*/totalLoopCount,
                                            /*symbolCount=*/0, dimExprs,
                                            op.getContext()));
    }

    Operation *linalgOp = rewriter.create<linalg::GenericOp>(
        loc, /*resultTensorTypes=*/TypeRange{outputType},
        /*inputs=*/ValueRange{adaptor.getLhs(), adaptor.getRhs()},
        /*outputBuffers=*/ValueRange{zeroTensor}, indexingMaps,
        getParallelAndReductionIterators(
            /*nLoops=*/totalLoopCount,
            /*nReduction=*/numContracting),
        [](OpBuilder &b, Location loc, ValueRange) {
          ImplicitLocOpBuilder builder(loc, b);
          linalg::MatmulOp::regionBuilder(builder, *b.getInsertionBlock(), {});
        },
        linalg::getPrunedAttributeList(op));

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

} // namespace

namespace detail {
void populateStableHloDotProdToLinalgConversionPatterns(
    MLIRContext *context, TypeConverter &typeConverter,
    RewritePatternSet *patterns) {
  // Ensure specialized patterns are higher priority than their generic
  // versions.
  patterns
      ->add<DotOpConversion<DotOperationType::kMatrixMatrix, linalg::MatmulOp>,
            DotOpConversion<DotOperationType::kMatrixVector, linalg::MatvecOp>,
            DotOpConversion<DotOperationType::kVectorMatrix, linalg::VecmatOp>,
            DotOpConversion<DotOperationType::kVectorDot, linalg::DotOp>,
            DotGeneralBatchMatMulOpConversion>(typeConverter, context,
                                               PatternBenefit(2));
  patterns->add<DotGeneralOpConversion>(typeConverter, context,
                                        PatternBenefit(1));
}
} // namespace detail
} // namespace mlir::iree_compiler::stablehlo