#include "lib/Transforms/LayoutOptimization/LayoutOptimization.h"

#include <cassert>
#include <limits>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "lib/Analysis/LayoutFoldingAnalysis/LayoutFoldingAnalysis.h"
#include "lib/Dialect/Secret/IR/SecretAttributes.h"
#include "lib/Dialect/Secret/IR/SecretDialect.h"
#include "lib/Dialect/TensorExt/IR/TensorExtAttributes.h"
#include "lib/Dialect/TensorExt/IR/TensorExtDialect.h"
#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
#include "lib/Interface/HoistingInterfaces.h"
#include "lib/Kernel/AbstractValue.h"
#include "lib/Kernel/Kernel.h"
#include "lib/Kernel/KernelImplementation.h"
#include "lib/Kernel/KernelName.h"
#include "lib/Kernel/RotationCountVisitor.h"
#include "lib/Kernel/Utils.h"
#include "lib/Transforms/LayoutOptimization/Hoisting.h"
#include "lib/Transforms/LayoutOptimization/LayoutConversionCost.h"
#include "lib/Transforms/LayoutOptimization/Patterns.h"
#include "lib/Utils/AttributeUtils.h"
#include "llvm/include/llvm/ADT/STLExtras.h"               // from @llvm-project
#include "llvm/include/llvm/Support/Debug.h"               // from @llvm-project
#include "llvm/include/llvm/Support/raw_ostream.h"         // from @llvm-project
#include "mlir/include/mlir/Analysis/DataFlow/Utils.h"     // from @llvm-project
#include "mlir/include/mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"     // from @llvm-project
#include "mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/IR/AffineMap.h"          // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Diagnostics.h"        // from @llvm-project
#include "mlir/include/mlir/IR/Iterators.h"          // from @llvm-project
#include "mlir/include/mlir/IR/Operation.h"          // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"       // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"              // from @llvm-project
#include "mlir/include/mlir/IR/Visitors.h"           // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"          // from @llvm-project
#include "mlir/include/mlir/Support/WalkResult.h"    // from @llvm-project
#include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project

#define DEBUG_TYPE "layout-optimization"

namespace mlir {
namespace heir {

using ::mlir::heir::kernel::RotationCountVisitor;
using ::mlir::heir::kernel::SymbolicValue;
using ::mlir::heir::secret::KernelAttr;
using ::mlir::heir::tensor_ext::AssignLayoutOp;
using ::mlir::heir::tensor_ext::ConvertLayoutOp;
using ::mlir::heir::tensor_ext::LayoutAttr;

constexpr const static StringLiteral kKernelAttrName =
    ::mlir::heir::secret::SecretDialect::kKernelAttrName;

#define GEN_PASS_DEF_LAYOUTOPTIMIZATION
#include "lib/Transforms/LayoutOptimization/LayoutOptimization.h.inc"

namespace {

auto& kLayoutAttrName = tensor_ext::TensorExtDialect::kLayoutAttrName;

struct OperandChange {
  Attribute fromLayout;
  Attribute toLayout;
  Cost cost;
};

struct HoistOption {
  // Net increase in cost.
  Cost cost;
  HoistResult hoistResult;
};

}  // namespace

struct LayoutOptimization : impl::LayoutOptimizationBase<LayoutOptimization> {
  using LayoutOptimizationBase::LayoutOptimizationBase;

  enum OpHoistResult { UNHOISTABLE, SUCCESS, FAILURE };
  OpHoistResult hoistOp(Operation* op, IRRewriter& builder,
                        DataFlowSolver* solver);

  std::vector<HoistOption> computeHoistingOptions(
      Operation* op, ConvertLayoutOp convertLayoutOp, DataFlowSolver* solver);

  // Computes cost of layout conversion.
  Cost costOfLayoutConversion(Attribute fromLayout, Attribute toLayout);

  // Computes cost of changed operand.
  OperandChange costOfChangedOperand(OpOperand& operand, Operation* kernel,
                                     Attribute newLayout,
                                     DataFlowSolver* solver);

  // Computes cost of changed result.
  Cost costOfChangedResult(Operation* kernel, Attribute newLayout,
                           KernelAttr oldKernel, KernelName newKernel);

  void runOnOperation() override;
};

void LayoutOptimization::runOnOperation() {
  auto* ctx = &getContext();
  IRRewriter builder(ctx);

  RewritePatternSet preprocessingPatterns(ctx);
  preprocessingPatterns.add<FoldMatvecInputConversionIntoPlaintext>(ctx);
  if (failed(applyPatternsGreedily(getOperation(),
                                   std::move(preprocessingPatterns)))) {
    signalPassFailure();
    return;
  }

  DataFlowSolver solver;
  dataflow::loadBaselineAnalyses(solver);
  solver.load<LayoutIsFreeAnalysis>();
  auto solveResult = solver.initializeAndRun(getOperation());

  if (failed(solveResult)) {
    emitError(getOperation()->getLoc(), "Failed to run the analysis.\n");
    signalPassFailure();
    return;
  }

  SmallVector<Operation*> opsToErase;
  WalkResult result =
      getOperation()->walk<WalkOrder::PreOrder, ReverseIterator>(
          [&](Operation* op) {
            LLVM_DEBUG(llvm::dbgs()
                       << "Visiting op: " << op->getName() << " \n");
            if (auto hoistable =
                    dyn_cast<LayoutConversionHoistableOpInterface>(op)) {
              // TODO(#2317): figure out how to get OpInterface verifier to
              // run automatically instead of doing it manually here
              LLVM_DEBUG(llvm::dbgs() << "Visiting op: " << op->getName()
                                      << " with hoistable interface\n");
              [[maybe_unused]] KernelName kernelName = KernelName::Trivial;
              if (op->hasAttr(kKernelAttrName)) {
                kernelName =
                    op->getAttrOfType<KernelAttr>(kKernelAttrName).getName();
                LLVM_DEBUG(llvm::dbgs()
                           << "Op " << op->getName() << " has kernel attribute "
                           << kernelName << "\n");
              } else {
                LLVM_DEBUG(llvm::dbgs()
                           << "Op " << op->getName()
                           << " has no kernel attribute; using trivial\n");
              }
            }

            // Attempt to hoist layout conversions before this operation.
            OpHoistResult result = hoistOp(op, builder, &solver);
            if (result == FAILURE) {
              return WalkResult::interrupt();
            };
            if (result == UNHOISTABLE) {
              return WalkResult::advance();
            };

            // The above may results in sequences of convert_layout ops,
            // or convert_layout ops that occur directly after assign_layout
            // ops, and these can be eagerly folded.
            for (Value value : op->getOperands()) {
              if (auto convertLayoutOp =
                      value.getDefiningOp<ConvertLayoutOp>()) {
                (void)tryFoldLayoutConversionIntoPrevious(
                    builder, convertLayoutOp, opsToErase);
              }
            }

            LLVM_DEBUG(llvm::dbgs()
                       << "Dump after hoisting and eager folding: "
                       << op->getParentOfType<func::FuncOp>() << "\n\n");
            return WalkResult::advance();
          });

  if (result.wasInterrupted()) {
    signalPassFailure();
  }

  for (auto* op : opsToErase) {
    builder.eraseOp(op);
  }

  RewritePatternSet patterns(ctx);
  patterns.add<HoistArgLayouts, FoldLayoutConversions>(ctx);
  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
    signalPassFailure();
  }
};

LayoutOptimization::OpHoistResult LayoutOptimization::hoistOp(
    Operation* op, IRRewriter& builder, DataFlowSolver* solver) {
  // Folders will canonicalize assign_layout and convert_layout
  if (isa<ConvertLayoutOp, AssignLayoutOp>(op)) {
    return UNHOISTABLE;
  }

  if (op->getNumResults() != 1) {
    // Most ops we're interested in have a single result.
    return UNHOISTABLE;
  }

  // Check if any results are converted directly after.
  SmallVector<ConvertLayoutOp> resultLayoutConversions;
  for (auto* user : op->getResult(0).getUsers()) {
    if (auto convertLayoutOp = dyn_cast<ConvertLayoutOp>(user)) {
      resultLayoutConversions.push_back(convertLayoutOp);
    }
  }

  // No result has a layout conversion directly after.
  if (resultLayoutConversions.empty()) {
    LLVM_DEBUG(llvm::dbgs()
               << "Skipping op, no results followed by convert_layout\n");
    return UNHOISTABLE;
  }

  LLVM_DEBUG({
    llvm::dbgs() << "Considering hoisting " << resultLayoutConversions.size()
                 << " layout conversion through op: " << op->getName()
                 << "\n\n\t" << *op;
    for (auto convertLayoutOp : resultLayoutConversions) {
      llvm::dbgs() << "\n\t" << *convertLayoutOp << "\n";
    }
    llvm::dbgs() << "\n";
  });

  // Now compute the cost of hoisting each conversion layout.
  SmallVector<HoistOption> hoistingOptions;
  for (auto resultLayoutConversion : resultLayoutConversions) {
    auto options = computeHoistingOptions(op, resultLayoutConversion, solver);
    for (auto& option : options) {
      LLVM_DEBUG(llvm::dbgs()
                 << "\tHoisting layout " << option.hoistResult.newOutputLayout
                 << " via kernel " << option.hoistResult.newKernel
                 << " will cost " << option.cost << "\n");

      if (option.hoistResult.newInputLayouts.size() != op->getNumOperands()) {
        LLVM_DEBUG({
          std::string kernelName;
          llvm::raw_string_ostream os(kernelName);
          os << option.hoistResult.newKernel;
          op->emitOpError()
              << "Found invalid hoist result; op " << *op << " has "
              << op->getNumOperands() << " operands, but hoist result has "
              << option.hoistResult.newInputLayouts.size()
              << " operand layouts for kernel " << kernelName << "\n";
        });
        return UNHOISTABLE;
      }

      hoistingOptions.push_back(option);
    }
  }

  // Select the least costly layout conversion to hoist.
  auto* minHoistingCost = llvm::min_element(
      hoistingOptions, [](const HoistOption& a, const HoistOption& b) {
        return a.cost < b.cost;
      });

  if (minHoistingCost->cost > 0) {
    LLVM_DEBUG(llvm::dbgs()
               << "Skipping op, all hoisting results have positive net cost\n");
    return UNHOISTABLE;
  }
  HoistResult minHoistResult = minHoistingCost->hoistResult;

  LLVM_DEBUG(llvm::dbgs()
             << "Minimum cost layout conversion converts output layout from "
             << op->getAttr(kLayoutAttrName) << " to layout "
             << minHoistingCost->hoistResult.newOutputLayout << "\n");

  builder.setInsertionPoint(op);
  for (auto i = 0; i < op->getNumOperands(); ++i) {
    // Convert each operand layout given the result.
    auto operand = op->getOperand(i);
    auto originalLayout = findAttributeAssociatedWith(operand, kLayoutAttrName);
    if (failed(originalLayout)) {
      LLVM_DEBUG(llvm::dbgs()
                 << "Failed to get layout for operand : " << i << "\n");
      return FAILURE;
    }

    LLVM_DEBUG(llvm::dbgs() << "Converting operand " << i << " with layout "
                            << originalLayout.value() << " to layout "
                            << minHoistResult.newInputLayouts[i] << "\n");
    Attribute newInputLayout = minHoistResult.newInputLayouts[i];
    auto newInput = ConvertLayoutOp::create(
        builder, op->getLoc(), operand, originalLayout.value(), newInputLayout);
    setAttributeAssociatedWith(newInput.getResult(), kLayoutAttrName,
                               newInputLayout);
    builder.replaceUsesWithIf(operand, newInput, [&](OpOperand& operand) {
      return operand.getOwner() == op;
    });
  }

  // Set new layout attribute for the op.
  Attribute newOutputLayout = minHoistResult.newOutputLayout;
  op->setAttr(kLayoutAttrName, newOutputLayout);

  // Replace uses of the convert layout op with its input.
  builder.replaceAllUsesWith(minHoistResult.convertLayoutOp.getResult(),
                             minHoistResult.convertLayoutOp.getValue());

  // Ensure that any other convert_layout ops have their from_layouts updated.
  // Update downstream ops.
  for (auto* user : op->getResult(0).getUsers()) {
    if (auto convertLayoutOp = dyn_cast<ConvertLayoutOp>(user)) {
      // Update any convert_layout uses of the op result.
      convertLayoutOp.setFromLayoutAttr(newOutputLayout);
    }
  }

  builder.eraseOp(minHoistResult.convertLayoutOp);
  return SUCCESS;
}

Cost LayoutOptimization::costOfLayoutConversion(Attribute fromLayout,
                                                Attribute toLayout) {
  LayoutAttr fromLayoutAttr = dyn_cast<LayoutAttr>(fromLayout);
  LayoutAttr toLayoutAttr = dyn_cast<LayoutAttr>(toLayout);

  if (!fromLayoutAttr || !toLayoutAttr) {
    return fromLayout == toLayout ? 0 : 1;
  }

  if (fromLayoutAttr == toLayoutAttr) {
    return 0;
  }

  return computeCostOfLayoutConversion(minSlotCount, fromLayoutAttr,
                                       toLayoutAttr, vveRandomSeed,
                                       vveRandomTries);
}

OperandChange LayoutOptimization::costOfChangedOperand(OpOperand& operand,
                                                       Operation* kernel,
                                                       Attribute newLayout,
                                                       DataFlowSolver* solver) {
  auto value = operand.get();

  LLVM_DEBUG(llvm::dbgs() << "Checking if operand has free layout: " << value
                          << "\n");
  if (isLayoutFree(value, solver)) {
    LLVM_DEBUG(llvm::dbgs() << "Layout is free!\n");
    return OperandChange{Attribute(), Attribute(), 0};
  }
  LLVM_DEBUG(llvm::dbgs() << "Layout is not free\n");

  if (auto convertLayoutOp = value.getDefiningOp<ConvertLayoutOp>()) {
    // If the operand came from convert_layout, the cost of the change is
    // (folded conversion - original conversion).
    auto fromLayout = convertLayoutOp.getFromLayout();
    Cost originalConversion =
        costOfLayoutConversion(fromLayout, convertLayoutOp.getToLayout());
    Cost foldedConversion = costOfLayoutConversion(fromLayout, newLayout);
    return OperandChange{fromLayout, newLayout,
                         foldedConversion - originalConversion};
  }

  // Otherwise, we need to insert a new convert_layout op.
  auto originalLayoutResult =
      findAttributeAssociatedWith(value, kLayoutAttrName);
  assert(succeeded(originalLayoutResult) &&
         "Operand does not have a layout attribute");
  auto originalLayout = originalLayoutResult.value();
  return OperandChange{originalLayout, newLayout,
                       costOfLayoutConversion(originalLayout, newLayout)};
}

Cost LayoutOptimization::costOfChangedResult(Operation* kernel,
                                             Attribute newLayout,
                                             KernelAttr oldKernel,
                                             KernelName newKernel) {
  LLVM_DEBUG(llvm::dbgs() << "Computing cost of changed result to new layout "
                          << newLayout << "\n");

  // Note: this assumes that trivial kernels have no result layout cost
  // change. This is true of all current trivial kernels (tensor.collapse_shape
  // and all arithmetic trivial kernels) and will continue to hold as long as
  // trivial kernels are implemented with equivalent ciphertext-semantic
  // operations.
  if ((oldKernel == nullptr && newKernel == KernelName::Trivial) ||
      (oldKernel != nullptr && oldKernel.getName() == KernelName::Trivial &&
       newKernel == KernelName::Trivial)) {
    return 0;
  }

  Cost totalCost = 0;
  for (auto* user : kernel->getResult(0).getUsers()) {
    LLVM_DEBUG(llvm::dbgs() << "\tConvert layout use: " << *user << "\n");
    if (auto convertLayoutOp = dyn_cast<ConvertLayoutOp>(user)) {
      Cost originalConversion = costOfLayoutConversion(
          convertLayoutOp.getFromLayout(), convertLayoutOp.getToLayout());
      Cost foldedConversion =
          costOfLayoutConversion(newLayout, convertLayoutOp.getToLayout());
      totalCost += foldedConversion - originalConversion;
    }
  }
  return totalCost;
}

// Build a symbolic DAG for a kernel and count rotations. Returns
// FailureOr<Cost> to allow caller to handle unsupported kernels.
static FailureOr<Cost> computeKernelCostFromDAG(KernelName kernel,
                                                Operation* op) {
  // In each of the cases below, the symbolic DAG constructed must match the
  // lowering of this kernel in ConvertToCiphertextSemantics.cpp.

  RotationCountVisitor costModel;
  switch (kernel) {
    case KernelName::Trivial:
      return 0;

    case KernelName::MatvecDiagonal: {
      auto matrixType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
      if (!matrixType) return failure();
      auto shape = matrixType.getShape();
      auto dagType = kernel::mlirTypeToDagType(matrixType);

      SymbolicValue symbolicVector({shape[1]},
                                   /*isSecret=*/true);
      SymbolicValue symbolicMatrix({shape[0], shape[1]},
                                   /*isSecret=*/false);
      auto kernelDag = kernel::implementHaleviShoup(
          symbolicVector, symbolicMatrix, shape.vec(), dagType);

      if (!kernelDag) return failure();
      return costModel.process(kernelDag);
    }

    case KernelName::MatvecNaive:
      // TODO(#1589): implement MatvecNaive kernel cost model
      return failure();

    case KernelName::VecmatDiagonal: {
      auto matrixType = dyn_cast<RankedTensorType>(op->getOperand(1).getType());
      if (!matrixType) return failure();
      auto shape = matrixType.getShape();
      auto dagType = kernel::mlirTypeToDagType(matrixType);

      SymbolicValue symbolicVector({shape[0]}, /*isSecret=*/true);
      SymbolicValue symbolicMatrix({shape[0], shape[1]},
                                   /*isSecret=*/false);
      auto kernelDag = kernel::implementHaleviShoup(
          symbolicVector, symbolicMatrix, shape.vec(), dagType);

      if (!kernelDag) return failure();
      return costModel.process(kernelDag);
    }

    case KernelName::MatmulBicyclic: {
      // TODO(#2386): ensure this reflects required zero-padding of inputs
      auto lhsType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
      if (!lhsType) return failure();
      auto lhsShape = lhsType.getShape();
      auto rhsType = dyn_cast<RankedTensorType>(op->getOperand(1).getType());
      if (!rhsType) return failure();
      auto rhsShape = rhsType.getShape();
      auto dagType = kernel::mlirTypeToDagType(lhsType);

      SymbolicValue lhsLeaf(lhsShape, /*isSecret=*/true);
      SymbolicValue rhsLeaf(rhsShape, /*isSecret=*/true);
      auto implementedKernel = kernel::implementBicyclicMatmul(
          lhsLeaf, rhsLeaf, lhsShape[0], lhsShape[1], rhsShape[1], dagType);

      if (!implementedKernel) return failure();
      return costModel.process(implementedKernel);
    }

    case KernelName::Dot: {
      auto lhsType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
      if (!lhsType) return failure();
      auto rhsType = dyn_cast<RankedTensorType>(op->getOperand(1).getType());
      if (!rhsType) return failure();
      auto shape = lhsType.getShape();
      auto dagType = kernel::mlirTypeToDagType(lhsType);

      SymbolicValue lhsVector({shape[0]}, /*isSecret=*/true);
      SymbolicValue rhsVector({shape[0]}, /*isSecret=*/true);

      auto kernelDag =
          kernel::implementDot(lhsVector, rhsVector, shape[0], dagType);

      if (!kernelDag) return failure();
      return costModel.process(kernelDag);
    }

    default:
      return failure();
  }
}

static Cost costOfKernelChange(Operation* op, KernelName oldKernel,
                               const HoistResult& hoistResult) {
  LLVM_DEBUG(llvm::dbgs() << "\tComputing kernel change cost: " << oldKernel
                          << "\n");
  KernelName newKernel = hoistResult.newKernel;

  // No cost if kernel isn't changing
  if (oldKernel == newKernel) {
    return 0;
  }

  // Trivial kernels are free (elementwise operations)
  if (newKernel == KernelName::Trivial) {
    return 0;
  }

  // Build symbolic DAG and count rotations
  // TODO(#2351): enhance cost model to include multiplicative depth
  FailureOr<Cost> costResult = computeKernelCostFromDAG(newKernel, op);

  if (failed(costResult)) {
    LLVM_DEBUG(llvm::dbgs()
               << "Failed to compute cost for kernel " << newKernel << "\n");
    // Skip unsupported kernels by returning large cost
    return std::numeric_limits<Cost>::max() / 2;
  }

  Cost cost = *costResult;
  LLVM_DEBUG(llvm::dbgs() << "Kernel " << newKernel << " cost: " << cost
                          << " rotations\n");

  return cost;
}

std::vector<HoistOption> LayoutOptimization::computeHoistingOptions(
    Operation* op, ConvertLayoutOp convertLayoutOp, DataFlowSolver* solver) {
  LayoutConversionHoistableOpInterface hoistableInterface =
      dyn_cast<LayoutConversionHoistableOpInterface>(op);

  if (!hoistableInterface) {
    // TODO(#1597): Add linalg::reduce op kernel.
    LLVM_DEBUG(llvm::dbgs() << "Encountered op with no hoisting interface. "
                               "Returning empty HoistOption\n");
    return {HoistOption()};
  }

  auto hoisters = hoistableInterface.getHoisters(convertLayoutOp);
  LLVM_DEBUG(llvm::dbgs() << "Evaluating " << hoisters.size()
                          << " hoisting options for " << op->getName() << "\n");

  std::vector<HoistResult> results;
  for (auto& hoister : hoisters) {
    auto result = hoister(convertLayoutOp);
    if (failed(result)) {
      continue;
    }
    results.push_back(result.value());
  }

  LLVM_DEBUG(llvm::dbgs() << results.size()
                          << " successful hoisting options\n");

  auto outputLayout = convertLayoutOp.getToLayout();
  KernelAttr oldKernel = op->getAttrOfType<KernelAttr>(kKernelAttrName);
  std::vector<HoistOption> options;
  for (HoistResult& result : results) {
    HoistOption& option = options.emplace_back();
    option.hoistResult = result;
    option.cost = 0;

    LLVM_DEBUG(llvm::dbgs()
               << "Computing cost of hoisting layout " << result.newOutputLayout
               << " via kernel " << result.newKernel << "\n");

    // A map is used to deduplicate operand changes.
    DenseMap<std::tuple<Value, Attribute, Attribute>, Cost> operandChangeMap;
    SmallVector<Cost> operandChangeCosts;
    for (auto& operand : op->getOpOperands()) {
      auto computedCost =
          costOfChangedOperand(operand, op, outputLayout, solver);
      operandChangeCosts.push_back(computedCost.cost);
      auto key = std::make_tuple(operand.get(), computedCost.fromLayout,
                                 computedCost.toLayout);
      operandChangeMap[key] = computedCost.cost;
    }

    LLVM_DEBUG({
      llvm::dbgs() << "\toperand change costs (with duplicates): ";
      for (auto cost : operandChangeCosts) {
        llvm::dbgs() << cost << ", ";
      }
      llvm::dbgs() << "\n";
    });

    Cost totalOperandChangeCost = 0;
    for (auto& [_, cost] : operandChangeMap) {
      totalOperandChangeCost += cost;
    }
    LLVM_DEBUG(llvm::dbgs() << "\toperand change total cost: "
                            << totalOperandChangeCost << "\n");
    option.cost += totalOperandChangeCost;

    Cost resultChangeCost =
        costOfChangedResult(op, outputLayout, oldKernel, result.newKernel);
    LLVM_DEBUG(llvm::dbgs()
               << "\tresult change cost: " << resultChangeCost << "\n");
    option.cost += resultChangeCost;

    // The op may not have a kernel set, in which case the kernel may be
    // trivial and not explicitly marked; in this case we can ignore kernel
    // costs. Otherwise, we can ignore a kernel cost if this hoisting option
    // doesn't change the kernel.
    if ((oldKernel == nullptr &&
         option.hoistResult.newKernel == KernelName::Trivial) ||
        (oldKernel != nullptr &&
         oldKernel.getName() == option.hoistResult.newKernel))
      continue;

    Cost kernelChangeCost =
        costOfKernelChange(op, oldKernel.getName(), option.hoistResult);
    LLVM_DEBUG(llvm::dbgs()
               << "\tkernel change cost: " << kernelChangeCost << "\n");
    option.cost += kernelChangeCost;
  }
  return options;
}

}  // namespace heir
}  // namespace mlir
