NiccoloN
117 lines
3.1 KiB
C++
117 lines
3.1 KiB
C++
#pragma once
|
|
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/Value.h"
|
|
|
|
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.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 problem;
|
|
mlir::Operation* diagnosticAnchor = nullptr;
|
|
mlir::Value x;
|
|
mlir::Value w;
|
|
mlir::Value b;
|
|
const spatial::SpatialTargetResources* target = nullptr;
|
|
const ONNXToSpatialPlanningOptions* options = nullptr;
|
|
|
|
const spatial::SpatialTargetResources& targetInfo() const { return *target; }
|
|
const ONNXToSpatialPlanningOptions& planningOptions() const;
|
|
};
|
|
|
|
struct ConvGeometry {
|
|
int64_t k;
|
|
int64_t c;
|
|
int64_t p;
|
|
int64_t xbarSize;
|
|
int64_t matrixUnitsPerProcessor;
|
|
int64_t pack;
|
|
uint64_t im2colElements;
|
|
};
|
|
|
|
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,
|
|
PackedIm2Col,
|
|
StreamedPatch,
|
|
StreamedPacked,
|
|
InputKTiled,
|
|
};
|
|
|
|
struct ConvPlan {
|
|
ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
|
|
};
|
|
|
|
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::SpatialTargetResources& target);
|
|
|
|
mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
|
|
spatial::ConvLoweringStrategy strategy,
|
|
const spatial::SpatialTargetResources& target,
|
|
const ONNXToSpatialPlanningOptions& options);
|
|
|
|
uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
|
|
int64_t packFactor,
|
|
const ONNXToSpatialPlanningOptions& options);
|
|
|
|
RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
|
|
|
|
ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem);
|
|
|
|
} // namespace onnx_mlir
|