// 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 pointwise 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/MapStableHLOToScalarOp.h"
#include "compiler/plugins/input/StableHLO/Conversion/Rewriters.h"
#include "compiler/plugins/input/StableHLO/Conversion/TypeConversion.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "stablehlo/dialect/StablehloOps.h"

namespace mlir::iree_compiler::stablehlo {
namespace {
int64_t getRank(Value v) { return cast<ShapedType>(v.getType()).getRank(); }

int64_t getMaxRank(ValueRange operands) {
  int64_t maxRank = 0;
  for (Value operand : operands) {
    maxRank = std::max(maxRank, getRank(operand));
  }
  return maxRank;
}

bool isScalar(Value v) { return getRank(v) == 0; }

/// Inserts block arguments in places where scalar inputs have a nullptr.
SmallVector<Value> interleaveScalarAndBlockArgs(ValueRange scalarInputs,
                                                ValueRange blockArgs) {
  SmallVector<Value> result;
  auto argsIter = blockArgs.begin();
  for (Value scalarInput : scalarInputs) {
    if (scalarInput) {
      result.push_back(scalarInput);
    } else {
      result.push_back(*argsIter);
      ++argsIter;
    }
  }
  return result;
}

struct PointwiseConversionInfo {
  int64_t maxOperandRank = 0;
  ShapedType resultType;
};

/// Checks the preconditions for conversion of pointwise HLO ops to linalg.
/// Returns the max operand rank and the result type on success.
FailureOr<PointwiseConversionInfo>
checkOperandsAndResults(Operation *op, ValueRange operands,
                        const TypeConverter &typeConverter,
                        ConversionPatternRewriter &rewriter) {
  int64_t maxRank = getMaxRank(operands);

  // Apply only if all operands are scalar or have the same rank. Some ops,
  // like `stablehlo.select`, support implicit broadcasting of scalars.
  if (!llvm::all_of(operands, [&](Value v) {
        int64_t r = getRank(v);
        return r == 0 || r == maxRank;
      })) {
    return rewriter.notifyMatchFailure(
        op, "Operands must be of same rank or scalar.");
  }

  // Find result type, if on tensors.
  auto resultTy = dyn_cast_or_null<ShapedType>(
      typeConverter.convertType(op->getResultTypes().front()));

  // Check result type compatibility.
  if (!resultTy || !resultTy.hasRank() || resultTy.getRank() != maxRank ||
      !(resultTy.getElementType().isSignlessIntOrFloat() ||
        isa<ComplexType>(resultTy.getElementType()))) {
    return rewriter.notifyMatchFailure(
        op, "mismatched operand/result types or iterator count");
  }

  // All-scalar pointwise ops inside of linalg ops are processes by
  // ScalarHloToArithmeticPattern.
  if (maxRank == 0 && isInBodyOfLinalgOps(op))
    return failure();

  return PointwiseConversionInfo{maxRank, resultTy};
}

/// Converts a HLO operation to a linalg.map op that contains the corresponding
/// scalar operations.
template <typename OpTy>
struct PointwiseToLinalgMapConverter final : OpConversionPattern<OpTy> {
  using OpConversionPattern<OpTy>::OpConversionPattern;
  using OpAdaptor = typename OpTy::Adaptor;

  LogicalResult
  matchAndRewrite(OpTy op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto conversionInfo = checkOperandsAndResults(
        op, adaptor.getOperands(), *this->typeConverter, rewriter);
    if (failed(conversionInfo)) {
      return failure();
    }

    int64_t maxRank = conversionInfo->maxOperandRank;
    ShapedType resultTy = conversionInfo->resultType;
    Location loc = op.getLoc();

    // Find input/output values and types.
    Value emptyTensor =
        getEmptyTensorFor(rewriter, loc, resultTy, op, adaptor.getOperands());

    // Mapped inputs are cast to the same shape as the init tensor.
    // Values from scalar inputs are extracted and used directly in the block.
    SmallVector<Value> mappedInputs;
    SmallVector<Value> scalarInputs;
    for (Value input : adaptor.getOperands()) {
      if (getRank(input) == maxRank) {
        mappedInputs.push_back(coerceTensorShape(
            rewriter, loc, cast<TypedValue<ShapedType>>(input),
            cast<ShapedType>(emptyTensor.getType())));
        scalarInputs.push_back(nullptr);
      } else {
        scalarInputs.push_back(rewriter.create<tensor::ExtractOp>(loc, input));
      }
    }

    auto mapOp = rewriter.create<linalg::MapOp>(
        loc, mappedInputs, emptyTensor,
        [&](OpBuilder &b, Location loc, ValueRange args) {
          Value innerResult = mlir::stablehlo::StableHloOpToStdScalarOp::mapOp(
              op, getElementTypeOrSelf(emptyTensor),
              interleaveScalarAndBlockArgs(scalarInputs, args), &b);

          b.create<linalg::YieldOp>(loc, innerResult);
        },
        linalg::getPrunedAttributeList(op));

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

/// Converts a HLO operation to a linalg.generic op that contains the
/// corresponding scalar operations.
template <typename OpTy>
struct PointwiseToLinalgConverter final : OpConversionPattern<OpTy> {
  using OpConversionPattern<OpTy>::OpConversionPattern;
  using OpAdaptor = typename OpTy::Adaptor;

  LogicalResult
  matchAndRewrite(OpTy op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto conversionInfo = checkOperandsAndResults(
        op, adaptor.getOperands(), *this->typeConverter, rewriter);
    if (failed(conversionInfo)) {
      return failure();
    }

    int64_t maxRank = conversionInfo->maxOperandRank;
    ShapedType resultTy = conversionInfo->resultType;
    Location loc = op.getLoc();

    // Find input/output values and types.
    ValueRange inputs = adaptor.getOperands();
    Value output =
        getEmptyTensorFor(rewriter, loc, resultTy, op, adaptor.getOperands());

    // Create indexing maps.
    AffineMap scalarMap = AffineMap::get(maxRank, 0, rewriter.getContext());
    AffineMap idMap = rewriter.getMultiDimIdentityMap(maxRank);
    SmallVector<AffineMap> maps;
    for (Value v : inputs)
      maps.push_back(isScalar(v) ? scalarMap : idMap);
    maps.push_back(idMap);

    // Build `linalg.generic` op.
    bool failed = false;
    auto linalgOp = rewriter.create<linalg::GenericOp>(
        loc, resultTy ? resultTy : TypeRange{}, inputs, output, maps,
        getNParallelLoopsAttrs(maxRank),
        [&](OpBuilder &nestedBuilder, Location /*nested_loc*/,
            ValueRange args) {
          Type innerResultTy = getElementTypeOrSelf(output);
          auto argvec = llvm::to_vector<2>(args.take_front(inputs.size()));
          Value semiring = preSparsify(op, argvec, innerResultTy, &rewriter);
          Value innerResult = mlir::stablehlo::StableHloOpToStdScalarOp::mapOp(
              op, innerResultTy, argvec, &rewriter);
          if (!innerResult) {
            failed = true;
          } else {
            innerResult = postSparsify(op, semiring, innerResult, &rewriter);
            nestedBuilder.create<linalg::YieldOp>(loc, innerResult);
          }
        },
        linalg::getPrunedAttributeList(op));
    if (failed)
      return failure();

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

namespace detail {
void populatePointwiseStableHloToLinalgConversionPatterns(
    MLIRContext *context, TypeConverter &typeConverter,
    RewritePatternSet *patterns, bool enablePrimitiveOps) {
  if (enablePrimitiveOps) {
    patterns->add<
        PointwiseToLinalgMapConverter<mlir::stablehlo::AbsOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::AddOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::AndOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::Atan2Op>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::BitcastConvertOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::CbrtOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::CeilOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ClampOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ClzOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::CompareOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ComplexOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ConvertOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::CosineOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::DivOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ExpOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::Expm1Op>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::FloorOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ImagOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::IsFiniteOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::Log1pOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::LogOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::LogisticOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::MaxOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::MinOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::MulOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::NegOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::NotOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::OrOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::PopulationCountOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::PowOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::RealOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ReducePrecisionOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::RemOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::RoundNearestEvenOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::RoundOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::RsqrtOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::SelectOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ShiftLeftOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ShiftRightArithmeticOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::ShiftRightLogicalOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::SignOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::SineOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::SqrtOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::SubtractOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::TanhOp>,
        PointwiseToLinalgMapConverter<mlir::stablehlo::XorOp>>(typeConverter,
                                                               context);
    return;
  }

  patterns
      ->add<PointwiseToLinalgConverter<mlir::stablehlo::AbsOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::AddOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::AndOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::Atan2Op>,
            PointwiseToLinalgConverter<mlir::stablehlo::BitcastConvertOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::CbrtOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::CeilOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ClampOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ClzOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::CompareOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ComplexOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ConvertOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::CosineOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::DivOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ExpOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::Expm1Op>,
            PointwiseToLinalgConverter<mlir::stablehlo::FloorOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ImagOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::IsFiniteOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::Log1pOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::LogOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::LogisticOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::MaxOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::MinOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::MulOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::NegOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::NotOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::OrOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::PopulationCountOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::PowOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::RealOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ReducePrecisionOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::RemOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::RoundNearestEvenOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::RoundOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::RsqrtOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::SelectOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ShiftLeftOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ShiftRightArithmeticOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::ShiftRightLogicalOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::SignOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::SineOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::SqrtOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::SubtractOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::TanhOp>,
            PointwiseToLinalgConverter<mlir::stablehlo::XorOp>>(typeConverter,
                                                                context);
}
} // namespace detail
} // namespace mlir::iree_compiler::stablehlo