refactorone

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -128,12 +128,20 @@ FailureOr<DenseWeightView> resolveDenseWeightView(ModuleOp moduleOp, mlir::Value
 
 SmallVector<unsigned, 8> getUsedWeightIndices(Block& block) {
   SmallVector<unsigned, 8> indices;
-  auto addIndex = [&](unsigned weightIndex) {
-    if (!llvm::is_contained(indices, weightIndex))
-      indices.push_back(weightIndex);
+  auto coreOp = dyn_cast<pim::PimCoreOp>(block.getParentOp());
+  auto addWeight = [&](mlir::Value weight) {
+    if (!coreOp)
+      return;
+    for (unsigned weightIndex = 0; weightIndex < coreOp.getWeights().size(); ++weightIndex) {
+      if (coreOp.getWeightArgument(weightIndex) != weight)
+        continue;
+      if (!llvm::is_contained(indices, weightIndex))
+        indices.push_back(weightIndex);
+      return;
+    }
   };
 
-  block.walk([&](pim::PimVMMOp vmmOp) { addIndex(vmmOp.getWeightIndex()); });
+  block.walk([&](pim::PimVMMOp vmmOp) { addWeight(vmmOp.getWeight()); });
   llvm::sort(indices);
   return indices;
 }