#include "ContractionPlanning.hpp"

#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"

#include <algorithm>

namespace onnx_mlir {

namespace {

static int64_t ceilDivide(int64_t value, int64_t divisor) {
  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
}

} // namespace

ContractionPlan makeContractionPlan(
    const ContractionProblem& problem,
    const spatial::SpatialTargetResources& target,
    ContractionPlanKind kind,
    int64_t laneCount,
    int64_t fragmentRows) {
  ContractionPlan plan;
  plan.tileK = std::max<int64_t>(1, target.matrixShape.rows);
  plan.tileN = std::max<int64_t>(1, target.matrixShape.columns);
  plan.reductionSlices = std::max<int64_t>(1, ceilDivide(problem.k, plan.tileK));
  plan.outputTiles = std::max<int64_t>(1, ceilDivide(problem.n, plan.tileN));
  const int64_t rowsPerLane = std::max<int64_t>(
      1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows);

  if (laneCount != 0)
    plan.laneCount = laneCount;
  else if (kind == ContractionPlanKind::StaticTiled)
    plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles;
  else if (kind == ContractionPlanKind::GroupedRowDynamicVVD)
    plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane);
  else
    plan.laneCount = problem.batch * problem.m * problem.n;
  return plan;
}

} // namespace onnx_mlir