NiccoloN
151 lines
4.0 KiB
C++
151 lines
4.0 KiB
C++
#pragma once
|
|
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/Value.h"
|
|
|
|
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
|
|
#include <cstdint>
|
|
|
|
namespace mlir {
|
|
class Operation;
|
|
} // namespace mlir
|
|
|
|
namespace onnx_mlir {
|
|
|
|
struct ConvProblem {
|
|
mlir::RankedTensorType xType;
|
|
mlir::RankedTensorType wType;
|
|
mlir::RankedTensorType outType;
|
|
int64_t batchSize;
|
|
int64_t numChannelsIn;
|
|
int64_t xHeight;
|
|
int64_t xWidth;
|
|
int64_t numChannelsOut;
|
|
int64_t wHeight;
|
|
int64_t wWidth;
|
|
int64_t outHeight;
|
|
int64_t outWidth;
|
|
int64_t group;
|
|
int64_t numChannelsInPerGroup;
|
|
int64_t numChannelsOutPerGroup;
|
|
int64_t padHeightBegin;
|
|
int64_t padHeightEnd;
|
|
int64_t padWidthBegin;
|
|
int64_t padWidthEnd;
|
|
int64_t strideHeight;
|
|
int64_t strideWidth;
|
|
int64_t dilationHeight;
|
|
int64_t dilationWidth;
|
|
bool hasBias;
|
|
bool isDepthwise = false;
|
|
bool isGrouped = false;
|
|
bool isPointwise = false;
|
|
};
|
|
|
|
struct ConvLoweringState : ConvProblem {
|
|
mlir::Operation* diagnosticAnchor = nullptr;
|
|
mlir::Value x;
|
|
mlir::Value w;
|
|
mlir::Value b;
|
|
const spatial::SpatialTargetInfo* target = nullptr;
|
|
|
|
const spatial::SpatialTargetInfo& targetInfo() const { return *target; }
|
|
};
|
|
|
|
struct ConvGeometry {
|
|
int64_t batchSize;
|
|
int64_t numChannelsIn;
|
|
int64_t xHeight;
|
|
int64_t xWidth;
|
|
int64_t numChannelsOut;
|
|
int64_t wHeight;
|
|
int64_t wWidth;
|
|
int64_t outHeight;
|
|
int64_t outWidth;
|
|
int64_t group;
|
|
int64_t numChannelsInPerGroup;
|
|
int64_t numChannelsOutPerGroup;
|
|
int64_t k;
|
|
int64_t c;
|
|
int64_t p;
|
|
int64_t xbarSize;
|
|
int64_t matrixUnitsPerProcessor;
|
|
int64_t pack;
|
|
uint64_t im2colElements;
|
|
bool hasBias;
|
|
bool isDepthwise;
|
|
};
|
|
|
|
struct RowInterval {
|
|
int64_t begin = 0;
|
|
int64_t end = 0;
|
|
};
|
|
|
|
struct ConvRowDemand {
|
|
RowInterval outputRows;
|
|
RowInterval neededInputRows;
|
|
RowInterval acquiredInputRows;
|
|
int64_t topHaloRows = 0;
|
|
int64_t bottomHaloRows = 0;
|
|
};
|
|
|
|
enum class ConvMaterializationKind : uint8_t {
|
|
StructuredDepthwise,
|
|
PointwiseContraction,
|
|
PackedIm2Col,
|
|
StreamedPatch,
|
|
StreamedPacked,
|
|
OutputChannelTiled,
|
|
InputKTiled,
|
|
Tiled2D,
|
|
};
|
|
|
|
struct ConvPlan {
|
|
ConvGeometry geometry;
|
|
spatial::ConvLoweringStrategy strategy = spatial::ConvLoweringStrategy::Auto;
|
|
ConvMaterializationKind materializationKind = ConvMaterializationKind::PackedIm2Col;
|
|
int64_t laneCount = 0;
|
|
int64_t mvmCount = 0;
|
|
int64_t vectorCount = 0;
|
|
int64_t reductionCount = 0;
|
|
uint64_t weightElements = 0;
|
|
uint64_t scratchElements = 0;
|
|
uint64_t materializationElements = 0;
|
|
uint64_t communicationElements = 0;
|
|
spatial::PhysicalLayout resultLayout = spatial::PhysicalLayout::DenseNCHW;
|
|
bool consumesRowStrip = false;
|
|
bool producesRowStrip = false;
|
|
bool requiresInputMaterialization = false;
|
|
bool requiresOutputMaterialization = false;
|
|
bool usesContraction = false;
|
|
bool hasContractionPlan = false;
|
|
ContractionPlan contraction;
|
|
};
|
|
|
|
bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
|
|
|
|
void classifyConvProblem(ConvProblem& problem);
|
|
|
|
ConvGeometry buildConvGeometry(const ConvProblem& problem,
|
|
const spatial::SpatialTargetInfo& target);
|
|
|
|
mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
|
|
spatial::ConvLoweringStrategy strategy,
|
|
const spatial::SpatialTargetInfo& target);
|
|
|
|
llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
|
|
const ConvProblem& problem, const spatial::SpatialTargetInfo& target);
|
|
|
|
uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
|
|
int64_t packFactor,
|
|
const spatial::SpatialTargetInfo& target);
|
|
|
|
RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
|
|
|
|
ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem);
|
|
|
|
} // namespace onnx_mlir
|