teh only weight (WIP)

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -1,5 +1,6 @@
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SetVector.h"
 
 #include <string>
 
@@ -35,6 +36,17 @@ void setComputeOperandSegmentSizes(Operation* op, int32_t weightCount, int32_t i
   cast<SpatComputeBatch>(op).getProperties().setOperandSegmentSizes({weightCount, inputCount});
 }
 
+using CrossbarWeightSet = llvm::SetVector<Value, llvm::SmallVector<Value, 4>, llvm::SmallDenseSet<Value, 4>>;
+
+CrossbarWeightSet collectCrossbarWeights(Region& body) {
+  CrossbarWeightSet weights;
+  body.walk([&](SpatVMMOp vmmOp) {
+    Value weight = vmmOp.getWeight();
+    weights.insert(weight);
+  });
+  return weights;
+}
+
 } // namespace
 
 std::optional<BlockArgument> SpatCompute::getWeightArgument(unsigned idx) { return getBlockArgument(getBody(), idx); }
@@ -45,7 +57,6 @@ std::optional<BlockArgument> SpatCompute::getInputArgument(unsigned idx) {
 
 std::optional<std::tuple<Value, BlockArgument>> SpatCompute::insertWeight(unsigned idx, Value weight, Location loc) {
   if (auto existing = llvm::find(getWeights(), weight); existing != getWeights().end()) {
-    llvm::dbgs() << "Disse netanyao\n";
     auto index = std::distance(getWeights().begin(), existing);
     return {
       {*existing, *getWeightArgument(index)}
@@ -75,6 +86,8 @@ std::optional<std::tuple<Value, BlockArgument>> SpatCompute::insertInput(unsigne
   return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
 }
 
+CrossbarWeightSet SpatCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
+
 FailureOr<std::tuple<OpResult, SpatCompute>>
 SpatCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
   if (idx > getNumResults())
@@ -127,7 +140,6 @@ std::optional<std::tuple<Value, BlockArgument>>
 SpatComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
   if (auto existing = llvm::find(getWeights(), weight); existing != getWeights().end()) {
     auto index = std::distance(getWeights().begin(), existing);
-    llvm::dbgs() << "Bum bum bum bum\n";
     return {
       {*existing, *getWeightArgument(index)}
     };
@@ -156,6 +168,8 @@ std::optional<std::tuple<Value, BlockArgument>> SpatComputeBatch::insertInput(un
   return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
 }
 
+CrossbarWeightSet SpatComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
+
 FailureOr<std::tuple<OpResult, BlockArgument, SpatComputeBatch>>
 SpatComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
   if (idx > getNumResults())