faster pim VerificationPass.cpp and pim code emission

src/PIM/Common/IR/WeightUtils.hpp

@@ -3,9 +3,15 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/StringRef.h"
 
+#include <optional>
+
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
 inline constexpr llvm::StringRef PimWeightAlwaysAttrName = "weightAlways";
 
 namespace onnx_mlir {
@@ -26,4 +32,24 @@ bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value);
 /// passes can identify globals that must remain weight-backed.
 void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback);
 
+template <typename CoreLikeOpTy>
+llvm::SmallVector<unsigned, 8> getUsedWeightIndices(CoreLikeOpTy coreLikeOp) {
+  llvm::SmallVector<unsigned, 8> indices;
+  auto addWeight = [&](mlir::Value weight) {
+    for (unsigned weightIndex = 0; weightIndex < coreLikeOp.getWeights().size(); ++weightIndex) {
+      if (coreLikeOp.getWeightArgument(weightIndex) != weight)
+        continue;
+      if (!llvm::is_contained(indices, weightIndex))
+        indices.push_back(weightIndex);
+      return;
+    }
+  };
+
+  coreLikeOp.walk([&](pim::PimVMMOp vmmOp) { addWeight(vmmOp.getWeight()); });
+  llvm::sort(indices);
+  return indices;
+}
+
+std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, pim::PimVMMOp vmmOp);
+
 } // namespace onnx_mlir