NiccoloN
43 lines
1.3 KiB
C++
43 lines
1.3 KiB
C++
#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
|