faster pim VerificationPass.cpp and pim code emission

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -3,17 +3,19 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <cassert>
+#include <type_traits>
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimBatchEmission.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
@@ -126,30 +128,6 @@ FailureOr<DenseWeightView> resolveDenseWeightView(ModuleOp moduleOp, mlir::Value
   return view;
 }
 
-SmallVector<unsigned, 8> getUsedWeightIndices(Block& block) {
-  SmallVector<unsigned, 8> indices;
-  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) { addWeight(vmmOp.getWeight()); });
-  llvm::sort(indices);
-  return indices;
-}
-
-SmallVector<unsigned, 8> getUsedWeightIndices(pim::PimCoreOp coreOp) {
-  return getUsedWeightIndices(coreOp.getBody().front());
-}
-
 SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
   SmallVector<Operation*> coreLikeOps;
   for (Operation& op : funcOp.getBody().front())
@@ -171,86 +149,117 @@ createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
   int64_t xbarSize = crossbarSize.getValue();
   llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>> mapCoreWeightToFileName;
   llvm::DenseMap<memref::GlobalOp, std::string> mapGlobalOpToFileName;
+  llvm::DenseMap<mlir::Value, std::string> mapWeightValueToFileName;
 
   SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
 
   for (Operation* op : coreLikeOps) {
-    auto processCore = [&](pim::PimCoreOp coreOp) {
-      size_t coreId = static_cast<size_t>(coreOp.getCoreId());
-      for (unsigned index : getUsedWeightIndices(coreOp)) {
-        if (index >= coreOp.getWeights().size()) {
-          coreOp.emitWarning("Weight index " + std::to_string(index) + " is out of range");
-          assert(index < coreOp.getWeights().size() && "Weight index is out of range");
-        }
-        mlir::Value weight = coreOp.getWeights()[index];
+    auto processWeight = [&](Operation* ownerOp,
+                             mlir::Value weight,
+                             size_t weightIndex,
+                             size_t coreId) -> LogicalResult {
+      auto weightView = resolveDenseWeightView(moduleOp, weight);
+      if (failed(weightView)) {
+        ownerOp->emitWarning("Weight is not from a memref.get_global at index " + std::to_string(weightIndex));
+        assert(succeeded(weightView) && "Weight is not from a dense memref.global view");
+      }
 
-        auto weightView = resolveDenseWeightView(moduleOp, weight);
-        if (failed(weightView)) {
-          coreOp.emitWarning("Weight is not from a memref.get_global at index " + std::to_string(index));
-          assert(succeeded(weightView) && "Weight is not from a dense memref.global view");
-        }
+      if (mapCoreWeightToFileName[coreId].contains(weight))
+        return success();
 
-        if (mapCoreWeightToFileName[coreId].contains(weight))
-          continue;
+      if (auto weightFile = mapWeightValueToFileName.find(weight); weightFile != mapWeightValueToFileName.end()) {
+        mapCoreWeightToFileName[coreId].insert({weight, weightFile->second});
+        return success();
+      }
 
-        auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
-        auto globalOp = getGlobalOp ? lookupGlobalForGetGlobal(moduleOp, getGlobalOp) : memref::GlobalOp {};
-        if (globalOp && mapGlobalOpToFileName.contains(globalOp)) {
-          auto& fileName = mapGlobalOpToFileName[globalOp];
-          mapCoreWeightToFileName[coreId].insert({weight, fileName});
-          continue;
-        }
+      auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
+      auto globalOp = getGlobalOp ? lookupGlobalForGetGlobal(moduleOp, getGlobalOp) : memref::GlobalOp {};
+      if (globalOp && mapGlobalOpToFileName.contains(globalOp)) {
+        auto& fileName = mapGlobalOpToFileName[globalOp];
+        mapWeightValueToFileName[weight] = fileName;
+        mapCoreWeightToFileName[coreId].insert({weight, fileName});
+        return success();
+      }
 
-        DenseElementsAttr denseAttr = weightView->denseAttr;
-        ArrayRef<int64_t> shape = weightView->shape;
-        assert(isMatrixShape(shape) && "Weight matrix must be 2-dimensional");
-        int64_t numRows = shape[0];
-        int64_t numCols = shape[1];
-        assert(numRows <= xbarSize && numCols <= xbarSize && "Weight dimensions must not exceed crossbar size");
+      DenseElementsAttr denseAttr = weightView->denseAttr;
+      ArrayRef<int64_t> shape = weightView->shape;
+      assert(isMatrixShape(shape) && "Weight matrix must be 2-dimensional");
+      int64_t numRows = shape[0];
+      int64_t numCols = shape[1];
+      assert(numRows <= xbarSize && numCols <= xbarSize && "Weight dimensions must not exceed crossbar size");
 
-        size_t elementByteWidth = getElementTypeSizeInBytes(denseAttr.getElementType());
+      size_t elementByteWidth = getElementTypeSizeInBytes(denseAttr.getElementType());
 
-        std::string newFileName = "crossbar_" + std::to_string(indexFileName++) + ".bin";
-        auto weightFilePath = (coreWeightsDirPath + "/" + newFileName).str();
-        std::error_code errorCode;
-        raw_fd_ostream weightFileStream(weightFilePath, errorCode, sys::fs::OF_None);
-        if (errorCode) {
-          errs() << "Error while opening weight file `" << weightFilePath << "`: " << errorCode.message() << '\n';
-          assert(errorCode);
-        }
+      std::string newFileName = "crossbar_" + std::to_string(indexFileName++) + ".bin";
+      auto weightFilePath = (coreWeightsDirPath + "/" + newFileName).str();
+      std::error_code errorCode;
+      raw_fd_ostream weightFileStream(weightFilePath, errorCode, sys::fs::OF_None);
+      if (errorCode) {
+        errs() << "Error while opening weight file `" << weightFilePath << "`: " << errorCode.message() << '\n';
+        assert(errorCode);
+      }
 
-        uint64_t zero = 0;
-        for (int64_t row = 0; row < xbarSize; row++) {
-          for (int64_t col = 0; col < xbarSize; col++) {
-            if (row < numRows && col < numCols) {
-              int64_t elementIndex = weightView->offset + row * weightView->strides[0] + col * weightView->strides[1];
-              APInt bits = denseAttr.getValues<APFloat>()[elementIndex].bitcastToAPInt();
-              uint64_t word = bits.getZExtValue();
-              weightFileStream.write(reinterpret_cast<const char*>(&word), elementByteWidth);
-            }
-            else {
-              weightFileStream.write(reinterpret_cast<const char*>(&zero), elementByteWidth);
-            }
+      uint64_t zero = 0;
+      for (int64_t row = 0; row < xbarSize; row++) {
+        for (int64_t col = 0; col < xbarSize; col++) {
+          if (row < numRows && col < numCols) {
+            int64_t elementIndex = weightView->offset + row * weightView->strides[0] + col * weightView->strides[1];
+            APInt bits = denseAttr.getValues<APFloat>()[elementIndex].bitcastToAPInt();
+            uint64_t word = bits.getZExtValue();
+            weightFileStream.write(reinterpret_cast<const char*>(&word), elementByteWidth);
+          }
+          else {
+            weightFileStream.write(reinterpret_cast<const char*>(&zero), elementByteWidth);
           }
         }
-
-        weightFileStream.close();
-        if (globalOp)
-          mapGlobalOpToFileName.insert({globalOp, newFileName});
-        mapCoreWeightToFileName[coreId].insert({weight, newFileName});
       }
+
+      weightFileStream.close();
+      if (globalOp)
+        mapGlobalOpToFileName.insert({globalOp, newFileName});
+      mapWeightValueToFileName[weight] = newFileName;
+      mapCoreWeightToFileName[coreId].insert({weight, newFileName});
       return success();
     };
 
+    auto processCoreLike = [&](auto coreLikeOp) {
+      auto usedIndices = getUsedWeightIndices(coreLikeOp);
+      for (unsigned index : usedIndices) {
+        if (index >= coreLikeOp.getWeights().size()) {
+          coreLikeOp.emitWarning("Weight index " + std::to_string(index) + " is out of range");
+          assert(index < coreLikeOp.getWeights().size() && "Weight index is out of range");
+        }
+      }
+
+      if constexpr (std::is_same_v<std::decay_t<decltype(coreLikeOp)>, pim::PimCoreOp>) {
+        size_t coreId = static_cast<size_t>(coreLikeOp.getCoreId());
+        for (unsigned index : usedIndices)
+          if (failed(processWeight(coreLikeOp, coreLikeOp.getWeights()[index], index, coreId)))
+            return failure();
+        return success();
+      }
+      else {
+        auto batchCoreIds = getBatchCoreIds(coreLikeOp);
+        SmallVector<size_t> orderedCoreIds;
+        llvm::SmallSet<size_t, 8> seenCoreIds;
+        for (int32_t coreId : batchCoreIds)
+          if (seenCoreIds.insert(static_cast<size_t>(coreId)).second)
+            orderedCoreIds.push_back(static_cast<size_t>(coreId));
+
+        for (size_t coreId : orderedCoreIds)
+          for (unsigned index : usedIndices)
+            if (failed(processWeight(coreLikeOp, coreLikeOp.getWeights()[index], index, coreId)))
+              return failure();
+        return success();
+      }
+    };
+
     if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-      (void) processCore(coreOp);
+      (void) processCoreLike(coreOp);
       continue;
     }
 
-    auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
-    for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
-      if (failed(withScalarCoreFromBatchLane(coreBatchOp, lane, processCore)))
-        return mapCoreWeightToFileName;
+    (void) processCoreLike(cast<pim::PimCoreBatchOp>(op));
   }
   return mapCoreWeightToFileName;
 }