better memory report

capped vector allocations at u32::MAX in rust simulator

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -7,6 +7,7 @@
 
 #include <cassert>
 
+#include "Common/Support/CheckedArithmetic.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -18,7 +19,7 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {} // namespace
 
-llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>>
+WeightEmissionResult
 createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef outputDirPath) {
   auto coreWeightsDirPath = outputDirPath + "/weights";
   auto error = sys::fs::create_directory(coreWeightsDirPath);
@@ -26,7 +27,7 @@ createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef ou
   size_t indexFileName = 0;
 
   int64_t xbarSize = crossbarSize.getValue();
-  llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>> mapCoreWeightToFileName;
+  WeightEmissionResult result;
   llvm::SmallVector<std::pair<ResolvedWeightView, std::string>, 16> materializedWeights;
 
   auto materializeWeight = [&](const ResolvedWeightView& weightView) -> std::string {
@@ -72,17 +73,22 @@ createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef ou
 
     weightFileStream.close();
     materializedWeights.push_back({weightView, newFileName});
+    uint64_t weightBytes = pim::checkedMulOrCrash(
+      pim::checkedMulOrCrash(static_cast<size_t>(xbarSize), static_cast<size_t>(xbarSize), "weight element count"),
+      elementByteWidth,
+      "weight byte size");
+    result.totalWeightBytes = pim::checkedAddOrCrash(result.totalWeightBytes, weightBytes, "total weight bytes");
     return newFileName;
   };
 
   for (const WeightFileRequest& request : requests) {
-    auto& coreFiles = mapCoreWeightToFileName[request.coreId];
+    auto& coreFiles = result.mapCoreWeightToFileName[request.coreId];
     coreFiles.reserve(request.weights.size());
     for (const ResolvedWeightView& weight : request.weights)
       coreFiles.push_back(materializeWeight(weight));
   }
 
-  return mapCoreWeightToFileName;
+  return result;
 }
 
 } // namespace onnx_mlir