big refactor

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp

@@ -3,14 +3,19 @@
 #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 ConvLoweringState {
-  mlir::Value x;
-  mlir::Value w;
-  mlir::Value b;
+struct ConvProblem {
   mlir::RankedTensorType xType;
   mlir::RankedTensorType wType;
   mlir::RankedTensorType outType;
@@ -35,29 +40,32 @@ struct ConvLoweringState {
   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 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 {
@@ -73,14 +81,36 @@ struct ConvRowDemand {
   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);
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state);
+void classifyConvProblem(ConvProblem& problem);
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor);
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetResources& target);
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options);
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+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