#pragma once

#include "ContractionProblem.hpp"

#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"

namespace onnx_mlir {

enum class ContractionPlanKind {
  StaticTiled,
  BatchedDynamicVVD,
  GroupedRowDynamicVVD,
};

struct ContractionPlan {
  int64_t tileK = 1;
  int64_t tileN = 1;
  int64_t reductionSlices = 1;
  int64_t outputTiles = 1;
  int64_t laneCount = 0;
};

ContractionPlan makeContractionPlan(
    const ContractionProblem& problem,
    const spatial::SpatialTargetResources& target,
    ContractionPlanKind kind,
    int64_t laneCount = 0,
    int64_t fragmentRows = 0);

} // namespace onnx_mlir