#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/raw_os_ostream.h"

#include <fstream>
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
#include "src/Accelerators/PIM/Pass/PIMPasses.h"

using namespace mlir;

namespace onnx_mlir {
namespace spatial {
namespace {

template <typename ComputeOp>
static bool hasOnlyStructuralAttrs(ComputeOp op) {
  return llvm::all_of(op->getAttrs(), [&](NamedAttribute attr) {
    return attr.getName() == op.getOperandSegmentSizesAttrName()
        || (isa<SpatGraphComputeBatch>(op.getOperation())
            && attr.getName() == cast<SpatGraphComputeBatch>(op.getOperation()).getLaneCountAttrName());
  });
}

static bool hasCapacityFor(Operation* producer, Operation* consumer,
                           size_t residentWeightCapacity) {
  ResidentWeightSet producerWeights =
    collectDistinctResidentWeights(producer);
  ResidentWeightSet consumerWeights =
    collectDistinctResidentWeights(consumer);
  return getResidentWeightUnionSize(producerWeights, consumerWeights)
      <= residentWeightCapacity;
}

template <typename ConsumerOp>
static bool isUniqueGraphComputePredecessor(Operation *candidate, ConsumerOp consumer) {
  llvm::SmallSetVector<Operation *, 4> predecessors;
  for (Value input : consumer.getInputs()) {
    Operation *producer = input.getDefiningOp();
    if (producer && isGraphComputeLike(producer))
      predecessors.insert(producer);
  }
  return predecessors.size() == 1 && predecessors.front() == candidate;
}

struct TrivialGraphMergeStats {
  size_t scalarBefore = 0;
  size_t batchBefore = 0;
  size_t scalarAfter = 0;
  size_t batchAfter = 0;
  size_t scalarProducerConsumerMerges = 0;
  size_t batchProducerConsumerMerges = 0;
  size_t leadingUnitNormalizationFolds = 0;
};

static std::pair<size_t, size_t> countGraphComputes(ModuleOp module) {
  size_t scalar = 0, batch = 0;
  module.walk([&](Operation *op) {
    scalar += isa<SpatGraphCompute>(op);
    batch += isa<SpatGraphComputeBatch>(op);
  });
  return {scalar, batch};
}

static void dumpTrivialMergeReport(const TrivialGraphMergeStats &stats,
                                   size_t largestBatchLaneCount) {
  std::fstream file = openDialectDumpFileWithExtension("spatial2_trivial_merged", "/reports", "txt");
  if (!file.is_open())
    return;
  llvm::raw_os_ostream os(file);
  size_t before = stats.scalarBefore + stats.batchBefore;
  size_t after = stats.scalarAfter + stats.batchAfter;
  size_t removed = before - after;
  double percentage = before ? 100.0 * removed / before : 0.0;
  os << "Summary\n"
     << "  graph computes: " << before << " -> " << after << "\n"
     << "    removed: " << removed << " (" << llvm::formatv("{0:F2}", percentage) << "%)\n"
     << "  scalar: " << stats.scalarBefore << " -> " << stats.scalarAfter << "\n"
     << "  batch: " << stats.batchBefore << " -> " << stats.batchAfter << "\n"
     << "  transformations:\n"
     << "    scalar producer-consumer merges: " << stats.scalarProducerConsumerMerges << "\n"
     << "    batch producer-consumer merges: " << stats.batchProducerConsumerMerges << "\n"
     << "    leading-unit normalization folds: " << stats.leadingUnitNormalizationFolds << "\n\n"
     << "Resulting graph\n"
     << "  graph computes: " << after << "\n"
     << "  scalar: " << stats.scalarAfter << "\n"
     << "  batch: " << stats.batchAfter << "\n"
     << "  largest batch lane count: " << largestBatchLaneCount << "\n";
}

template <typename ProducerOp, typename ConsumerOp>
static bool isExclusivelyConsumedBy(ProducerOp producer, ConsumerOp consumer) {
  bool hasDependency = false;
  for (Value result : producer.getResults()) {
    for (OpOperand& use : result.getUses()) {
      if (use.getOwner() != consumer.getOperation() || !llvm::is_contained(consumer.getInputs(), result))
        return false;
      hasDependency = true;
    }
  }
  return hasDependency;
}

static bool hasNoNestedArgumentCaptures(SpatGraphCompute compute) {
  Block &body = compute.getBody().front();
  return llvm::all_of(body.getArguments(), [&](BlockArgument argument) {
    return llvm::all_of(argument.getUsers(), [&](Operation *user) { return user->getBlock() == &body; });
  });
}

template <typename ProducerOp, typename ConsumerOp>
static void collectExternalOperands(ProducerOp producer,
                                    ConsumerOp consumer,
                                    llvm::SetVector<Value>& weights,
                                    llvm::SetVector<Value>& inputs) {
  weights.insert(producer.getWeights().begin(), producer.getWeights().end());
  weights.insert(consumer.getWeights().begin(), consumer.getWeights().end());
  auto appendInput = [&](Value value) {
    if (value.getDefiningOp() != producer.getOperation() && !weights.contains(value))
      inputs.insert(value);
  };
  llvm::for_each(producer.getInputs(), appendInput);
  llvm::for_each(consumer.getInputs(), appendInput);
}

template <typename OldOp, typename NewOp>
static void mapExternalArguments(OldOp oldOp, NewOp newOp, IRMapping& mapper, bool mapForwardedInputs = true) {
  for (auto [index, operand] : llvm::enumerate(oldOp.getWeights())) {
    auto oldArg = oldOp.getWeightArgument(index);
    auto newOperand = llvm::find(newOp.getWeights(), operand);
    assert(oldArg && newOperand != newOp.getWeights().end());
    mapper.map(*oldArg, *newOp.getWeightArgument(std::distance(newOp.getWeights().begin(), newOperand)));
  }
  for (auto [index, operand] : llvm::enumerate(oldOp.getInputs())) {
    auto oldArg = oldOp.getInputArgument(index);
    assert(oldArg);
    if (mapper.contains(operand)) {
      if (mapForwardedInputs)
        mapper.map(*oldArg, mapper.lookup(operand));
      continue;
    }
    auto newOperand = llvm::find(newOp.getInputs(), operand);
    assert(oldArg && newOperand != newOp.getInputs().end());
    mapper.map(*oldArg, *newOp.getInputArgument(std::distance(newOp.getInputs().begin(), newOperand)));
  }
}

struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
  MergeTrivialScalarComputes(MLIRContext *context,
                             TrivialGraphMergeStats *stats,
                             size_t residentWeightCapacity)
      : OpRewritePattern(context), stats(stats),
        residentWeightCapacity(residentWeightCapacity) {}

  LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
    SpatGraphCompute producer;
    for (Value input : consumer.getInputs()) {
      auto candidate = input.getDefiningOp<SpatGraphCompute>();
      if (candidate && candidate->getBlock() == consumer->getBlock() && hasOnlyStructuralAttrs(candidate)
          && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
          && isExclusivelyConsumedBy(candidate, consumer)
          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
          && hasNoNestedArgumentCaptures(candidate)
          && hasNoNestedArgumentCaptures(consumer)) {
        producer = candidate;
        break;
      }
    }
    if (!producer)
      return failure();

    llvm::SetVector<Value> weights, inputs;
    collectExternalOperands(producer, consumer, weights, inputs);
    rewriter.setInsertionPoint(consumer);
    SpatGraphCompute merged = createEmptySpatGraphCompute(
      rewriter, consumer.getLoc(), consumer.getResultTypes(), weights.getArrayRef(), inputs.getArrayRef());

    IRMapping mapper;
    mapExternalArguments(producer, merged, mapper);
    for (Operation& op : producer.getBody().front().without_terminator())
      rewriter.clone(op, mapper);
    auto producerYield = cast<SpatYieldOp>(producer.getBody().front().getTerminator());
    for (auto [result, yielded] : llvm::zip(producer.getResults(), producerYield.getOutputs()))
      mapper.map(result, mapper.lookupOrDefault(yielded));

    mapExternalArguments(consumer, merged, mapper);
    for (Operation& op : consumer.getBody().front())
      rewriter.clone(op, mapper);

    rewriter.replaceOp(consumer, merged.getResults());
    rewriter.eraseOp(producer);
    ++stats->scalarProducerConsumerMerges;
    return success();
  }

private:
  TrivialGraphMergeStats *stats;
  size_t residentWeightCapacity;
};

static bool isLaneIndex(Value value, Value lane, int64_t laneCount) {
  if (!value)
    return false;
  if (value == lane)
    return true;
  auto apply = value.getDefiningOp<affine::AffineApplyOp>();
  if (!apply || apply.getMapOperands().size() != 1 || apply.getMapOperands().front() != lane)
    return false;
  AffineMap map = apply.getAffineMap();
  if (map.getNumDims() != 1 || map.getNumSymbols() != 0 || map.getNumResults() != 1)
    return false;
  AffineExpr expression = map.getResult(0);
  if (auto dim = dyn_cast<AffineDimExpr>(expression))
    return dim.getPosition() == 0;
  auto modulo = dyn_cast<AffineBinaryOpExpr>(expression);
  if (!modulo || modulo.getKind() != AffineExprKind::Mod)
    return false;
  auto dim = dyn_cast<AffineDimExpr>(modulo.getLHS());
  auto divisor = dyn_cast<AffineConstantExpr>(modulo.getRHS());
  return dim && dim.getPosition() == 0
      && divisor && divisor.getValue() >= laneCount;
}

static bool isDirectLaneSlot(ArrayRef<OpFoldResult> offsets,
                             ArrayRef<OpFoldResult> sizes,
                             ArrayRef<OpFoldResult> strides,
                             Value lane,
                             RankedTensorType physicalType) {
  size_t rank = physicalType.getRank();
  if (offsets.size() != rank || sizes.size() != rank || strides.size() != rank
      || !isLaneIndex(dyn_cast<Value>(offsets.front()), lane, physicalType.getDimSize(0)))
    return false;
  for (size_t dim = 0; dim < rank; ++dim) {
    int64_t expectedOffset = 0;
    int64_t expectedSize = dim == 0 ? 1 : physicalType.getDimSize(dim);
    if ((dim != 0 && getConstantIntValue(offsets[dim]) != expectedOffset)
        || getConstantIntValue(sizes[dim]) != expectedSize || getConstantIntValue(strides[dim]) != 1)
      return false;
  }
  return true;
}

static FailureOr<SmallVector<Value>> collectPublishedFragments(SpatGraphComputeBatch producer) {
  auto terminator = dyn_cast<SpatInParallelOp>(producer.getBody().front().getTerminator());
  auto lane = producer.getLaneArgument();
  if (!terminator || !lane)
    return failure();
  SmallVector<Value> fragments(producer.getNumResults());
  for (Operation& op : terminator.getRegion().front()) {
    auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
    auto destination = insert ? dyn_cast<BlockArgument>(insert.getDest()) : BlockArgument();
    if (!insert || !destination)
      return failure();
    unsigned resultIndex = 0;
    while (resultIndex < producer.getNumResults() && producer.getOutputArgument(resultIndex) != destination)
      ++resultIndex;
    if (resultIndex == producer.getNumResults() || fragments[resultIndex])
      return failure();
    auto physicalType = dyn_cast<RankedTensorType>(producer.getResult(resultIndex).getType());
    auto fragmentType = dyn_cast<RankedTensorType>(insert.getSource().getType());
    auto expectedFragmentType = getGraphBatchFragmentType(physicalType, producer.getLaneCount());
    if (!physicalType || !fragmentType || failed(expectedFragmentType) || *expectedFragmentType != fragmentType
        || !isDirectLaneSlot(
          insert.getMixedOffsets(), insert.getMixedSizes(), insert.getMixedStrides(), *lane, physicalType))
      return failure();
    fragments[resultIndex] = insert.getSource();
  }
  if (!llvm::all_of(fragments, [](Value value) { return value; }))
    return failure();
  return fragments;
}

static bool matchLeadingUnitNormalization(SpatGraphComputeBatch producer, SpatGraphCompute consumer) {
  if (producer.getNumResults() != 1 || !consumer.getWeights().empty()
      || consumer.getInputs().size() != 1 || consumer.getInputs().front() != producer.getResult(0)
      || consumer.getNumResults() != 1)
    return false;
  auto yield = dyn_cast<SpatYieldOp>(consumer.getBody().front().getTerminator());
  auto loopResult = yield && yield.getOutputs().size() == 1
    ? dyn_cast<OpResult>(yield.getOutputs().front())
    : OpResult();
  auto loop = loopResult ? dyn_cast<scf::ForOp>(loopResult.getOwner()) : scf::ForOp();
  auto targetType = dyn_cast<RankedTensorType>(consumer.getResult(0).getType());
  if (!loop || llvm::range_size(consumer.getBody().front().without_terminator()) != 2
      || loop->getBlock() != &consumer.getBody().front() || loopResult.getResultNumber() != 0
      || loop.getNumResults() != 1 || loop.getResult(0).getType() != targetType
      || !targetType || !targetType.hasStaticShape()
      || getConstantIntValue(loop.getLowerBound()) != 0
      || getConstantIntValue(loop.getUpperBound()) != producer.getLaneCount()
      || getConstantIntValue(loop.getStep()) != 1
      || llvm::range_size(loop.getBody()->without_terminator()) != 2)
    return false;
  auto empty = loop.getInitArgs().front().getDefiningOp<tensor::EmptyOp>();
  auto bodyIt = loop.getBody()->begin();
  auto extract = dyn_cast<tensor::ExtractSliceOp>(&*bodyIt++);
  auto insert = dyn_cast<tensor::InsertSliceOp>(&*bodyIt);
  auto loopYield = dyn_cast<scf::YieldOp>(loop.getBody()->getTerminator());
  auto sourceType = dyn_cast<RankedTensorType>(producer.getResult(0).getType());
  auto sourceFragmentType = sourceType
    ? getGraphBatchFragmentType(sourceType, producer.getLaneCount())
    : FailureOr<RankedTensorType>(failure());
  auto targetFragmentType = getGraphBatchFragmentType(targetType, producer.getLaneCount());
  return empty && empty->getBlock() == &consumer.getBody().front() && empty.getType() == targetType
      && extract && insert && loopYield && loopYield.getResults().size() == 1
      && extract.getSource() == consumer.getInputArgument(0)
      && insert.getSource() == extract.getResult() && insert.getDest() == loop.getRegionIterArgs().front()
      && loopYield.getResults().front() == insert.getResult()
      && succeeded(sourceFragmentType) && succeeded(targetFragmentType)
      && extract.getType() == *sourceFragmentType
      && sourceFragmentType->getRank() == targetFragmentType->getRank() + 1
      && sourceFragmentType->getDimSize(0) == 1
      && sourceFragmentType->getShape().drop_front() == targetFragmentType->getShape()
      && sourceFragmentType->getElementType() == targetFragmentType->getElementType()
      && isDirectLaneSlot(extract.getMixedOffsets(), extract.getMixedSizes(), extract.getMixedStrides(),
                          loop.getInductionVar(), sourceType)
      && isDirectLaneSlot(insert.getMixedOffsets(), insert.getMixedSizes(), insert.getMixedStrides(),
                          loop.getInductionVar(), targetType);
}

struct FoldBatchLeadingUnitNormalization : OpRewritePattern<SpatGraphCompute> {
  FoldBatchLeadingUnitNormalization(MLIRContext *context, TrivialGraphMergeStats *stats)
      : OpRewritePattern(context), stats(stats) {}

  LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
    auto producer = consumer.getInputs().empty()
      ? SpatGraphComputeBatch()
      : consumer.getInputs().front().getDefiningOp<SpatGraphComputeBatch>();
    if (!producer || producer->getBlock() != consumer->getBlock() || !hasOnlyStructuralAttrs(producer)
        || !hasOnlyStructuralAttrs(consumer) || !isUniqueGraphComputePredecessor(producer, consumer)
        || !isExclusivelyConsumedBy(producer, consumer))
      return failure();
    auto fragments = collectPublishedFragments(producer);
    if (!matchLeadingUnitNormalization(producer, consumer) || failed(fragments))
      return failure();

    rewriter.setInsertionPoint(consumer);
    auto folded = createEmptySpatGraphComputeBatch(rewriter, consumer.getLoc(), consumer.getResultTypes(),
      producer.getLaneCount(), producer.getWeights(), producer.getInputs());
    if (failed(folded))
      return failure();
    IRMapping mapper;
    mapper.map(*producer.getLaneArgument(), *folded->getLaneArgument());
    mapExternalArguments(producer, *folded, mapper);
    for (Operation& op : producer.getBody().front().without_terminator())
      rewriter.clone(op, mapper);

    auto outputType = cast<RankedTensorType>(consumer.getResult(0).getType());
    auto fragmentType = *getGraphBatchFragmentType(outputType, producer.getLaneCount());
    auto fragment = removeLeadingUnitTensorDimension(
      rewriter, consumer.getLoc(), mapper.lookup(fragments->front()), fragmentType);
    assert(succeeded(fragment) && "normalization fragment types were prechecked");
    publishGraphBatchPhysicalFragment(rewriter, consumer.getLoc(), *fragment,
      *folded->getOutputArgument(0), *folded->getLaneArgument());
    rewriter.replaceOp(consumer, folded->getResults());
    rewriter.eraseOp(producer);
    ++stats->leadingUnitNormalizationFolds;
    return success();
  }

private:
  TrivialGraphMergeStats *stats;
};

static bool hasDirectLaneConsumers(SpatGraphComputeBatch producer, SpatGraphComputeBatch consumer) {
  auto lane = consumer.getLaneArgument();
  if (!lane)
    return false;
  for (auto [inputIndex, input] : llvm::enumerate(consumer.getInputs())) {
    if (input.getDefiningOp() != producer.getOperation())
      continue;
    auto inputArg = consumer.getInputArgument(inputIndex);
    auto physicalType = dyn_cast<RankedTensorType>(input.getType());
    if (!inputArg || !physicalType || inputArg->use_empty())
      return false;
    for (Operation* user : inputArg->getUsers()) {
      auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
      if (!extract || extract.getSource() != *inputArg
          || !isDirectLaneSlot(
            extract.getMixedOffsets(), extract.getMixedSizes(), extract.getMixedStrides(), *lane, physicalType))
        return false;
      auto resultType = dyn_cast<RankedTensorType>(extract.getType());
      auto fragmentType = getGraphBatchFragmentType(physicalType, producer.getLaneCount());
      if (failed(fragmentType) || !resultType
          || (resultType != *fragmentType
              && (resultType.getRank() != fragmentType->getRank() + 1 || resultType.getDimSize(0) != 1
                  || resultType.getShape().drop_front() != fragmentType->getShape()
                  || resultType.getElementType() != fragmentType->getElementType())))
        return false;
    }
  }
  return true;
}

struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
  MergeTrivialBatchComputes(MLIRContext *context,
                            TrivialGraphMergeStats *stats,
                            size_t residentWeightCapacity)
      : OpRewritePattern(context), stats(stats),
        residentWeightCapacity(residentWeightCapacity) {}

  LogicalResult matchAndRewrite(SpatGraphComputeBatch consumer, PatternRewriter& rewriter) const override {
    SpatGraphComputeBatch producer;
    FailureOr<SmallVector<Value>> fragments = failure();
    for (Value input : consumer.getInputs()) {
      auto candidate = input.getDefiningOp<SpatGraphComputeBatch>();
      if (candidate && candidate->getBlock() == consumer->getBlock()
          && candidate.getLaneCount() == consumer.getLaneCount() && hasOnlyStructuralAttrs(candidate)
          && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
          && isExclusivelyConsumedBy(candidate, consumer)
          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
          && hasDirectLaneConsumers(candidate, consumer)
          && succeeded(fragments = collectPublishedFragments(candidate))) {
        producer = candidate;
        break;
      }
    }
    if (!producer)
      return failure();

    llvm::SetVector<Value> weights, inputs;
    collectExternalOperands(producer, consumer, weights, inputs);
    rewriter.setInsertionPoint(consumer);
    SpatGraphComputeBatch merged = *createEmptySpatGraphComputeBatch(rewriter, consumer.getLoc(), consumer.getResultTypes(),
      consumer.getLaneCount(), weights.getArrayRef(), inputs.getArrayRef());
    IRMapping mapper;
    mapper.map(*producer.getLaneArgument(), *merged.getLaneArgument());
    mapExternalArguments(producer, merged, mapper);
    for (Operation& op : producer.getBody().front().without_terminator())
      rewriter.clone(op, mapper);
    for (auto [result, fragment] : llvm::zip(producer.getResults(), *fragments))
      mapper.map(result, mapper.lookupOrDefault(fragment));

    mapper.map(*consumer.getLaneArgument(), *merged.getLaneArgument());
    mapExternalArguments(consumer, merged, mapper, /*mapForwardedInputs=*/false);
    for (auto [index, output] : llvm::enumerate(consumer.getResults()))
      mapper.map(*consumer.getOutputArgument(index), *merged.getOutputArgument(index));
    for (Operation& op : consumer.getBody().front().without_terminator()) {
      auto extract = dyn_cast<tensor::ExtractSliceOp>(op);
      auto sourceArg = extract ? dyn_cast<BlockArgument>(extract.getSource()) : BlockArgument();
      unsigned firstInputArg = 1 + consumer.getWeights().size();
      if (!sourceArg || sourceArg.getOwner() != &consumer.getBody().front() || sourceArg.getArgNumber() < firstInputArg
          || sourceArg.getArgNumber() >= firstInputArg + consumer.getInputs().size()) {
        rewriter.clone(op, mapper);
        continue;
      }
      unsigned inputIndex = sourceArg.getArgNumber() - firstInputArg;
      Value input = consumer.getInputs()[inputIndex];
      if (input.getDefiningOp() != producer.getOperation()) {
        rewriter.clone(op, mapper);
        continue;
      }
      Value fragment = mapper.lookup(input);
      if (fragment.getType() != extract.getType()) {
        auto expanded = addLeadingUnitTensorDimension(rewriter, extract.getLoc(), fragment);
        assert(succeeded(expanded) && expanded->getType() == extract.getType()
               && "prechecked physical fragment type must be forwardable");
        fragment = *expanded;
      }
      mapper.map(extract.getResult(), fragment);
    }
    rewriter.clone(*consumer.getBody().front().getTerminator(), mapper);

    rewriter.replaceOp(consumer, merged.getResults());
    rewriter.eraseOp(producer);
    ++stats->batchProducerConsumerMerges;
    return success();
  }

private:
  TrivialGraphMergeStats *stats;
  size_t residentWeightCapacity;
};

struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMergePass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrivialGraphComputeMergePass)

  TrivialGraphComputeMergePass() = default;
  explicit TrivialGraphComputeMergePass(size_t residentWeightCapacity)
      : residentWeightCapacity(residentWeightCapacity) {}

  StringRef getArgument() const override { return "pim-trivial-graph-compute-merge"; }
  StringRef getDescription() const override {
    return "Inline linear exclusive graph compute chains while preserving fan-in boundaries.";
  }

  void runOnOperation() override {
    ModuleOp module = getOperation();
    if (residentWeightCapacity == 0) {
      module.emitError(
        "TrivialGraphComputeMerge requires an explicit valid resident-weight capacity");
      signalPassFailure();
      return;
    }
    TrivialGraphMergeStats stats;
    std::tie(stats.scalarBefore, stats.batchBefore) = countGraphComputes(module);
    RewritePatternSet patterns(&getContext());
    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(
      &getContext(), &stats, residentWeightCapacity);
    if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
      signalPassFailure();
      return;
    }
    RewritePatternSet normalization(&getContext());
    normalization.add<FoldBatchLeadingUnitNormalization>(&getContext(), &stats);
    if (failed(applyPatternsGreedily(module, std::move(normalization)))) {
      signalPassFailure();
      return;
    }
    std::tie(stats.scalarAfter, stats.batchAfter) = countGraphComputes(module);
    size_t largestBatchLaneCount = 0;
    module.walk([&](SpatGraphComputeBatch batch) {
      largestBatchLaneCount = std::max(largestBatchLaneCount, static_cast<size_t>(batch.getLaneCount()));
    });
    dumpTrivialMergeReport(stats, largestBatchLaneCount);
    dumpModule(module, "spatial2_trivial_merged");
    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial2)) {
      auto entryFunc = getPimEntryFunc(module);
      if (failed(entryFunc)
          || failed(exportSpatialDataflowCsvGraph(*entryFunc, "spatial2_trivial_merged", "spatial2")))
        signalPassFailure();
    }
  }

private:
  size_t residentWeightCapacity = 0;
};

} // namespace
} // namespace spatial

std::unique_ptr<Pass> createTrivialGraphComputeMergePass() {
  return std::make_unique<spatial::TrivialGraphComputeMergePass>();
}

std::unique_ptr<Pass> createTrivialGraphComputeMergePass(
    size_t residentWeightCapacity) {
  return std::make_unique<spatial::TrivialGraphComputeMergePass>(
    residentWeightCapacity);
}

} // namespace onnx_mlir