// 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/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Pass/Pass.h"

namespace mlir::iree_compiler::IREE::LinalgExt {

#define GEN_PASS_DEF_CONVERTATTENTIONTOONLINEATTENTIONPASS
#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h.inc"

namespace {

struct ConvertAttentionToOnlineAttentionPass final
    : impl::ConvertAttentionToOnlineAttentionPassBase<
          ConvertAttentionToOnlineAttentionPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<IREE::LinalgExt::IREELinalgExtDialect,
                    linalg::LinalgDialect, tensor::TensorDialect>();
  }
  void runOnOperation() override;
};

} // namespace

void convertToOnlineAttention(IREE::LinalgExt::AttentionOp attnOp,
                              SmallVectorImpl<Operation *> &ops,
                              RewriterBase &rewriter) {
  rewriter.setInsertionPoint(attnOp);

  Location loc = attnOp.getLoc();
  MLIRContext *ctx = attnOp.getContext();

  FailureOr<AttentionOpDetail> maybeOpInfo =
      AttentionOpDetail::get(attnOp.getQueryMap(), attnOp.getKeyMap(),
                             attnOp.getValueMap(), attnOp.getOutputMap());
  assert(succeeded(maybeOpInfo) && "Invalid attention indexing maps");
  AttentionOpDetail opInfo = maybeOpInfo.value();

  // Create standard maps for max and sum: (batch, m)
  int64_t rank = opInfo.getDomainRank();
  AffineMap maxMap = AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, ctx);
  for (auto dim :
       llvm::concat<const int64_t>(opInfo.getBatchDims(), opInfo.getMDims())) {
    maxMap = maxMap.insertResult(rewriter.getAffineDimExpr(dim),
                                 maxMap.getNumResults());
  }
  AffineMap sumMap = maxMap;

  AffineMap accMap = attnOp.getOutputMap();

  SmallVector<Range> domain = attnOp.getIterationDomain(rewriter);

  // Create fill for acc, max and sum.
  // TODO: Acc should not need a fill. The attention op should get a filled
  // input instead of an empty input.

  SmallVector<OpFoldResult> sizes =
      llvm::map_to_vector(domain, [](Range x) { return x.size; });
  SmallVector<OpFoldResult> accSize =
      applyPermutationMap<OpFoldResult>(accMap, sizes);
  SmallVector<OpFoldResult> rowRedSize =
      applyPermutationMap<OpFoldResult>(maxMap, sizes);

  Type f32Type = rewriter.getF32Type();
  Value acc = tensor::EmptyOp::create(rewriter, loc, accSize, f32Type);
  Value rowRedEmpty =
      tensor::EmptyOp::create(rewriter, loc, rowRedSize, f32Type);

  Value accInit =
      arith::getIdentityValue(arith::AtomicRMWKind::addf, f32Type, rewriter,
                              loc, /*useOnlyFiniteValue=*/true);
  Value maxInit =
      arith::getIdentityValue(arith::AtomicRMWKind::maximumf, f32Type, rewriter,
                              loc, /*useOnlyFiniteValue=*/true);
  Value sumInit = arith::getIdentityValue(arith::AtomicRMWKind::addf, f32Type,
                                          rewriter, loc);

  Value accFill =
      linalg::FillOp::create(rewriter, loc, ValueRange{accInit}, acc)
          .getResult(0);
  Value maxFill =
      linalg::FillOp::create(rewriter, loc, ValueRange{maxInit}, rowRedEmpty)
          .getResult(0);
  Value sumFill =
      linalg::FillOp::create(rewriter, loc, ValueRange{sumInit}, rowRedEmpty)
          .getResult(0);

  // Create online attention op.
  SmallVector<AffineMap> indexingMaps = attnOp.getIndexingMapsArray();
  indexingMaps.push_back(maxMap);
  indexingMaps.push_back(sumMap);

  std::optional<Value> mask =
      attnOp.getMask() ? std::optional<Value>(attnOp.getMask()) : std::nullopt;

  // Use the custom builder which correctly handles the Optional mask.
  // ODS operand order: Q, K, V, scale, [mask], output, max, sum.
  // Custom builder takes: Q, K, V, scale, output, max, sum, maps, mask.
  OnlineAttentionOp onlineAttn = OnlineAttentionOp::create(
      rewriter, loc,
      TypeRange{accFill.getType(), maxFill.getType(), sumFill.getType()},
      attnOp.getQuery(), attnOp.getKey(), attnOp.getValue(), attnOp.getScale(),
      accFill, maxFill, sumFill, rewriter.getAffineMapArrayAttr(indexingMaps),
      mask);
  if (attnOp.getDecompositionConfigAttr()) {
    onlineAttn.setDecompositionConfigAttr(attnOp.getDecompositionConfigAttr());
  }

  rewriter.cloneRegionBefore(attnOp.getRegion(), onlineAttn.getRegion(),
                             onlineAttn.getRegion().begin());
  onlineAttn->setDiscardableAttrs(attnOp->getDiscardableAttrDictionary());
  ops.push_back(onlineAttn);

  Value x = onlineAttn.getResult(0);
  Value sum = onlineAttn.getResult(2);

  // Merge the outputs of online attention:
  //  x = (1 / sum) * x

  // Compress the indexing maps.
  SmallVector<AffineMap> compressedMaps =
      compressUnusedDims(SmallVector<AffineMap>{sumMap, accMap, accMap});

  SmallVector<utils::IteratorType> iteratorTypes(compressedMaps[0].getNumDims(),
                                                 utils::IteratorType::parallel);

  auto genericOp = linalg::GenericOp::create(
      rewriter, loc, attnOp.getOutput().getType(), ValueRange{sum, x},
      attnOp.getOutput(), compressedMaps, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value one = arith::ConstantOp::create(
            b, loc, b.getFloatAttr(args[0].getType(), 1.0));
        Value reciprocal = arith::DivFOp::create(b, loc, one, args[0]);
        // Both sum and x are in fp32, as created earlier, so we only need
        // to cast after the mul.
        Value result = arith::MulFOp::create(b, loc, reciprocal, args[1]);
        // Cast result to the required type by attention output.
        result = convertScalarToDtype(b, loc, result, args[2].getType(),
                                      /*isUnsignedCast=*/false);
        linalg::YieldOp::create(b, loc, result);
      });
  ops.push_back(genericOp);

  rewriter.replaceOp(attnOp, genericOp);
}

void ConvertAttentionToOnlineAttentionPass::runOnOperation() {
  MLIRContext *context = &getContext();
  IRRewriter rewriter(context);
  getOperation()->walk([&](AttentionOp attnOp) {
    SmallVector<Operation *> ops;
    convertToOnlineAttention(attnOp, ops, rewriter);
  });
}

} // namespace mlir::iree_compiler::IREE::LinalgExt
