//===- EmptyTensorElimination.cpp - tensor.empty op elimination -----------===//
//
// Part of the LLVM Project, 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 "mlir/Dialect/Linalg/Transforms/Transforms.h"

#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

using namespace mlir;
using namespace mlir::bufferization;
using namespace mlir::linalg;

/// Get an output operand that matches the given input operand and can be used
/// to eliminate a tensor.empty op.
static OpOperand *getUnusedOutOperand(LinalgOp op, OpOperand *in) {
  for (OpOperand &operand : op.getDpsInitsMutable()) {
    // Operand must be unused.
    if (op.payloadUsesValueFromOperand(&operand))
      continue;
    // Types must match.
    if (operand.get().getType() != in->get().getType())
      continue;
    // Indexing maps must match.
    if (op.getMatchingIndexingMap(&operand) != op.getMatchingIndexingMap(in))
      continue;
    return &operand;
  }
  return nullptr;
}

LogicalResult linalg::linalgOpAnchoredEmptyTensorEliminationStep(
    RewriterBase &rewriter, Operation *op, OneShotAnalysisState &state) {
  OpBuilder::InsertionGuard g(rewriter);
  DominanceInfo domInfo;

  op->walk([&](LinalgOp op) {
    // Only ops with all "parallel" iterator types are supported.
    if (op.getNumParallelLoops() != op.getNumLoops())
      return WalkResult::skip();

    for (OpOperand *in : op.getDpsInputOperands()) {
      // Skip non-tensor operands.
      if (!isa<RankedTensorType>(in->get().getType()))
        continue;

      // Find tensor.empty ops on the reverse SSA use-def chain. Only follow
      // equivalent tensors. I.e., stop when there are ops such as extract_slice
      // on the path.
      TraversalConfig config;
      config.followEquivalentOnly = true;
      config.alwaysIncludeLeaves = false;
      SetVector<Value> emptyTensors = state.findValueInReverseUseDefChain(
          in->get(), /*condition=*/
          [&](Value val) {
            return val.getDefiningOp<tensor::EmptyOp>() &&
                   val.getType() == in->get().getType();
          },
          config);
      if (emptyTensors.empty())
        continue;

      // Find matching out operand.
      OpOperand *out = getUnusedOutOperand(op, in);
      if (!out)
        continue;

      // Check if this transform would violate dominance.
      if (!llvm::all_of(emptyTensors, [&](Value v) {
            return domInfo.properlyDominates(out->get(), v.getDefiningOp());
          }))
        continue;

      // Replace all uses of the tensor.empty, but do not delete it yet. It will
      // fold away later (to not invalidate DominanceInfo).
      for (Value v : emptyTensors) {
        assert(v.getDefiningOp<tensor::EmptyOp>() && "expected tensor.empty");
        rewriter.replaceAllUsesWith(v, out->get());
      }

      // Turn the "in" into an "out".
      rewriter.modifyOpInPlace(op, [&]() {
        out->set(in->get());
        // The original "in" could be removed entirely here (because it will no
        // longer have any uses in the payload), but we delegate this to
        // existing cleanup patterns that remove unused operands.
        in->set(emptyTensors.front());
        BlockArgument outArg = op.getMatchingBlockArgument(out);
        assert(outArg.getUses().empty() && "expected that out has no uses");
        BlockArgument inArg = op.getMatchingBlockArgument(in);
        rewriter.replaceAllUsesWith(inArg, outArg);
        assert(!op.payloadUsesValueFromOperand(in) &&
               "expected that the in operand is now unused");
      });

      state.resetCache();
    }

    return WalkResult::advance();
  });
  return success();
}