#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"

#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
#include "mlir/Transforms/Passes.h"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

static constexpr StringLiteral kDenseLayout = "dense_nchw";
static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";

static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
                                                         Value value) {
  auto it = rowStripValues.find(value);
  if (it == rowStripValues.end())
    return failure();
  return it->second;
}

static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
                                                           Value storage) {
  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
  if (!logicalType)
    return blueprint.emitOpError("requires ranked logical output type"), failure();
  RowStripPhysicalValue value;
  value.storage = storage;
  value.logicalType = logicalType;
  value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
  value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
  if (blueprint.getIndexMap() != kRowStripIndexMap)
    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
  if (!storageType || storageType != getRowStripStorageType(logicalType))
    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
  return value;
}

static FailureOr<Value>
lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
  return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc());
}

static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
                                             spatial::SpatBiasAddPlanOp planOp,
                                             PatternRewriter& rewriter) {
  return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
}

static FailureOr<Value>
materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
    return failure();
  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
    return failure();
  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
}

static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
                                               PatternRewriter& rewriter, Location loc) {
  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
  auto inputType = dyn_cast<RankedTensorType>(input.getType());
  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
    return failure();
  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
    return failure();
  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
    if (outputDim > inputDim)
      return failure();

  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
      if (failed(fragment))
        return failure();
      MixedSliceGeometry biasSlice;
      for (int64_t dim : inputFragmentType->getShape()) {
        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
      }
      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
      if (!biasFragment)
        return failure();
      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
      MixedSliceGeometry outputSlice;
      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
      for (int64_t dim : outputFragmentType->getShape())
        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
      if (!output)
        return failure();
      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
      return success();
    });
  if (failed(batch))
    return failure();
  return batch->getResult(0);
}

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

  StringRef getArgument() const override { return "lower-spatial-plans"; }
  StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }

  void runOnOperation() override {
    ModuleOp moduleOp = getOperation();
    MLIRContext* ctx = moduleOp.getContext();
    auto entryFunc = getPimEntryFunc(moduleOp);
    if (failed(entryFunc)) {
      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
      signalPassFailure();
      return;
    }
    func::FuncOp funcOp = *entryFunc;
    PatternRewriter rewriter(ctx);
    llvm::DenseMap<Value, RowStripPhysicalValue> rowStripValues;
    llvm::SmallPtrSet<Operation*, 16> eraseAfterLowering;
    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
      if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
        return true;
      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
      signalPassFailure();
      return false;
    };

    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
      return;
    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
      if (auto planOp = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
        FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(rowStripValues, planOp.getInput());
        auto rowStripBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
        });
        if (rowStripBlueprint != planOp.getResult().getUsers().end()) {
          rewriter.setInsertionPoint(planOp);
          FailureOr<Value> lowered = lowerSelectedConv2DPlan(
            planOp,
            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
            /*emitRowStripLayout=*/true,
            rewriter);
          if (failed(lowered)) {
            planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
            signalPassFailure();
            return;
          }
          auto blueprint = cast<spatial::SpatBlueprintOp>(*rowStripBlueprint);
          FailureOr<RowStripPhysicalValue> rowStripValue = buildRowStripValue(blueprint, *lowered);
          if (failed(rowStripValue)) {
            signalPassFailure();
            return;
          }
          rowStripValues[blueprint.getResult()] = *rowStripValue;
          eraseAfterLowering.insert(planOp);
          eraseAfterLowering.insert(blueprint);
          continue;
        }
        rewriter.setInsertionPoint(planOp);
        FailureOr<Value> lowered =
          lowerSelectedConv2DPlan(planOp, std::nullopt, /*emitRowStripLayout=*/false, rewriter);
        if (failed(lowered)) {
          planOp.emitOpError("failed to lower selected Spatial Conv plan");
          signalPassFailure();
          return;
        }
        rewriter.replaceOp(planOp, *lowered);
        continue;
      }

      if (auto planOp = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
          });
          if (outputBlueprint == planOp.getResult().getUsers().end()) {
            planOp.emitOpError("row-strip Relu plan requires a row-strip blueprint result");
            signalPassFailure();
            return;
          }

          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
          rewriter.setInsertionPoint(planOp);
          FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
          if (failed(lowered)) {
            planOp.emitOpError("failed to lower selected row-strip Spatial Relu plan");
            signalPassFailure();
            return;
          }
          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
          if (failed(output)) {
            signalPassFailure();
            return;
          }
          rowStripValues[blueprint.getResult()] = *output;
          eraseAfterLowering.insert(planOp);
          eraseAfterLowering.insert(blueprint);
          continue;
        }

        rewriter.setInsertionPoint(planOp);
        auto computeOp = createSpatCompute<1>(
          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
            auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
          });
        rewriter.replaceOp(planOp, computeOp.getResults());
        continue;
      }
      if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
        });
        if (outputBlueprint == planOp.getResult().getUsers().end()) {
          planOp.emitOpError("selected MaxPool plan requires a row-strip blueprint result");
          signalPassFailure();
          return;
        }

        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
        rewriter.setInsertionPoint(planOp);
        FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
          planOp, succeeded(input) ? std::optional<Value> {input->storage} : std::nullopt, rewriter);
        if (failed(lowered)) {
          planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
          signalPassFailure();
          return;
        }
        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
        FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
        if (failed(output)) {
          signalPassFailure();
          return;
        }
        rowStripValues[blueprint.getResult()] = *output;
        eraseAfterLowering.insert(planOp);
        eraseAfterLowering.insert(blueprint);
        continue;
      }
      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
          });
          if (outputBlueprint == planOp.getResult().getUsers().end()) {
            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
            signalPassFailure();
            return;
          }

          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
          rewriter.setInsertionPoint(planOp);
          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
          if (failed(lowered)) {
            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
            signalPassFailure();
            return;
          }
          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
          if (failed(output)) {
            signalPassFailure();
            return;
          }
          rowStripValues[blueprint.getResult()] = *output;
          eraseAfterLowering.insert(planOp);
          eraseAfterLowering.insert(blueprint);
          continue;
        }

        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
        if (!resultType) {
          planOp.emitOpError("requires ranked output type");
          signalPassFailure();
          return;
        }
        rewriter.setInsertionPoint(planOp);
        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
        if (failed(denseBias)) {
          planOp.emitOpError("failed to materialize dense Conv-style bias");
          signalPassFailure();
          return;
        }
        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
          if (succeeded(lowered)) {
            rewriter.replaceOp(planOp, *lowered);
            continue;
          }
        }
        auto computeOp = createSpatCompute<2>(rewriter,
                                              planOp.getLoc(),
                                              planOp.getOutput().getType(),
                                              {},
                                              ValueRange {planOp.getInput(), *denseBias},
                                              [&](Value x, Value y) {
                                                auto added = spatial::SpatVAddOp::create(
                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
                                              });
        rewriter.replaceOp(planOp, computeOp.getResults());
        continue;
      }
      if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
        if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
            && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
          rewriter.replaceOp(materializeOp, materializeOp.getInput());
          continue;
        }
        if (materializeOp.getSourcePhysicalLayout() != kRowStripLayout
            || materializeOp.getTargetPhysicalLayout() != kDenseLayout) {
          materializeOp.emitOpError("non-dense materialize_layout lowering is not supported yet");
          signalPassFailure();
          return;
        }
        FailureOr<RowStripPhysicalValue> rowStripValue = getRowStripValue(rowStripValues, materializeOp.getInput());
        if (failed(rowStripValue)) {
          materializeOp.emitOpError("expected a row-strip blueprint input during row-strip materialization");
          signalPassFailure();
          return;
        }
        rewriter.setInsertionPoint(materializeOp);
        FailureOr<Value> dense = materializeRowStripToDense(*rowStripValue, materializeOp.getLoc(), rewriter);
        if (failed(dense)) {
          materializeOp.emitOpError("failed to materialize selected row-strip layout back to dense NCHW");
          signalPassFailure();
          return;
        }
        rewriter.replaceOp(materializeOp, *dense);
        continue;
      }
      if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
        if (std::optional<StringRef> mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly")
          continue;
        if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
          rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
          continue;
        }
        if (blueprintOp.getPhysicalLayout() != kRowStripLayout) {
          blueprintOp.emitOpError("non-dense blueprint lowering is not supported yet");
          signalPassFailure();
          return;
        }
        if (!eraseAfterLowering.contains(blueprintOp)) {
          blueprintOp.emitOpError("unhandled row-strip blueprint remained during LowerSpatialPlans");
          signalPassFailure();
          return;
        }
      }
    }
    bool erasedAny = true;
    while (erasedAny) {
      erasedAny = false;
      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
        if (!eraseAfterLowering.contains(&op))
          continue;
        if (!op.use_empty())
          continue;
        eraseAfterLowering.erase(&op);
        rewriter.eraseOp(&op);
        erasedAny = true;
      }
    }
    if (!eraseAfterLowering.empty()) {
      for (Operation& op : funcOp.getBody().front())
        if (eraseAfterLowering.contains(&op))
          op.emitOpError("selected row-strip planning op could not be fully eliminated during LowerSpatialPlans");
      signalPassFailure();
      return;
    }
    ConversionTarget helperTarget(*ctx);
    helperTarget.addLegalDialect<spatial::SpatialDialect,
                                 tensor::TensorDialect,
                                 linalg::LinalgDialect,
                                 affine::AffineDialect,
                                 arith::ArithDialect,
                                 scf::SCFDialect,
                                 func::FuncDialect>();
    helperTarget.addLegalOp<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
    helperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
    helperTarget.markOpRecursivelyLegal<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();

    RewritePatternSet helperPatterns(ctx);
    populateGemmPatterns(helperPatterns, ctx);
    populateTransposePatterns(helperPatterns, ctx);
    FrozenRewritePatternSet frozenHelperPatterns(
      std::move(helperPatterns));
    SmallVector<Operation*> topLevelHelperOps;
    funcOp.walk([&](Operation* op) {
      if (isa<spatial::SpatGraphCompute,
              spatial::SpatGraphComputeBatch>(op))
        return WalkResult::skip();
      if (isa<ONNXGemmOp, ONNXTransposeOp>(op))
        topLevelHelperOps.push_back(op);
      return WalkResult::advance();
    });
    for (Operation *helper : topLevelHelperOps) {
      if (failed(applyPartialConversion(
            helper, helperTarget, frozenHelperPatterns))) {
        moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
        signalPassFailure();
        return;
      }
    }
    ConversionTarget nestedHelperTarget(*ctx);
    nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
                                       tensor::TensorDialect,
                                       linalg::LinalgDialect,
                                       affine::AffineDialect,
                                       arith::ArithDialect,
                                       scf::SCFDialect,
                                       func::FuncDialect>();
    nestedHelperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
    SmallVector<Operation*> computeLikeOps;
    funcOp.walk([&](Operation* op) {
      if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(op))
        computeLikeOps.push_back(op);
    });
    for (Operation* op : computeLikeOps) {
      if (failed(applyFullConversion(
            op, nestedHelperTarget, frozenHelperPatterns))) {
        op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
        signalPassFailure();
        return;
      }
    }
    if (!verifyLogicalPhase("after nested helper conversions"))
      return;
    bool hasIllegalOps = false;
    moduleOp.walk([&](Operation* op) {
      if (isa<ONNXEntryPointOp>(op))
        return;
      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
        if (std::optional<StringRef> mode = blueprint.getMode(); mode && *mode == "fragment_assembly")
          return;
        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
        hasIllegalOps = true;
      } else if (isa<spatial::SpatConv2DPlanOp,
                     spatial::SpatBiasAddPlanOp,
                     spatial::SpatReluPlanOp,
                     spatial::SpatMaxPool2DPlanOp,
                     spatial::SpatMaterializeLayoutOp>(op)
          || op->getDialect()->getNamespace() == "onnx") {
        op->emitOpError("operation must not remain after LowerSpatialPlans");
        hasIllegalOps = true;
      }
    });

    PassManager canonicalizationPM(ctx);
    canonicalizationPM.addPass(createCanonicalizerPass());
    if (failed(canonicalizationPM.run(moduleOp)))
      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");

    if (hasIllegalOps) {
      signalPassFailure();
    } else {
      dumpModule(moduleOp, "spatial1_graph");
      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
        signalPassFailure();
        return;
      }
    }

    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
      return;
  }
};

} // namespace

std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }

} // namespace onnx_mlir