add PIM accelerator

src/PIM/Conversion/ONNXToSpatial/Utils/WeightSubdivider.hpp

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/SmallVector.h"
+#include <map>
+
+using namespace mlir;
+using namespace std;
+
+namespace onnx_mlir {
+
+/**
+ * @brief A helper struct to store a group of weights.
+ *
+ */
+struct TaggedWeights {
+  long inputTile;
+  long outputTile;
+  size_t startingCrossbarIndex;
+  SmallVector<Value> weights;
+};
+
+/**
+ * @brief A helper class to subdivide weights into groups.
+ *
+ * Weights are stored as a map of maps of SmallVectors. The outer map is indexed
+ * by input tile, the inner map is indexed by output tile, and the SmallVector
+ * contains the weights for the filter. This class allows us to extract groups
+ * of weights from the map until we've extracted a certain number of elements,
+ * namely as many as we need to fill a compute unit.
+ */
+class WeightSubdivider {
+private:
+  map<long, map<long, SmallVector<Value>>> weights;
+  size_t crossbarsUsed = 0;
+
+  TaggedWeights popGroup(size_t amount);
+
+public:
+  WeightSubdivider(map<long, map<long, SmallVector<Value>>> weights);
+
+  bool isEmpty() const;
+  SmallVector<TaggedWeights> popGroups(size_t n);
+};
+
+} // namespace onnx_mlir